Use capstone to implement ELF.libc_start_main_return

Author: tesuji

CHANGELOG.md

@@ -75,6 +75,7 @@ The table below shows which release corresponds to each branch, and what date th
 
 ## 5.0.0 (`dev`)
 
+- [#2580][2580] Use capstone to implement `ELF.libc_start_main_return`
 - [#2419][2419] riscv: avoid compressed instructions (if you need compressed, use .option rvc)
 - [#2551][2551] Detect when kitty is being used as terminal
 - [#2519][2519] Drop Python 2.7 support / Require Python 3.10
@@ -94,6 +95,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2575][2575] Detect when Terminator is being used as terminal
 - [#2578][2578] Add gnome-terminal, Alacritty, Ttilix for run_in_new_terminal
 
+[2580]: https://github.com/Gallopsled/pwntools/pull/2580
 [2419]: https://github.com/Gallopsled/pwntools/pull/2419
 [2551]: https://github.com/Gallopsled/pwntools/pull/2551
 [2519]: https://github.com/Gallopsled/pwntools/pull/2519

pwnlib/asm.py

@@ -388,10 +388,10 @@ def _bfdname():
         'sparc64' : 'elf64-sparc',
     }
 
-    if arch in bfdnames:
-        return bfdnames[arch]
-    else:
+    name = bfdnames.get(arch)
+    if not name:
         raise Exception("Cannot find bfd name for architecture %r" % arch)
+    return name
 
 
 def _bfdarch():
@@ -409,10 +409,7 @@ def _bfdarch():
         'loongarch64': 'loongarch64'
     }
 
-    if arch in convert:
-        return convert[arch]
-
-    return arch
+    return convert.get(arch, arch)
 
 def _run(cmd, stdin = None):
     log.debug('%s', subprocess.list2cmdline(cmd))
@@ -1015,3 +1012,52 @@ def disasm(data, vma = 0, byte = True, offset = True, instructions = True):
         lines.append(line)
 
     return re.sub(',([^ ])', r', \1', '\n'.join(lines))
+
+@LocalContext
+def get_cs_disassembler(eabi=None):
+    import capstone as cs
+    E = {
+        'big':    cs.CS_MODE_BIG_ENDIAN,
+        'little': cs.CS_MODE_LITTLE_ENDIAN,
+    }[context.endianness]
+
+    B = {16: cs.CS_MODE_16, 32: cs.CS_MODE_32, 64: cs.CS_MODE_64}[context.bits]
+
+    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),
+        'armhf'  : (cs.CS_ARCH_ARM, cs.CS_MODE_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),
+        'sparc':   (cs.CS_ARCH_SPARC, cs.CS_MODE_32 + E),
+        'sparc64': (cs.CS_ARCH_SPARC, cs.CS_MODE_64 + E),
+        '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),
+        #'ia64': None,
+        #'sysz': cs.CS_ARCH_SYSZ,
+        #'m68k': cs.CS_ARCH_M68K,
+        #'xcore': cs.CS_ARCH_XCORE,
+        #'tms320c64x': cs.CS_ARCH_TMS320C64X,
+        #'m680x': cs.CS_ARCH_M680X,
+        #'evm': cs.CS_ARCH_EVM,
+        #'mos65xx': cs.CS_ARCH_MOS65XX,
+        #'bpf': cs.CS_ARCH_BPF,
+        #'riscv': cs.CS_ARCH_RISCV,
+        #'tricore': cs.CS_ARCH_TRICORE,
+        #'wasm': cs.CS_ARCH_WASM,
+        #'sh': cs.CS_ARCH_SH,
+    }
+
+    arch = context.arch
+    if arch == 'arm' and eabi == 'hf': arch = 'armhf'
+    param = params.get(arch)
+    if not param:
+        raise Exception(f"unsupported {context.arch} for capstone")
+    arch, mode = param
+    md = cs.Cs(arch, mode)
+    md.detail = True
+    return md

pwnlib/elf/elf.py

@@ -54,6 +54,7 @@
 import re
 import subprocess
 import tempfile
+import capstone as cs
 
 from io import BytesIO
 
@@ -1143,6 +1144,12 @@ def _populate_kernel_version(self):
 
         self.config['version'] = self.version
 
+    def cs_disasm(self, md: cs.Cs, address, n_bytes):
+        if self.arch == 'arm' and address & 1:
+            address -= 1
+
+        return md.disasm(self.read(address, n_bytes), address)
+
     @property
     def libc_start_main_return(self):
         """:class:`int`: Address of the return address into __libc_start_main from main.
@@ -1163,56 +1170,66 @@ def libc_start_main_return(self):
         if 'exit' not in self.symbols:
             return 0
 
+        eabi = None
         # If there's no delay slot, execution continues on the next instruction after a call.
         call_return_offset = 1
         if self.arch in ['arm', 'thumb']:
-            call_instructions = set(['blx', 'bl'])
+            if b'armhf' in self.linker:
+                eabi = 'hf'
+            call_instructions = set([cs.CS_GRP_CALL])
         elif self.arch == 'aarch64':
-            call_instructions = set(['blr', 'bl'])
+            call_instructions = set([cs.CS_GRP_CALL])
         elif self.arch in ['mips', 'mips64']:
-            call_instructions = set(['bal', 'jalr'])
+            # FIXME: `bal` was not included in CS_GRP_CALL. This is fixed on capstone v6.alpha
+            #call_instructions = set([cs.CS_GRP_CALL])
+            call_instructions = set([cs.CS_GRP_CALL, cs.CS_GRP_BRANCH_RELATIVE])
             # Account for the delay slot.
             call_return_offset = 2
         elif self.arch in ['i386', 'amd64', 'ia64']:
-            call_instructions = set(['call'])
+            call_instructions = set([cs.CS_GRP_CALL])
         else:
             log.error('Unsupported architecture %s in ELF.libc_start_main_return', self.arch)
             return 0
 
-        lines = self.functions['__libc_start_main'].disasm().split('\n')
-        exit_addr = hex(self.symbols['exit'])
-        calls = [(index, line) for index, line in enumerate(lines) if set(line.split()) & call_instructions]
+        from pwnlib.asm import get_cs_disassembler
+        md = get_cs_disassembler(arch=self.arch, endian=self.endian, bits=self.bits, eabi=eabi)
+        func = self.functions['__libc_start_main']
+        dis = list(self.cs_disasm(md, func.address, func.size))
 
-        def find_ret_main_addr(lines, calls):
-            exit_calls = [index for index, line in enumerate(calls) if exit_addr in line[1]]
-            if len(exit_calls) != 1:
+        exit_addr = self.symbols['exit']
+        if self.arch == 'arm' and exit_addr & 1: exit_addr -= 1
+
+        calls = [(i, x) for i, x in enumerate(dis) if call_instructions & set(x.groups)]
+
+        def find_ret_main_addr(caller_dis, calls):
+            call_to_main = -1
+            for i, insn in calls:
+                if insn.operands[0].imm == exit_addr: break
+                call_to_main = i
+            else:
                 return 0
 
-            call_to_main = calls[exit_calls[0] - 1]
-            return_from_main = lines[call_to_main[0] + call_return_offset].lstrip()
-            return_from_main = int(return_from_main[ : return_from_main.index(':') ], 16)
-            return return_from_main
+            return_from_main = caller_dis[call_to_main + call_return_offset]
+            return return_from_main.address
 
         # Starting with glibc-2.34 calling `main` is split out into `__libc_start_call_main`
-        ret_addr = find_ret_main_addr(lines, calls)
+        ret_addr = find_ret_main_addr(dis, calls)
         # Pre glibc-2.34 case - `main` is called directly
         if ret_addr:
             return ret_addr
 
         # `__libc_start_main` -> `__libc_start_call_main` -> `main`
         # Find a direct call which calls `exit` once. That's probably `__libc_start_call_main`.
-        direct_call_pattern = re.compile(r'['+r'|'.join(call_instructions)+r']\s+(0x[0-9a-zA-Z]+)')
-        for line in calls:
-            match = direct_call_pattern.search(line[1])
-            if not match:
-                continue
+        for _, insn in calls:
+            op = insn.operands[0]
+            if op.type != cs.CS_OP_IMM: continue
 
-            target_addr = int(match.group(1), 0)
+            target_addr = op.imm
             # `__libc_start_call_main` is usually smaller than `__libc_start_main`, so
             # we might disassemble a bit too much, but it's a good dynamic estimate.
-            callee_lines = self.disasm(target_addr, self.functions['__libc_start_main'].size).split('\n')
-            callee_calls = [(index, line) for index, line in enumerate(callee_lines) if set(line.split()) & call_instructions]
-            ret_addr = find_ret_main_addr(callee_lines, callee_calls)
+            callee_dis = list(self.cs_disasm(md, target_addr, func.size))
+            callee_calls = [(i, x) for i, x in enumerate(callee_dis) if call_instructions & set(x.groups)]
+            ret_addr = find_ret_main_addr(callee_dis, callee_calls)
             if ret_addr:
                 return ret_addr
         return 0