Discovery phase reworked and implemented with the new IDAPython API.

1 files changed, 40 insertions, 32 deletions

diff --git a/SystemCalls.py b/SystemCalls.py
index 6a8d979..0494972 100644
--- a/SystemCalls.py
+++ b/SystemCalls.py
@@ -15,6 +15,10 @@ from miasm2.analysis.depgraph import DependencyGraph
 from utils import guess_machine
 
 from idaapi import *
+import ida_bytes
+import ida_idp
+import ida_search
+import ida_segment
 
 from SystemCalls_constants import *
 
@@ -202,6 +206,8 @@ class SystemCall():
         elements = set([self.mn.regs.RAX])
 
         if not func.f:
+            for call in func.calls:
+                sol.append([call.addr, call.sctype, ''])
             return sol
 
         blocs = self.mdis.dis_multibloc(func.f.startEA)
@@ -274,41 +280,43 @@ class SystemCall():
             arch[faddr].calls.add(self.Call(addr, sctype))
             arch[faddr].f = f
 
+    def __findCalls(self, seg, sbytes, slength, sctype, arch):
+        addr = ida_search.find_binary(seg.startEA, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)
+
+        while addr != BADADDR:
+            if (    ida_bytes.get_item_head(addr) == addr
+                and ida_bytes.get_item_size(addr) == slength):
+                self.__addCall(addr, sctype, arch)
+
+            addr = ida_search.find_binary(addr+1, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)
+
+        addr = ida_search.find_binary(seg.startEA, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)
+
     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.
-                        self.__addCall(head, 0, self.x86)
-
-                    elif  (    inst.size == 2
-                         and get_many_bytes(head, inst.size)
-                                    .startswith("\x0F\x34")):
-                        # sysenter. Just on 32bit.
-                        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:
-                            self.__addCall(head, 2, self.x86_64)
-                        else:
-                            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.
-                        self.__addCall(head, 3, self.x86)
+        seg = ida_segment.get_first_seg()
+
+        while seg:
+
+            # Check if segment is executable
+            if seg.perm & 1:
+                # int 80h. Just on 32bit.
+                self.__findCalls(seg, "CD 80", 2, 0, self.x86)
+
+                # sysenter. Just on 32bit.
+                self.__findCalls(seg, "0F 34", 2, 1, self.x86)
+
+                # syscall. 32bit just on AMD. 64bit on AMD and Intel.
+                if ida_idp.ph.flag & ida_idp.PR_USE64:
+                    self.__findCalls(seg, "0F 05", 2, 2, self.x86_64)
+                else:
+                    self.__findCalls(seg, "0F 05", 2, 2, self.x86)
+
+                # gs:[10h]. Just on 32bit.
+                self.__findCalls(seg, "65 FF 15 10 00 00 00", 7, 3, self.x86)
+
+            seg = ida_segment.get_next_seg(seg.startEA)
 
     def showView(self):
         self.systemCallView.show()