Some refactoring.

2 files changed, 299 insertions, 359 deletions

diff --git a/SystemCalls.py b/SystemCalls.py
new file mode 100644
index 0000000..4fcffd2
--- /dev/null
+++ b/SystemCalls.py
@@ -0,0 +1,297 @@
+""" IDAPython plugin to display all system calls of AMD/Intel 32/64bit
+    Linux userland binaries.
+
+    Work in progress...
+
+    The system call ABI from the following link are supported.
+    http://esec-lab.sogeti.com/post/2011/07/05/Linux-syscall-ABI
+    
+    by n0p
+"""
+
+from miasm2.core.bin_stream_ida import bin_stream_ida
+from miasm2.analysis.depgraph import DependencyGraph
+
+from utils import guess_machine
+
+from idaapi import *
+
+from SystemCalls_constants import *
+
+
+
+class SystemCallView(Choose2):
+
+    def __init__(self, systemCalls):
+
+        self.systemCalls = systemCalls
+
+        Choose2.__init__(self, 
+                         "System call", 
+                         [ ["Address", 13], 
+                           ["Type", 10], 
+                           ["Number", 10], 
+                           ["Name", 20], 
+                           ["Pointer Size", 12] ])
+
+        self.items = list()
+
+        self.nop_items = list()
+
+        self.icon = 150
+
+        self.cmd_nop = None
+
+        self.initialized = False
+
+    def __fillView(self):
+
+        self.systemCalls.searchSystemCalls()
+
+        self.items = list()
+
+        if len(self.systemCalls.x86) != 0:
+            for call in self.systemCalls.x86:
+                try:
+                    callNr = int(self.systemCalls.getSystemCallNumber(call[0])[0])
+                    self.items.append(["0x%X" % call[0],
+                                       systemCallTypes[call[1]],
+                                       "0x%03X" % callNr,
+                                       x86SystemCalls[callNr],
+                                       "32bit"])
+                except:
+                    # No hex system call number found.
+                    self.items.append(["0x%X" % call[0],
+                                       systemCallTypes[call[1]],
+                                       "",
+                                       "",
+                                       "32bit"])
+
+        if len(self.systemCalls.x86_64) != 0:
+            for call in self.systemCalls.x86_64:
+                try:
+                    callNr = int(self.systemCalls.getSystemCallNumber(call[0])[0])
+                    self.items.append(["0x%X" % call[0],
+                                       systemCallTypes[call[1]],
+                                       "0x%03X" % callNr,
+                                       x86_64SystemCalls[callNr],
+                                       "64bit"])
+                except:
+                    # No hex system call number found.
+                    self.items.append(["0x%X" % call[0],
+                                       systemCallTypes[call[1]],
+                                       "",
+                                       "",
+                                       "64bit"])
+
+    def OnClose(self):
+        pass
+
+    def OnCommand(self, n, cmd_id):
+        if cmd_id == self.cmd_nop:
+            start_ea = int(self.items[n][0], 16)
+            end_ea = start_ea+ItemSize(start_ea)
+
+            self.nop_items.append(self.items[n][0])
+
+            for i in xrange(start_ea, end_ea):
+                PatchByte(i, 0x90)
+
+    def OnGetIcon(self, n):
+        if not len(self.items) > 0:
+            return -1
+
+        if self.items[n][2] == "":
+            # No system call number found => display red icon.
+            return 59
+        else:
+            # Display green icon.
+            return 61
+
+    def OnGetLine(self, n):
+        return self.items[n]
+
+    def OnGetLineAttr(self, n):
+        if len(self.nop_items) > 0:
+            if self.items[n][0] in self.nop_items:
+                return [0xB6B6B4, 0]
+
+    def OnGetSize(self):
+        return len(self.items)
+
+    def OnRefresh(self, n):
+        self.__fillView()
+
+    def OnSelectLine(self, n):
+        idaapi.jumpto(int(self.items[n][0], 16))
+
+    def show(self):
+        if not self.initialized:
+            self.initialized = True
+            self.__fillView()
+
+        if self.Show() < 0: return False
+
+        if self.cmd_nop == None:
+            self.cmd_nop = self.AddCommand("NOP system call", 
+                                           flags = idaapi.CHOOSER_POPUP_MENU, 
+                                           icon=50)
+
+        return True
+
+
+class SystemCall():
+
+    def __init__(self):
+
+        # Init miasm stuff.
+        self.machine = guess_machine()
+        self.mn, self.dis_engine, self.ira = self.machine.mn, self.machine.dis_engine, self.machine.ira
+
+        self.mdis = self.dis_engine(bin_stream_ida(), dont_dis_nulstart_bloc=True)
+        self.ir_arch = self.ira(self.mdis.symbol_pool)
+
+        # Populate symbols with ida names
+        for ad, name in Names():
+            if name is None:
+                continue
+            self.mdis.symbol_pool.add_label(name, ad)
+
+        self.x86 = list()
+        self.x86_64 = list()
+
+        self.systemCallView = SystemCallView(self)
+
+    def getSystemCallNumber(self, addr):
+        """ Get the value of rax/eax at the time of the system call.
+        """
+
+        sol = list()
+
+        # Get the current function
+        f = get_func(addr)
+
+        if not f:
+            return sol
+
+        blocs = self.mdis.dis_multibloc(f.startEA)
+
+        # Generate IR
+        for bloc in blocs:
+            self.ir_arch.add_bloc(bloc)
+
+        # Check if addr is in a basic block without an entry.
+        if len(self.ir_arch.getby_offset(addr)) == 0:
+            fc = qflow_chart_t("", f, BADADDR, BADADDR, FC_PREDS)
+
+            # Iterate through all basic blocks.
+            for i in xrange(0, fc.size()):
+                if fc[i].startEA <= addr and addr < fc[i].endEA:
+                    # Basic block without entry found.
+                    blocs = self.mdis.dis_multibloc(fc[i].startEA)
+
+                    # Generate IR
+                    for bloc in blocs:
+                        self.ir_arch.add_bloc(bloc)
+
+        cur_bloc = list(self.ir_arch.getby_offset(addr))[0]
+        cur_label = cur_bloc.label
+
+        elements = set([self.mn.regs.RAX])
+
+        for line_nb, l in enumerate(cur_bloc.lines):
+            if l.offset == addr:
+                break
+
+        # Get dependency graphs
+        dg = DependencyGraph(self.ir_arch, follow_call=False)
+        graphs = dg.get(cur_label, elements, line_nb,
+                        set([self.ir_arch.symbol_pool.getby_offset(f.startEA)]))
+
+        while 1:
+            try:
+                graph = graphs.next()
+            except StopIteration:
+                break
+
+            if not graph.has_loop:
+                sol.append(graph.emul().values()[0])
+
+        return sol
+
+    def searchSystemCalls(self):
+        """ Looks for 'int 80', 'sysenter', 'syscall' and 'gs:[10h]' system calls.
+        """
+        # Iterating through all segments
+        for segment in Segments():
+            # Iterating through all instructions in each segment
+            for head in Heads(segment, SegEnd(segment)):
+                if isCode(GetFlags(head)):
+                    inst = DecodeInstruction(head)
+
+                    if  (    inst.size == 2
+                         and get_many_bytes(head, inst.size)
+                                    .startswith("\xCD\x80")):
+                        # int 80h. Just on 32bit.
+                        if not [head, 0] in self.x86:
+                            self.x86.append([head, 0])
+
+                    elif  (    inst.size == 2
+                         and get_many_bytes(head, inst.size)
+                                    .startswith("\x0F\x34")):
+                        # sysenter. Just on 32bit.
+                        if not [head, 1] in self.x86:
+                            self.x86.append([head, 1])
+
+                    elif  (    inst.size == 2
+                         and get_many_bytes(head, inst.size)
+                                    .startswith("\x0F\x05")):
+                        # syscall. 32bit just on AMD. 64bit on AMD and Intel.
+                        if idaapi.ph.flag & idaapi.PR_USE64:
+                            if not [head, 2] in self.x86_64:
+                                self.x86_64.append([head, 2])
+                        else:
+                            if not [head, 2] in self.x86:
+                                self.x86.append([head, 2])
+
+                    elif  (    inst.size == 7
+                         and get_many_bytes(head, inst.size)
+                                    .startswith("\x65\xFF\x15\x10\x00\x00\x00")):
+                        # gs:[10h]. Just on 32bit.
+                        if not [head, 3] in self.x86:
+                            self.x86.append([head, 3])
+
+    def showView(self):
+        self.systemCallView.show()
+
+
+class SystemCallPlugin_t(idaapi.plugin_t):
+    flags = 0
+    comment = ""
+    help = ""
+    wanted_name = "System Calls"
+    wanted_hotkey = ""
+
+    def init(self):
+        global systemCalls
+
+        if idaapi.ph_get_id() == idaapi.PLFM_386:
+            # Check if already initialized
+            if not 'systemCalls' in globals():
+                systemCalls = SystemCall()
+
+            return idaapi.PLUGIN_KEEP
+        else:
+            return idaapi.PLUGIN_SKIP
+
+    def run(self, arg):
+        global systemCalls
+
+        systemCalls.showView()
+
+    def term(self):
+        if 'systemCalls' in globals():
+            del globals()['systemCalls']
+
+def PLUGIN_ENTRY():
+    return SystemCallPlugin_t()
diff --git a/idaSystemCalls.py b/SystemCalls_constants.py
index 88bb374..4083060 100644
--- a/idaSystemCalls.py
+++ b/SystemCalls_constants.py
@@ -1,20 +1,5 @@
-""" IDAPython plugin to display all system calls of AMD/Intel 32/64bit
-    Linux userland binaries.
- 
-    Work in progress...
- 
-    The system call ABI from the following link are supported.
-    http://esec-lab.sogeti.com/post/2011/07/05/Linux-syscall-ABI
-    
-    by n0p
-"""
- 
-from miasm2.core.bin_stream_ida import bin_stream_ida
-from miasm2.analysis.depgraph import DependencyGraph
-from miasm2.analysis.machine import Machine
- 
-from idaapi import *
- 
+systemCallTypes = ["int 80h", "sysenter", "syscall", "gs:[10h]"]
+
 # Taken from /usr/include/asm/unistd_32.h at Arch Linux i686 3.16.3-1.
 x86SystemCalls = ["sys_restart_syscall",
                   "sys_exit",
@@ -684,345 +669,3 @@ x86_64SystemCalls = ["sys_read",
                      "sys_finit_module",
                      "sys_sched_setattr",
                      "sys_sched_getattr"]
- 
-systemCallTypes = ["int 80h", "sysenter", "syscall", "gs:[10h]"]
- 
-
-# max_size_to_size and guess_machine from miasm/example/ida/utils.py
-def max_size_to_size(max_size):
-    for size in [16, 32, 64]:
-        if (1 << size) - 1 == max_size:
-            return size
-    return None
- 
-def guess_machine():
-    "Return an instance of Machine corresponding to the IDA guessed processor"
- 
-    processor_name = GetLongPrm(INF_PROCNAME)
-    max_size = GetLongPrm(INF_START_SP)
-    size = max_size_to_size(max_size)
- 
-    if processor_name == "metapc":
- 
-        # HACK: check 32/64 using INF_START_SP
-        if max_size == 0x80:  # TODO XXX check
-            machine = Machine("x86_16")
-        elif size == 32:
-            machine = Machine("x86_32")
-        elif size == 64:
-            machine = Machine("x86_64")
-        else:
-            raise ValueError('cannot guess 32/64 bit! (%x)' % max_size)
-    elif processor_name == "ARM":
-        # TODO ARM/thumb
-        # hack for thumb: set armt = True in globals :/
-        # set bigendiant = True is bigendian
-        # Thumb, size, endian
-        info2machine = {(True, 32, True): "armtb",
-                        (True, 32, False): "armtl",
-                        (False, 32, True): "armb",
-                        (False, 32, False): "arml",
-                        (False, 64, True): "aarch64b",
-                        (False, 64, False): "aarch64l",
-                        }
-        is_armt = globals().get('armt', False)
-        is_bigendian = globals().get('bigendian', False)
-        infos = (is_armt, size, is_bigendian)
-        if not infos in info2machine:
-            raise NotImplementedError('not fully functional')
-        machine = Machine(info2machine[infos])
- 
-        from miasm2.analysis.disasm_cb import guess_funcs, guess_multi_cb
-        from miasm2.analysis.disasm_cb import arm_guess_subcall, arm_guess_jump_table
-        guess_funcs.append(arm_guess_subcall)
-        guess_funcs.append(arm_guess_jump_table)
- 
-    elif processor_name == "msp430":
-        machine = Machine("msp430")
-    elif processor_name == "mipsl":
-        machine = Machine("mips32l")
-    elif processor_name == "mipsb":
-        machine = Machine("mips32b")
-    else:
-        print repr(processor_name)
-        raise NotImplementedError('not fully functional')
- 
-    return machine
- 
- 
-class SystemCallView(Choose2):
- 
-    def __init__(self, systemCalls):
- 
-        self.systemCalls = systemCalls
- 
-        Choose2.__init__(self, 
-                         "System call", 
-                         [ ["Address", 13], 
-                           ["Type", 10], 
-                           ["Number", 10], 
-                           ["Name", 20], 
-                           ["Pointer Size", 12] ])
- 
-        self.items = list()
- 
-        self.nop_items = list()
- 
-        self.icon = 150
- 
-        self.cmd_nop = None
- 
-        self.initialized = False
- 
-    def __fillView(self):
- 
-        self.systemCalls.searchSystemCalls()
- 
-        self.items = list()
- 
-        if len(self.systemCalls.x86) != 0:
-            for call in self.systemCalls.x86:
-                try:
-                    callNr = int(self.systemCalls.getSystemCallNumber(call[0])[0])
-                    self.items.append(["0x%X" % call[0],
-                                       systemCallTypes[call[1]],
-                                       "0x%03X" % callNr,
-                                       x86SystemCalls[callNr],
-                                       "32bit"])
-                except:
-                    # No hex system call number found.
-                    self.items.append(["0x%X" % call[0],
-                                       systemCallTypes[call[1]],
-                                       "",
-                                       "",
-                                       "32bit"])
- 
-        if len(self.systemCalls.x86_64) != 0:
-            for call in self.systemCalls.x86_64:
-                try:
-                    callNr = int(self.systemCalls.getSystemCallNumber(call[0])[0])
-                    self.items.append(["0x%X" % call[0],
-                                       systemCallTypes[call[1]],
-                                       "0x%03X" % callNr,
-                                       x86_64SystemCalls[callNr],
-                                       "64bit"])
-                except:
-                    # No hex system call number found.
-                    self.items.append(["0x%X" % call[0],
-                                       systemCallTypes[call[1]],
-                                       "",
-                                       "",
-                                       "64bit"])
- 
-    def OnClose(self):
-        pass
- 
-    def OnCommand(self, n, cmd_id):
-        if cmd_id == self.cmd_nop:
-            start_ea = int(self.items[n][0], 16)
-            end_ea = start_ea+ItemSize(start_ea)
- 
-            self.nop_items.append(self.items[n][0])
- 
-            for i in xrange(start_ea, end_ea):
-                PatchByte(i, 0x90)
- 
-    def OnGetIcon(self, n):
-        if not len(self.items) > 0:
-            return -1
- 
-        if self.items[n][2] == "":
-            # No system call number found => display red icon.
-            return 59
-        else:
-            # Display green icon.
-            return 61
- 
-    def OnGetLine(self, n):
-        return self.items[n]
- 
-    def OnGetLineAttr(self, n):
-        if len(self.nop_items) > 0:
-            if self.items[n][0] in self.nop_items:
-                return [0xB6B6B4, 0]
- 
-    def OnGetSize(self):
-        return len(self.items)
- 
-    def OnRefresh(self, n):
-        self.__fillView()
- 
-    def OnSelectLine(self, n):
-        idaapi.jumpto(int(self.items[n][0], 16))
- 
-    def show(self):
-        if not self.initialized:
-            self.initialized = True
-            self.__fillView()
- 
-        if self.Show() < 0: return False
- 
-        if self.cmd_nop == None:
-            self.cmd_nop = self.AddCommand("NOP system call", 
-                                           flags = idaapi.CHOOSER_POPUP_MENU, 
-                                           icon=50)
- 
-        return True
- 
- 
-class SystemCall():
- 
-    def __init__(self):
- 
-        self.x86 = list()
-        self.x86_64 = list()
- 
-        self.systemCallView = SystemCallView(self)
- 
-    def getSystemCallNumber(self, addr):
-        """ Get the value of rax/eax at the time of the system call.
-        """
-
-        sol = list()
- 
-        # Init
-        machine = guess_machine()
-        mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
- 
-        bs = bin_stream_ida()
-        mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
-        ir_arch = ira(mdis.symbol_pool)
- 
-        # Populate symbols with ida names
-        for ad, name in Names():
-            if name is None:
-                continue
-            mdis.symbol_pool.add_label(name, ad)
- 
-        # Get the current function
-        f = get_func(addr)
-
-        if not f:
-            return sol
-
-        blocs = mdis.dis_multibloc(f.startEA)
-
-        # Generate IR
-        for bloc in blocs:
-            ir_arch.add_bloc(bloc)
-
-        # Check if addr is in a basic block without an entry.
-        if len(ir_arch.getby_offset(addr)) == 0:
-            fc = qflow_chart_t("", f, BADADDR, BADADDR, FC_PREDS)
-
-            # Iterate through all basic blocks.
-            for i in xrange(0, fc.size()):
-                if fc[i].startEA <= addr and addr < fc[i].endEA:
-                    # Basic block without entry found.
-                    blocs = mdis.dis_multibloc(fc[i].startEA)
-
-                    # Generate IR
-                    for bloc in blocs:
-                        ir_arch.add_bloc(bloc)
- 
-        cur_bloc = list(ir_arch.getby_offset(addr))[0]
-        cur_label = cur_bloc.label
- 
-        elements = set([mn.regs.RAX])
- 
-        for line_nb, l in enumerate(cur_bloc.lines):
-            if l.offset == addr:
-                break
- 
-        # Get dependency graphs
-        dg = DependencyGraph(ir_arch, follow_call=False)
-        graphs = dg.get(cur_label, elements, line_nb,
-                        set([ir_arch.symbol_pool.getby_offset(f.startEA)]))
- 
-        while 1:
-            try:
-                graph = graphs.next()
-            except StopIteration:
-                break
- 
-            if not graph.has_loop:
-                sol.append(graph.emul().values()[0])
- 
-        return sol
- 
-    def searchSystemCalls(self):
-        """ Looks for 'int 80', 'sysenter', 'syscall' and 'gs:[10h]' system calls.
-        """
-        # Iterating through all segments
-        for segment in Segments():
-            # Iterating through all instructions in each segment
-            for head in Heads(segment, SegEnd(segment)):
-                if isCode(GetFlags(head)):
-                    inst = DecodeInstruction(head)
- 
-                    if  (    inst.size == 2
-                         and get_many_bytes(head, inst.size)
-                                    .startswith("\xCD\x80")):
-                        # int 80h. Just on 32bit.
-                        if not [head, 0] in self.x86:
-                            self.x86.append([head, 0])
- 
-                    elif  (    inst.size == 2
-                         and get_many_bytes(head, inst.size)
-                                    .startswith("\x0F\x34")):
-                        # sysenter. Just on 32bit.
-                        if not [head, 1] in self.x86:
-                            self.x86.append([head, 1])
- 
-                    elif  (    inst.size == 2
-                         and get_many_bytes(head, inst.size)
-                                    .startswith("\x0F\x05")):
-                        # syscall. 32bit just on AMD. 64bit on AMD and Intel.
-                        if idaapi.ph.flag & idaapi.PR_USE64:
-                            if not [head, 2] in self.x86_64:
-                                self.x86_64.append([head, 2])
-                        else:
-                            if not [head, 2] in self.x86:
-                                self.x86.append([head, 2])
- 
-                    elif  (    inst.size == 7
-                         and get_many_bytes(head, inst.size)
-                                    .startswith("\x65\xFF\x15\x10\x00\x00\x00")):
-                        # gs:[10h]. Just on 32bit.
-                        if not [head, 3] in self.x86:
-                            self.x86.append([head, 3])
- 
-    def showView(self):
-        self.systemCallView.show()
- 
- 
-class SystemCallPlugin_t(idaapi.plugin_t):
-    flags = 0
-    comment = ""
-    help = ""
-    wanted_name = "System Calls"
-    wanted_hotkey = ""
- 
-    def init(self):
-        global systemCalls
- 
-        if idaapi.ph_get_id() == idaapi.PLFM_386:
-            # Check if already initialized
-            if not 'systemCalls' in globals():
-                systemCalls = SystemCall()
- 
-            return idaapi.PLUGIN_KEEP
-        else:
-            return idaapi.PLUGIN_SKIP
- 
-    def run(self, arg):
-        global systemCalls
- 
-        systemCalls.showView()
- 
-    def term(self):
-        if 'systemCalls' in globals():
-            del globals()['systemCalls']
- 
-def PLUGIN_ENTRY():
-    return SystemCallPlugin_t()