Some more refactoring.

1 files changed, 130 insertions, 81 deletions

diff --git a/SystemCalls.py b/SystemCalls.py
index 4fcffd2..6a8d979 100644
--- a/SystemCalls.py
+++ b/SystemCalls.py
@@ -46,43 +46,57 @@ class SystemCallView(Choose2):
 
     def __fillView(self):
 
+        start = time.time()
+
         self.systemCalls.searchSystemCalls()
 
+        end = time.time()
+        print ("[*] It took {} seconds to discover the system calls.".format(end-start))
+
         self.items = list()
 
+        start = time.time()
+
         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"])
+            for faddr in self.systemCalls.x86.iterkeys():
+                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86[faddr])
+                for call in calls:
+                    try:
+                        self.items.append(["0x%08X" % call[0],
+                                           systemCallTypes[call[1]],
+                                           "0x%03X" % int(call[2]),
+                                           x86SystemCalls[int(call[2])],
+                                           "32bit"])
+                    except:
+                        # No hex system call number found.
+                        self.items.append(["0x%08X" % call[0],
+                                           systemCallTypes[call[1]],
+                                           str(call[2]),
+                                           "",
+                                           "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"])
+            for faddr in self.systemCalls.x86_64.iterkeys():
+                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86_64[faddr])
+                for call in calls:
+                    try:
+                        self.items.append(["0x%08X" % call[0],
+                                           systemCallTypes[call[1]],
+                                           "0x%03X" % int(call[2]),
+                                           x86_64SystemCalls[int(call[2])],
+                                           "64bit"])
+                    except:
+                        # No hex system call number found.
+                        self.items.append(["0x%08X" % call[0],
+                                           systemCallTypes[call[1]],
+                                           str(call[2]),
+                                           "",
+                                           "64bit"])
+
+        end = time.time()
+        print ("[*] It took {} seconds to analyze the system calls.".format(end-start))
+
+        self.items.sort(key=lambda tup:tup[0])
 
     def OnClose(self):
         pass
@@ -101,7 +115,7 @@ class SystemCallView(Choose2):
         if not len(self.items) > 0:
             return -1
 
-        if self.items[n][2] == "":
+        if self.items[n][3] == "":
             # No system call number found => display red icon.
             return 59
         else:
@@ -142,6 +156,24 @@ class SystemCallView(Choose2):
 
 class SystemCall():
 
+    class Func():
+        def __init__(self):
+
+            # System calls inside the function.
+            self.calls = set()
+
+            # Return value of idaapi.get_func(addr)
+            self.f = None
+
+    class Call():
+        def __init__(self, addr, sctype):
+
+            self.addr = addr
+            self.sctype = sctype
+
+        def __hash__(self):
+            return self.addr
+
     def __init__(self):
 
         # Init miasm stuff.
@@ -157,68 +189,91 @@ class SystemCall():
                 continue
             self.mdis.symbol_pool.add_label(name, ad)
 
-        self.x86 = list()
-        self.x86_64 = list()
+        self.x86 = dict()
+        self.x86_64 = dict()
 
         self.systemCallView = SystemCallView(self)
 
-    def getSystemCallNumber(self, addr):
+    def getSystemCallNumber(self, func):
         """ Get the value of rax/eax at the time of the system call.
         """
 
         sol = list()
+        elements = set([self.mn.regs.RAX])
 
-        # Get the current function
-        f = get_func(addr)
-
-        if not f:
+        if not func.f:
             return sol
 
-        blocs = self.mdis.dis_multibloc(f.startEA)
+        blocs = self.mdis.dis_multibloc(func.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)
+        try:
+            for bloc in blocs:
+                self.ir_arch.add_bloc(bloc)
+        except:
+            for call in func.calls:
+                sol.append([call.addr, call.sctype, ''])
+            return sol
 
-            # 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)
+        for call in func.calls:
+            addr = call.addr
 
-                    # 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("", func.f, BADADDR, BADADDR, FC_PREDS)
 
-        cur_bloc = list(self.ir_arch.getby_offset(addr))[0]
-        cur_label = cur_bloc.label
+                try:
+                    # 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)
+                except:
+                    sol.append([addr, call.sctype, ''])
+                    continue
 
-        elements = set([self.mn.regs.RAX])
+            cur_bloc = list(self.ir_arch.getby_offset(addr))[0]
+            cur_label = cur_bloc.label
 
-        for line_nb, l in enumerate(cur_bloc.lines):
-            if l.offset == addr:
-                break
+            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)]))
+            # 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(func.f.startEA)]))
 
-        while 1:
-            try:
-                graph = graphs.next()
-            except StopIteration:
-                break
+            while 1:
+                try:
+                    graph = graphs.next()
+                except StopIteration:
+                    break
 
-            if not graph.has_loop:
-                sol.append(graph.emul().values()[0])
+                if not graph.has_loop:
+                    sol.append([addr, call.sctype, graph.emul().values()[0]])
 
         return sol
 
+    def __addCall(self, addr, sctype, arch):
+        f = get_func(addr)
+
+        if not f:
+            faddr = 0
+        else:
+            faddr = f.startEA
+
+        if faddr in arch:
+            arch[faddr].calls.add(self.Call(addr, sctype))
+        else:
+            arch[faddr] = self.Func()
+            arch[faddr].calls.add(self.Call(addr, sctype))
+            arch[faddr].f = f
+
     def searchSystemCalls(self):
         """ Looks for 'int 80', 'sysenter', 'syscall' and 'gs:[10h]' system calls.
         """
@@ -233,33 +288,27 @@ class SystemCall():
                          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])
+                        self.__addCall(head, 0, self.x86)
 
                     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])
+                        self.__addCall(head, 1, self.x86)
 
                     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])
+                            self.__addCall(head, 2, self.x86_64)
                         else:
-                            if not [head, 2] in self.x86:
-                                self.x86.append([head, 2])
-
+                            self.__addCall(head, 2, self.x86)
                     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])
+                        self.__addCall(head, 3, self.x86)
 
     def showView(self):
         self.systemCallView.show()