[+] Auto-renaming dummy-named functions, which have one API call

[+] Assigning TAGS to functions accordingly to called API-indicators inside
[*] Setting tags as repeatable function comments and displaying TAG tree in the separate view

auto_re.py

@@ -0,0 +1,368 @@
+# -*- coding: utf-8 -*
+__author__ = 'Trafimchuk Aliaksandr'
+
+import idaapi
+from idautils import FuncItems
+from idaapi import o_reg, o_imm, o_far, o_near, o_mem
+import traceback
+
+
+# HAS_PYSIDE = False
+# try:
+#     from PySide import QtGui, QtCore
+#     from PySide.QtGui import QTreeView, QVBoxLayout, QLineEdit
+#
+#     _slot = QtCore.Slot
+#     global HAS_PYSIDE
+#     HAS_PYSIDE = True
+# except ImportError:
+#     from PyQt5 import QtGui, QtCore
+#     from PyQt5.QtWidgets import QTreeView, QVBoxLayout, QLineEdit
+#
+#     # dummy
+#     def _slot(fn):
+#         def wrap(*args, **kwargs):
+#             return fn(*args, **kwargs)
+#         return wrap
+
+from PyQt5 import QtGui, QtCore
+from PyQt5.QtWidgets import QTreeView, QVBoxLayout, QLineEdit
+
+
+TAGS_IGNORE_LIST = {
+    'OpenProcessToken',
+    'DisconnectNamedPipe'
+}
+
+IGNORE_CALL_LIST = {
+    'RtlNtStatusToDosError',
+    'GetLastError',
+    'SetLastError'
+}
+
+TAGS = {
+    'net': ['WSAStartup', 'socket', 'recv', 'recvfrom', 'send', 'sendto', 'acccept', 'bind', 'listen', 'select',
+            'setsockopt', 'ioctlsocket', 'closesocket', 'WSAAccept', 'WSARecv', 'WSARecvFrom', 'WSASend', 'WSASendTo',
+            'WSASocket', 'WSAConnect', 'ConnectEx', 'TransmitFile', 'HTTPOpenRequest', 'HTTPSendRequest',
+            'URLDownloadToFile', 'InternetCrackUrl', 'InternetOpen', 'InternetOpen', 'InternetConnect',
+            'InternetOpenUrl', 'InternetQueryOption', 'InternetSetOption', 'InternetReadFile', 'InternetWriteFile',
+            'InternetGetConnectedState', 'InternetSetStatusCallback', 'DnsQuery', 'getaddrinfo', 'GetAddrInfo',
+            'GetAdaptersInfo', 'GetAdaptersAddresses', 'HttpQueryInfo', 'ObtainUserAgentString', 'WNetGetProviderName',
+            'GetBestInterfaceEx', 'gethostbyname', 'getsockname', 'connect'],
+    'spawn': ['CreateProcess', 'ShellExecute', 'ShellExecuteEx', 'system', 'CreateProcessInternal', 'NtCreateProcess',
+              'ZwCreateProcess', 'NtCreateProcessEx', 'ZwCreateProcessEx', 'NtCreateUserProcess', 'ZwCreateUserProcess',
+              'RtlCreateUserProcess', 'NtCreateSection', 'ZwCreateSection', 'NtOpenSection', 'ZwOpenSection',
+              'NtAllocateVirtualMemory', 'ZwAllocateVirtualMemory', 'NtWriteVirtualMemory', 'ZwWriteVirtualMemory',
+              'NtMapViewOfSection', 'ZwMapViewOfSection', 'OpenSCManager', 'CreateService', 'OpenService',
+              'StartService', 'ControlService'],
+    'inject': ['OpenProcess-disabled', 'ZwOpenProcess', 'WriteProcessMemory', 'CreateRemoteThread', 'QueueUserAPC'],
+    'com': ['CoCreateInstance', 'CoInitializeSecurity', 'CoGetClassObject', 'OleConvertOLESTREAMToIStorage'],
+    'crypto': ['CryptAcquireContext', 'CryptProtectData', 'CryptUnprotectData', 'CryptProtectMemory',
+               'CryptUnprotectMemory', 'CryptDecrypt', 'CryptEncrypt', 'CryptHashData', 'CryptDecodeMessage',
+               'CryptDecryptMessage', 'CryptEncryptMessage', 'CryptHashMessage', 'CryptExportKey', 'CryptGenKey',
+               'CryptCreateHash', 'CryptDecodeObjectEx', 'EncryptMessage', 'DecryptMessage']
+}
+
+blacklist = {'@__security_check_cookie@4', '__SEH_prolog4', '__SEH_epilog4'}
+replacements = [
+    ('??3@YAXPAX@Z', 'alloc'),
+    ('?', '')
+]
+
+
+def get_addr_width():
+    return '16' if idaapi.cvar.inf.is_64bit() else '8'
+
+
+class AutoREView(idaapi.PluginForm):
+    ADDR_ROLE = QtCore.Qt.UserRole + 1
+
+    def __init__(self, data):
+        super(AutoREView, self).__init__()
+        self._data = data
+
+    def Show(self):
+        return idaapi.PluginForm.Show(self, 'AutoRE', options=idaapi.PluginForm.FORM_PERSIST)
+
+    def OnCreate(self, form):
+        # if HAS_PYSIDE:
+        #     self.parent = self.FormToPySideWidget(form)
+        # else:
+        self.parent = self.FormToPyQtWidget(form)
+
+        self.tv = QTreeView()
+        self.tv.setExpandsOnDoubleClick(False)
+
+        root_layout = QVBoxLayout(self.parent)
+        # self.le_filter = QLineEdit(self.parent)
+
+        # root_layout.addWidget(self.le_filter)
+        root_layout.addWidget(self.tv)
+
+        self.parent.setLayout(root_layout)
+
+        self._model = QtGui.QStandardItemModel()
+        self._init_model()
+        self.tv.setModel(self._model)
+
+        self.tv.setColumnWidth(0, 200)
+        self.tv.setColumnWidth(1, 300)
+        self.tv.header().setStretchLastSection(True)
+
+        self.tv.expandAll()
+
+        self.tv.doubleClicked.connect(self.on_navigate_to_method_requested)
+        # self.le_filter.textChanged.connect(self.on_filter_text_changed)
+
+    def OnClose(self, form):
+        # print 'TODO: OnClose(): clear the pointer to form in the plugin'
+        pass
+
+    def _tv_init_header(self, model):
+        item_header = QtGui.QStandardItem("EA")
+        item_header.setToolTip("Address")
+        model.setHorizontalHeaderItem(0, item_header)
+
+        item_header = QtGui.QStandardItem("Funcion name")
+        model.setHorizontalHeaderItem(1, item_header)
+
+        item_header = QtGui.QStandardItem("API called")
+        model.setHorizontalHeaderItem(2, item_header)
+
+    def _tv_make_tag_item(self, name):
+        rv = QtGui.QStandardItem(name)
+
+        rv.setEditable(False)
+        return [rv, QtGui.QStandardItem(), QtGui.QStandardItem()]
+
+    def _tv_make_ref_item(self, tag, ref):
+        ea_item = QtGui.QStandardItem(('%#0' + get_addr_width() + 'X') % ref['ea'])
+        ea_item.setEditable(False)
+        ea_item.setData(ref['ea'], self.ADDR_ROLE)
+
+        name_item = QtGui.QStandardItem(ref['name'])
+        name_item.setEditable(False)
+        name_item.setData(ref['ea'], self.ADDR_ROLE)
+
+        api_name = QtGui.QStandardItem(ref['tags'][tag])
+        api_name.setEditable(False)
+        api_name.setData(ref['ea'], self.ADDR_ROLE)
+
+        return [ea_item, name_item, api_name]
+
+    def _init_model(self):
+        self._model.clear()
+
+        root_node = self._model.invisibleRootItem()
+        self._tv_init_header(self._model)
+
+        for tag, refs in self._data.items():
+            item_tag_list = self._tv_make_tag_item(tag)
+            item_tag = item_tag_list[0]
+
+            root_node.appendRow(item_tag_list)
+
+            for ref in refs:
+                ref_item_list = self._tv_make_ref_item(tag, ref)
+
+                item_tag.appendRow(ref_item_list)
+
+    def on_navigate_to_method_requested(self, index):
+        addr = index.data(role=self.ADDR_ROLE)
+        if addr is not None:
+            idaapi.jumpto(addr)
+
+    # def on_filter_text_changed(self, text):
+    #     print 'on_text_changed: %s' % text
+
+
+class auto_re_t(idaapi.plugin_t):
+    flags = idaapi.PLUGIN_UNL
+    comment = ""
+
+    help = ""
+    wanted_name = "Auto RE"
+    wanted_hotkey = "Ctrl+Shift+M"
+
+    _PREFIX_NAME = 'au_re_'
+    _MIN_MAX_MATH_OPS_TO_ALLOW_RENAME = 10
+
+    def init(self):
+        # self._cfg = None
+        self.view = None
+        # self._load_config()
+
+        return idaapi.PLUGIN_OK
+
+    # def _load_config(self):
+    #     self._cfg = {'auto_rename': False}
+
+    # def _store_config(self, cfg):
+    #     pass
+
+    def _handle_tags(self, fn, fn_an):
+        tags = fn_an['tags']
+        if not tags:
+            return
+        print 'fn: %#08x tags: %s' % (fn.startEA, tags)
+        cmt = idaapi.get_func_cmt(fn, True)
+        if cmt:
+            cmt += '\n'
+        s = str(tags.keys())
+        name = idaapi.get_ea_name(fn.startEA)
+        item = {'ea': fn.startEA, 'name': name, 'tags': tags}
+        if not cmt or s not in cmt:
+            idaapi.set_func_cmt(fn, '%sTAGS: %s' % (cmt or '', s), True)
+        # self.mark_position(fn.startEA, 'TAGS: %s' % s)
+        for tag in tags:
+            if tag not in self._data:
+                self._data[tag] = list()
+            self._data[tag].append(item)
+
+    def _handle_calls(self, fn, fn_an):
+        num_calls = len(fn_an['calls'])
+        if num_calls != 1:
+            return
+
+        dis = fn_an['calls'][0]
+        if dis.Op1.type not in (o_imm, o_far, o_near, o_mem):
+            return
+
+        ea = dis.Op1.value
+        if not ea and dis.Op1.addr:
+            ea = dis.Op1.addr
+
+        if idaapi.has_dummy_name(idaapi.getFlags(ea)):
+            return
+
+        possible_name = idaapi.get_ea_name(ea)
+        if not possible_name or possible_name in blacklist:
+            return
+
+        normalized = self.normalize_name(possible_name)
+
+        # if self._cfg.get('auto_rename'):
+        if len(fn_an['math']) < self._MIN_MAX_MATH_OPS_TO_ALLOW_RENAME:
+            idaapi.do_name_anyway(fn.startEA, normalized)
+        # TODO: add an API to the view
+        print 'fn: %#08x: %d calls, %d math%s possible name: %s, normalized: %s' % (
+            fn.startEA, len(fn_an['calls']), len(fn_an['math']), 'has bads' if fn_an['has_bads'] else '',
+            possible_name, normalized)
+
+    def run(self, arg):
+        try:
+            self._data = dict()
+            count = idaapi.get_func_qty()
+            for i in xrange(count):
+                fn = idaapi.getn_func(i)
+                fn_an = self.analyze_func(fn)
+
+                # if fn_an['math']:
+                # 	print 'fn: %#08x has math' % fn.startEA
+
+                if idaapi.has_dummy_name(idaapi.getFlags(fn.startEA)):
+                    self._handle_calls(fn, fn_an)
+
+                self._handle_tags(fn, fn_an)
+
+            if self.view:
+                self.view.Close(idaapi.PluginForm.FORM_NO_CONTEXT)
+            self.view = AutoREView(self._data)
+            self.view.Show()
+        except:
+            idaapi.msg('AutoRE: address not found\nerror: %s\n' % traceback.format_exc())
+
+    def term(self):
+        pass
+
+    @classmethod
+    def disasm_func(cls, fn):
+        rv = list()
+        items = list(FuncItems(fn.startEA))
+        for item_ea in items:
+            obj = {'ea': item_ea, 'fn_ea': fn.startEA, 'dis': None}
+            if idaapi.decode_insn(item_ea) > 0:
+                obj['dis'] = idaapi.cmd.copy()
+            rv.append(obj)
+        return rv
+
+    @classmethod
+    def _analysis_handle_call_insn(cls, dis, rv):
+        rv['calls'].append(dis)
+        if dis.Op1.type != o_mem or not dis.Op1.addr:
+            return
+
+        name = idaapi.get_ea_name(dis.Op1.addr)
+        name = name.replace('__imp__', '')
+
+        if '@' in name:
+            name = name.split('@')[0]
+
+        if not name:
+            return
+
+        if name in IGNORE_CALL_LIST:
+            rv['calls'].pop()
+            return
+
+        for tag, names in TAGS.items():
+            if tag in rv['tags']:
+                continue
+            if name in TAGS_IGNORE_LIST:
+                continue
+
+            for tag_api in names:
+                if tag_api in name:
+                    # print '%#08x: %s, tag: %s' % (dis.ea, name, tag)
+                    rv['tags'][tag] = name
+                    break
+
+    @classmethod
+    def analyze_func(cls, fn):
+        rv = {'fn': fn, 'calls': [], 'math': [], 'has_bads': False, 'tags': {}}
+        items = cls.disasm_func(fn)
+
+        for item in items:
+            dis = item['dis']
+            if dis is None:
+                rv['has_bads'] = True
+                continue
+
+            if dis.itype in (idaapi.NN_call, idaapi.NN_callfi, idaapi.NN_callni):
+                cls._analysis_handle_call_insn(dis, rv)
+            elif dis.itype == idaapi.NN_xor:
+                if dis.Op1.type == o_reg and dis.Op2.type == o_reg and dis.Op1.reg == dis.Op2.reg:
+                    continue
+                rv['math'].append(dis)
+            elif dis.itype in (idaapi.NN_shr, idaapi.NN_shl, idaapi.NN_sal, idaapi.NN_sar, idaapi.NN_ror,
+                               idaapi.NN_rol, idaapi.NN_rcl, idaapi.NN_rcl):
+                # TODO
+                rv['math'].append(dis)
+
+        return rv
+
+    @classmethod
+    def normalize_name(cls, n):
+        for repl in replacements:
+            n = n.replace(*repl)
+        if '@' in n:
+            n = n.split('@')[0]
+        if len(n) < 3:
+            return ''
+        if not n.startswith(cls._PREFIX_NAME):
+            n = cls._PREFIX_NAME + n
+        return n
+
+    # @classmethod
+    # def mark_position(cls, ea, name, slot=[0]):
+    #     curloc = idaapi.curloc()
+    #     curloc.ea = ea
+    #     curloc.lnnum = 0
+    #     curloc.x = 0
+    #     curloc.y = 0
+    #     slot[0] += 1
+    #     curloc.mark(slot[0], name, name)
+
+
+def PLUGIN_ENTRY():
+    return auto_re_t()