Reusing IDA's system call analysis results, and just utilizing miasm in case IDA failed to identify the system call.

1 files changed, 30 insertions, 5 deletions

diff --git a/SystemCalls.py b/SystemCalls.py
index 0494972..c564029 100644
--- a/SystemCalls.py
+++ b/SystemCalls.py
@@ -15,6 +15,7 @@ from miasm2.analysis.depgraph import DependencyGraph
 from utils import guess_machine
 
 from idaapi import *
+import idc
 import ida_bytes
 import ida_idp
 import ida_search
@@ -63,7 +64,7 @@ class SystemCallView(Choose2):
 
         if len(self.systemCalls.x86) != 0:
             for faddr in self.systemCalls.x86.iterkeys():
-                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86[faddr])
+                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86[faddr], x86SystemCalls)
                 for call in calls:
                     try:
                         self.items.append(["0x%08X" % call[0],
@@ -81,7 +82,7 @@ class SystemCallView(Choose2):
 
         if len(self.systemCalls.x86_64) != 0:
             for faddr in self.systemCalls.x86_64.iterkeys():
-                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86_64[faddr])
+                calls = self.systemCalls.getSystemCallNumber(self.systemCalls.x86_64[faddr], x86_64SystemCalls)
                 for call in calls:
                     try:
                         self.items.append(["0x%08X" % call[0],
@@ -193,18 +194,42 @@ class SystemCall():
                 continue
             self.mdis.symbol_pool.add_label(name, ad)
 
+        self.elements = set([self.mn.regs.RAX])
+
         self.x86 = dict()
         self.x86_64 = dict()
 
         self.systemCallView = SystemCallView(self)
 
-    def getSystemCallNumber(self, func):
+    def __getSystemCallNumberByComment(self, addr, scstrings):
+        cmt = idc.Comment(addr)
+
+        if cmt and cmt.startswith("LINUX - "):
+            try:
+                return scstrings.index(cmt.replace("LINUX - ", ""))
+            except:
+                return None
+
+
+    def getSystemCallNumber(self, func, scstrings):
         """ Get the value of rax/eax at the time of the system call.
         """
 
         sol = list()
-        elements = set([self.mn.regs.RAX])
 
+        # Get the analysis results from IDA, by reading IDA's comments at system calls.
+        calls = set()
+        
+        for call in func.calls:
+            number = self.__getSystemCallNumberByComment(call.addr, scstrings)
+
+            if number:
+                sol.append([call.addr, call.sctype, str(number)])
+                calls.add(call)
+
+        func.calls -= calls
+
+        # Just proceed with depgraph if IDA detected a function.
         if not func.f:
             for call in func.calls:
                 sol.append([call.addr, call.sctype, ''])
@@ -251,7 +276,7 @@ class SystemCall():
 
             # Get dependency graphs
             dg = DependencyGraph(self.ir_arch, follow_call=False)
-            graphs = dg.get(cur_label, elements, line_nb,
+            graphs = dg.get(cur_label, self.elements, line_nb,
                             set([self.ir_arch.symbol_pool.getby_offset(func.f.startEA)]))
 
             while 1: