old capstone

Author: tesuji

pwnlib/asm.py

@@ -1023,12 +1023,22 @@ def get_cs_disassembler(eabi=None):
 
     B = {16: cs.CS_MODE_16, 32: cs.CS_MODE_32, 64: cs.CS_MODE_64}[context.bits]
 
+    try:
+        CS_ARCH_AARCH64 = cs.CS_ARCH_AARCH64
+    except Exception:
+        CS_ARCH_AARCH64 = cs.CS_ARCH_ARM64
+
+    try:
+        CS_ARCH_SYSTEMZ = cs.CS_ARCH_SYSTEMZ
+    except Exception:
+        CS_ARCH_SYSTEMZ = cs.CS_ARCH_SYSZ
+
     params = {
         'i386'   : (cs.CS_ARCH_X86, B),
         'amd64'  : (cs.CS_ARCH_X86, B),
         'thumb'  : (cs.CS_ARCH_ARM, cs.CS_MODE_THUMB + E),
         'arm'    : (cs.CS_ARCH_ARM, cs.CS_MODE_ARM + E),
-        'aarch64': (cs.CS_ARCH_AARCH64, cs.CS_MODE_ARM + E),
+        'aarch64': (CS_ARCH_AARCH64, cs.CS_MODE_ARM + E),
         'armhf'  : (cs.CS_ARCH_ARM, cs.CS_MODE_THUMB + E),
         'mips'   : (cs.CS_ARCH_MIPS, cs.CS_MODE_32 + E),
         'mips64' : (cs.CS_ARCH_MIPS, cs.CS_MODE_64 + E),
@@ -1037,7 +1047,7 @@ def get_cs_disassembler(eabi=None):
         'ppc'    : (cs.CS_ARCH_PPC, B + E),
         'powerpc':   (cs.CS_ARCH_PPC, E + cs.CS_MODE_32),
         'powerpc64': (cs.CS_ARCH_PPC, E + cs.CS_MODE_64),
-        'em_s390': (cs.CS_ARCH_SYSTEMZ, cs.CS_MODE_BIG_ENDIAN + cs.CS_MODE_64),
+        'em_s390': (CS_ARCH_SYSTEMZ, cs.CS_MODE_BIG_ENDIAN + cs.CS_MODE_64),
         #'ia64': None,
         #'m68k': cs.CS_ARCH_M68K,
         #'xcore': cs.CS_ARCH_XCORE,

pwnlib/elf/elf.py

@@ -1191,8 +1191,6 @@ def libc_start_main_return(self):
         elif self.arch == 'aarch64':
             pass
         elif self.arch in ['mips', 'mips64']:
-            # FIXME: `bal` was not included in CS_GRP_CALL. This is fixed on capstone v6.alpha
-            call_instructions = call_instructions.add(cs.CS_GRP_BRANCH_RELATIVE)
             # Account for the delay slot.
             call_return_offset = 2
         elif self.arch in ['i386', 'amd64', 'ia64']:
@@ -1215,6 +1213,13 @@ def libc_start_main_return(self):
         dis = list(self.cs_disasm(md, func.address, func.size))
 
         filter_calls = lambda dis: ((i, x) for i, x in enumerate(dis) if call_instructions & set(x.groups))
+
+        if self.arch in ['ppc', 'powerpc', 'powerpc64']:
+            filter_calls = lambda dis: ((i, x) for i, x in enumerate(dis) if set([x.mnemonic]) & set(['bctrl', 'bl']))
+        # FIXME: `bal` was not included in CS_GRP_CALL. This is fixed on capstone v6.alpha
+        elif self.arch in ['mips', 'mips64']:
+            filter_calls = lambda dis: ((i, x) for i, x in enumerate(dis) if set([x.mnemonic]) & set(['bal', 'jalr']))
+
         calls = list(filter_calls(dis))
 
         def find_ret_main_addr(caller_dis, calls):
@@ -1235,9 +1240,6 @@ def find_ret_main_addr(caller_dis, calls):
         if ret_addr:
             return ret_addr
 
-        if self.arch in ['ppc', 'powerpc', 'powerpc64']:
-            filter_calls = lambda dis: ((i, x) for i, x in enumerate(dis) if set([x.mnemonic]) & set(['bctrl', 'bl']))
-
         # `__libc_start_main` -> `__libc_start_call_main` -> `main`
         # Find a direct call which calls `exit` once. That's probably `__libc_start_call_main`.
         for _, insn in calls: