From 4930805acbdad346acf8aa106bfed9ae79884a50 Mon Sep 17 00:00:00 2001
From: kumarak <iit.akshay@gmail.com>
Date: Wed, 10 Feb 2021 22:28:04 -0500
Subject: [PATCH] Fix unhandled type class (#95)

---
 python/anvill/__main__.py | 12 ++++++++++++
 python/anvill/binja.py    | 22 +++++++++++++++++++---
 python/anvill/util.py     | 35 +++++++++++++++++++++++++++++++++++
 python/anvill/var.py      |  4 +++-
 setup.py                  |  2 +-
 5 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 python/anvill/util.py

diff --git a/python/anvill/__main__.py b/python/anvill/__main__.py
index 94dbd7b28..d6a07ee79 100644
--- a/python/anvill/__main__.py
+++ b/python/anvill/__main__.py
@@ -15,9 +15,12 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
+import sys
 import argparse
 import json
 
+from .util import INIT_DEBUG_FILE
 from .binja import get_program
 
 
@@ -43,8 +46,17 @@ def main():
         default=False,
     )
 
+    arg_parser.add_argument(
+        "--log_file",
+        type=argparse.FileType('w'), default=os.devnull,
+        help="Log to a specific file."
+    )
+
     args = arg_parser.parse_args()
 
+    if args.log_file != os.devnull:
+        INIT_DEBUG_FILE(args.log_file)
+
     p = get_program(args.bin_in)
 
     ep = None
diff --git a/python/anvill/binja.py b/python/anvill/binja.py
index a51896975..4d33cca83 100644
--- a/python/anvill/binja.py
+++ b/python/anvill/binja.py
@@ -24,6 +24,7 @@
 from .os import *
 from .type import *
 from .program import *
+from .util import *
 
 
 def is_valid_addr(bv, addr):
@@ -243,9 +244,13 @@ def _convert_bn_type(tinfo: bn.types.Type, cache):
         bn.TypeClass.NamedTypeReferenceClass,
         bn.TypeClass.WideCharTypeClass,
     ]:
-        raise UnhandledTypeException(
-            "Unhandled VarArgs, Value, or WideChar type: {}".format(str(tinfo)), tinfo
-        )
+        err_type_class = {
+            bn.TypeClass.VarArgsTypeClass : "VarArgsTypeClass",
+            bn.TypeClass.ValueTypeClass : "ValueTypeClass",
+            bn.TypeClass.NamedTypeReferenceClass : "NamedTypeReferenceClass",
+            bn.TypeClass.WideCharTypeClass : "WideCharTypeClass",
+        }
+        DEBUG("WARNING: Unhandled type class {}".format(err_type_class[tinfo.type_class]))
 
     else:
         raise UnhandledTypeException("Unhandled type: {}".format(str(tinfo)), tinfo)
@@ -472,6 +477,10 @@ def visit(self, program, is_definition, add_refs_as_defs):
         if not is_definition:
             return
 
+        # type could be None if type class not handled
+        if self._type is None:
+            return
+
         if isinstance(self._type, VoidType):
             return
 
@@ -484,6 +493,11 @@ def visit(self, program, is_definition, add_refs_as_defs):
         for ea in range(begin, end):
             br.seek(ea)
             seg = bv.get_segment_at(ea)
+            # _elf_header is getting recovered as variable
+            # get_segment_at(...) returns None for elf_header
+            if seg is None:
+                continue
+
             mem.map_byte(ea, br.read8(), seg.writable, seg.executable)
 
 
@@ -598,6 +612,8 @@ def get_program(*args, **kargs):
         return _PROGRAM
     assert len(args) == 1
 
+    DEBUG("Recovering program {}".format(args[0]))
+
     prog = BNProgram(args[0])
     if "cache" not in kargs or kargs["cache"]:
         _PROGRAM = prog
diff --git a/python/anvill/util.py b/python/anvill/util.py
new file mode 100644
index 000000000..d9f9407e8
--- /dev/null
+++ b/python/anvill/util.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2021 Trail of Bits, Inc.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+_DEBUG_FILE = None
+_DEBUG_PREFIX = ""
+
+def INIT_DEBUG_FILE(file):
+	global _DEBUG_FILE
+	_DEBUG_FILE = file
+
+def DEBUG_PUSH():
+	global _DEBUG_PREFIX
+	_DEBUG_PREFIX += "  "
+
+def DEBUG_POP():
+	global _DEBUG_PREFIX
+	_DEBUG_PREFIX = _DEBUG_PREFIX[:-2]
+
+def DEBUG(s):
+	global _DEBUG_FILE
+	if _DEBUG_FILE:
+		_DEBUG_FILE.write("{}{}\n".format(_DEBUG_PREFIX, str(s)))
diff --git a/python/anvill/var.py b/python/anvill/var.py
index 85c20045d..89ea2f61c 100644
--- a/python/anvill/var.py
+++ b/python/anvill/var.py
@@ -44,5 +44,7 @@ def is_declaration(self):
     def proto(self):
         proto = {}
         proto["address"] = self.address()
-        proto["type"] = self.type().proto(self._arch)
+        if self.type() != None:
+            proto["type"] = self.type().proto(self._arch)
+
         return proto
diff --git a/setup.py b/setup.py
index b477df23a..b0c96267c 100644
--- a/setup.py
+++ b/setup.py
@@ -32,4 +32,4 @@
         "anvill.__init__", "anvill.__main__", "anvill.arch", "anvill.binja",
         "anvill.exc", "anvill.function", "anvill.ida", "anvill.loc",
         "anvill.mem", "anvill.os", "anvill.program", "anvill.type",
-        "anvill.var"])
+        "anvill.var", "anvill.util"])