Merge pull request #99 from imbillow/get_signature

Enhanced Function.get_signature

Author: XVilka

idb/analysis.py

@@ -3,8 +3,7 @@
 import itertools
 import logging
 import types
-from collections import Counter
-from collections import namedtuple
+from collections import Counter, namedtuple
 from random import randint
 
 import six
@@ -13,6 +12,8 @@
 
 import idb
 import idb.netnode
+import idb.typeinf
+import idb.typeinf_flags
 
 logger = logging.getLogger(__name__)
 
@@ -1217,6 +1218,18 @@ def pairs(l):
 StackChangePoint = namedtuple("StackChangePoint", ["effective_address", "change"])
 
 
+def create_pstring_list(buf):
+    _lst = []
+    ofs = 0
+    while ofs < len(buf.strip(b"\x00")):
+        _len = struct.unpack_from("<B", buf, ofs)[0]
+        _lst.append(buf[ofs : ofs + _len][1:].decode("utf-8"))
+        ofs += _len
+        if _len == 0:
+            ofs += 1
+    return _lst
+
+
 class Function:
     """
     Example::
@@ -1238,41 +1251,24 @@ def get_name(self):
             return "sub_%X" % (self.nodeid)
 
     def get_signature(self):
-        typebuf = self.netnode.supval(tag="S", index=0x3000)
-        namebuf = self.netnode.supval(tag="S", index=0x3001)
-
-        if six.indexbytes(typebuf, 0x0) != 0xC:
-            raise RuntimeError("unexpected signature header")
-
-        if six.indexbytes(typebuf, 0x1) == ord("S"):
-            # this is just a guess...
-            conv = "__stdcall"
-        else:
-            raise NotImplementedError()
-
-        rtype = TypeString()
-        rtype.vsParse(typebuf, offset=2)
-
-        # this is a guess???
-        sp_delta = typebuf[2 + len(rtype)]
-
-        params = []
-        typeoffset = 0x2 + len(rtype) + 0x1
-        nameoffset = 0x0
-        while typeoffset < len(typebuf):
-            if six.indexbytes(typebuf, typeoffset) == 0x0:
-                break
-            typename = TypeString()
-            typename.vsParse(typebuf, offset=typeoffset)
-            typeoffset += len(typename)
-
-            paramname = PString()
-            paramname.vsParse(namebuf, offset=nameoffset)
-            nameoffset += len(paramname)
-
-            params.append(FunctionParameter(typename.s, paramname.s))
-
-        return FunctionSignature(conv, rtype.s, sp_delta, params)
+        try:
+            typebuf = self.netnode.supval(tag="S", index=0x3000)
+            namebuf = self.netnode.supval(tag="S", index=0x3001)
+            typ = six.indexbytes(typebuf, 0x0)
+            if not idb.typeinf_flags.is_type_func(typ):
+                raise RuntimeError("is not function")
+
+            names = create_pstring_list(namebuf)
+            typedata = idb.typeinf.FuncTypeData()
+            ts = idb.typeinf.TypeString(typebuf)
+            typedata.deserialize(self.idb.til, ts, names, [])
+            inf = Root(self.idb).idainfo
+
+            return idb.typeinf.TInfo(
+                typ, typedata, til=self.idb.til, name=self.get_name(), inf=inf
+            )
+        except KeyError:
+            return None
 
     def get_chunks(self):
         v = self.netnode.supval(tag="S", index=0x7000)

idb/fileformat.py

@@ -13,6 +13,7 @@
 
 import idb
 import idb.netnode
+from idb.analysis import Root
 from idb.typeinf import TIL
 
 try:
@@ -1194,6 +1195,8 @@ def pcb_header(self):
 
             s = sectiondef.cls(buf=section.contents, wordsize=self.wordsize)
             s.vsParse(section.contents)
+            if isinstance(s, TIL):
+                s.inf = Root(self).idainfo
             # vivisect doesn't allow you to assign vstructs to
             #  attributes that are not part of the struct,
             # so we need to override and use the default object behavior.

idb/idapython.py

@@ -1,10 +1,9 @@
 # -*- coding: utf-8 -*-
 import collections
+import logging
 import os
 import re
 import struct
-import logging
-import re
 import weakref
 
 import six
@@ -1213,23 +1212,8 @@ def GetType(self, ea):
         except Exception as e:
             logger.warning("failed to fetch function for GetType: %s", e)
             return None
-
-        try:
-            name = f.get_name()
-            sig = f.get_signature()
-        except KeyError:
-            return None
-
-        params = []
-        for param in sig.parameters:
-            params.append("%s %s" % (param.type, param.name))
-
-        return "{rtype:s} {cc:s} {name:s}({params:s})".format(
-            rtype=sig.rtype,
-            cc=sig.calling_convention,
-            name=name,
-            params=", ".join(params),
-        )
+        sig = f.get_signature()
+        return sig.get_typestr() if sig is not None else None
 
     @staticmethod
     def hasValue(flags):