tools: Deduplicate code

tools/util/elf.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+
+from typing import Any, Dict, Set
+
+import capstone as cs
+from elftools.elf.elffile import ELFFile
+
+import diff_settings
+from util import utils
+
+_config: Dict[str, Any] = {}
+diff_settings.apply(_config, {})
+
+_root = utils.get_repo_root()
+
+base_elf = ELFFile((_root / _config["baseimg"]).open("rb"))
+my_elf = ELFFile((_root / _config["myimg"]).open("rb"))
+my_symtab = my_elf.get_section_by_name(".symtab")
+if not my_symtab:
+    utils.fail(f'{_config["myimg"]} has no symbol table')
+
+md = cs.Cs(cs.CS_ARCH_ARM64, cs.CS_MODE_ARM)
+md.detail = True
+
+
+def get_file_offset(elf, addr: int) -> int:
+    for seg in elf.iter_segments():
+        if seg.header["p_type"] != "PT_LOAD":
+            continue
+        if seg["p_vaddr"] <= addr < seg["p_vaddr"] + seg["p_filesz"]:
+            return addr - seg["p_vaddr"] + seg["p_offset"]
+    assert False
+
+
+def get_symbol_file_offset_and_size(elf, table, name: str) -> (int, int):
+    syms = table.get_symbol_by_name(name)
+    if not syms or len(syms) != 1:
+        raise KeyError(name)
+    return get_file_offset(elf, syms[0]["st_value"]), syms[0]["st_size"]
+
+
+def build_symbol_table(symtab) -> Dict[int, str]:
+    table = dict()
+    for sym in symtab.iter_symbols():
+        addr = sym["st_value"]
+        existing_value = table.get(addr, None)
+        if existing_value is None or not existing_value.startswith("_Z"):
+            table[addr] = sym.name
+    return table
+
+
+def get_fn_from_base_elf(addr: int, size: int) -> bytes:
+    offset = get_file_offset(base_elf, addr)
+    base_elf.stream.seek(offset)
+    return base_elf.stream.read(size)
+
+
+def get_fn_from_my_elf(name: str) -> bytes:
+    offset, size = get_symbol_file_offset_and_size(my_elf, my_symtab, name)
+    my_elf.stream.seek(offset)
+    return my_elf.stream.read(size)
+
+
+def check_function_ex(addr: int, size: int, base_fn: bytes, my_fn: bytes) -> bool:
+    adrp_pair_registers: Set[int] = set()
+
+    for i1, i2 in zip(md.disasm(base_fn, addr), md.disasm(my_fn, addr)):
+        if i1.bytes == i2.bytes:
+            if i1.mnemonic == 'adrp':
+                adrp_pair_registers.add(i1.operands[0].reg)
+            elif i1.mnemonic == 'ldr':
+                reg = i1.operands[1].value.mem.base
+                if reg in adrp_pair_registers:
+                    adrp_pair_registers.remove(reg)
+            elif i1.mnemonic == 'ldp':
+                reg = i1.operands[2].value.mem.base
+                if reg in adrp_pair_registers:
+                    adrp_pair_registers.remove(reg)
+            elif i1.mnemonic == 'add':
+                reg = i1.operands[1].reg
+                if reg in adrp_pair_registers:
+                    adrp_pair_registers.remove(reg)
+            continue
+
+        if i1.mnemonic != i2.mnemonic:
+            return False
+
+        # Ignore some address differences until a fully matching executable can be generated.
+
+        if i1.mnemonic == 'bl':
+            continue
+
+        if i1.mnemonic == 'b':
+            # Needed for tail calls.
+            branch_target = int(i1.op_str[1:], 16)
+            if not (addr <= branch_target < addr + size):
+                continue
+
+        if i1.mnemonic == 'adrp':
+            if i1.operands[0].reg != i2.operands[0].reg:
+                return False
+            adrp_pair_registers.add(i1.operands[0].reg)
+            continue
+
+        if i1.mnemonic == 'ldr' or i1.mnemonic == 'str':
+            if i1.operands[0].reg != i2.operands[0].reg:
+                return False
+            if i1.operands[1].value.mem.base != i2.operands[1].value.mem.base:
+                return False
+            reg = i1.operands[1].value.mem.base
+            if reg not in adrp_pair_registers:
+                return False
+            adrp_pair_registers.remove(reg)
+            continue
+
+        if i1.mnemonic == 'ldp' or i1.mnemonic == 'stp':
+            if i1.operands[0].reg != i2.operands[0].reg:
+                return False
+            if i1.operands[1].reg != i2.operands[1].reg:
+                return False
+            if i1.operands[2].value.mem.base != i2.operands[2].value.mem.base:
+                return False
+            reg = i1.operands[2].value.mem.base
+            if reg not in adrp_pair_registers:
+                return False
+            adrp_pair_registers.remove(reg)
+            continue
+
+        if i1.mnemonic == 'add':
+            if i1.operands[0].reg != i2.operands[0].reg:
+                return False
+            if i1.operands[1].reg != i2.operands[1].reg:
+                return False
+            reg = i1.operands[1].reg
+            if reg not in adrp_pair_registers:
+                return False
+            adrp_pair_registers.remove(reg)
+            continue
+
+        return False
+
+    return True

tools/util/utils.py

@@ -0,0 +1,94 @@
+from colorama import Fore, Style
+import csv
+import cxxfilt
+import enum
+from pathlib import Path
+import sys
+import typing as tp
+
+
+class FunctionStatus(enum.Enum):
+    Matching = 0
+    Equivalent = 1  # semantically equivalent but not perfectly matching
+    NonMatching = 2
+    Wip = 3
+    NotDecompiled = 4
+
+
+class FunctionInfo(tp.NamedTuple):
+    addr: int  # without the 0x7100000000 base
+    name: str
+    size: int
+    decomp_name: str
+    status: FunctionStatus
+
+
+_markers = {
+    "?": FunctionStatus.Equivalent,
+    "!": FunctionStatus.NonMatching,
+    "|": FunctionStatus.Wip,
+}
+
+
+def parse_function_csv_entry(row) -> FunctionInfo:
+    ea, name, size, decomp_name = row
+    if decomp_name:
+        status = FunctionStatus.Matching
+
+        for marker, new_status in _markers.items():
+            if decomp_name[-1] == marker:
+                status = new_status
+                decomp_name = decomp_name[:-1]
+                break
+    else:
+        status = FunctionStatus.NotDecompiled
+
+    addr = int(ea, 16) - 0x7100000000
+    return FunctionInfo(addr, name, int(size, 0), decomp_name, status)
+
+
+def get_functions(path: tp.Optional[Path] = None) -> tp.Iterable[FunctionInfo]:
+    if path is None:
+        path = get_repo_root() / "data" / "uking_functions.csv"
+    with path.open() as f:
+        for row in csv.reader(f):
+            yield parse_function_csv_entry(row)
+
+
+def format_symbol_name(name: str) -> str:
+    try:
+        return f"{cxxfilt.demangle(name)} {Style.DIM}({name}){Style.RESET_ALL}"
+    except:
+        return name
+
+
+def format_symbol_name_for_msg(name: str) -> str:
+    try:
+        return f"{Fore.BLUE}{cxxfilt.demangle(name)}{Fore.RESET} {Style.DIM}({name}){Style.RESET_ALL}{Style.BRIGHT}"
+    except:
+        return name
+
+
+def are_demangled_names_equal(name1: str, name2: str):
+    return cxxfilt.demangle(name1) == cxxfilt.demangle(name2)
+
+
+def print_note(msg: str, prefix: str = ""):
+    sys.stderr.write(f"{Style.BRIGHT}{prefix}{Fore.CYAN}note:{Fore.RESET} {msg}{Style.RESET_ALL}\n")
+
+
+def warn(msg: str, prefix: str = ""):
+    sys.stderr.write(f"{Style.BRIGHT}{prefix}{Fore.MAGENTA}warning:{Fore.RESET} {msg}{Style.RESET_ALL}\n")
+
+
+def print_error(msg: str, prefix: str = ""):
+    sys.stderr.write(f"{Style.BRIGHT}{prefix}{Fore.RED}error:{Fore.RESET} {msg}{Style.RESET_ALL}\n")
+
+
+def fail(msg: str, prefix: str = ""):
+    print_error(msg, prefix)
+    sys.exit(1)
+
+
+def get_repo_root() -> Path:
+    return Path(__file__).parent.parent.parent