from __future__ import annotations

import argparse
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

from capstone import Cs, CS_ARCH_ARM64, CS_MODE_ARM
from capstone.arm64 import ARM64_OP_IMM
from elftools.elf.elffile import ELFFile
from elftools.elf.relocation import RelocationSection


@dataclass(frozen=True)
class Symbol:
    name: str
    addr: int
    size: int
    bind: str
    typ: str


class ElfAnalyzer:
    def __init__(self, path: str | Path):
        self.path = Path(path)
        self._fp = self.path.open("rb")
        self.elf = ELFFile(self._fp)
        self._symbols = self._collect_symbols()
        self._symbols_by_name = {s.name: s for s in self._symbols}
        self._symbols_by_addr = {s.addr: s for s in self._symbols if s.addr}
        self._plt_symbols = self._collect_plt_symbols()
        self._cs = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
        self._cs.detail = True

    def close(self) -> None:
        self._fp.close()

    def _collect_symbols(self) -> List[Symbol]:
        out: List[Symbol] = []
        for section in self.elf.iter_sections():
            if section["sh_type"] not in ("SHT_SYMTAB", "SHT_DYNSYM"):
                continue
            for sym in section.iter_symbols():
                name = sym.name
                if not name:
                    continue
                info = sym["st_info"]
                out.append(
                    Symbol(
                        name=name,
                        addr=int(sym["st_value"]),
                        size=int(sym["st_size"]),
                        bind=str(info["bind"]),
                        typ=str(info["type"]),
                    )
                )
        out.sort(key=lambda s: (s.addr, s.name))
        return out

    def _collect_plt_symbols(self) -> Dict[int, str]:
        # AArch64 ELF PLT layout: the resolver entry is 32 bytes, then one
        # 16-byte entry per relocation in .rela.plt order.
        plt = self.elf.get_section_by_name(".plt")
        if plt is None:
            return {}
        out: Dict[int, str] = {}
        plt_addr = int(plt["sh_addr"])
        entry_addr = plt_addr + 32
        for section in self.elf.iter_sections():
            if not isinstance(section, RelocationSection):
                continue
            if section.name not in (".rela.plt", ".rela.plt.sec", ".rel.plt"):
                continue
            symtab = self.elf.get_section(section["sh_link"])
            for index, rel in enumerate(section.iter_relocations()):
                sym = symtab.get_symbol(rel["r_info_sym"])
                if sym.name:
                    out[entry_addr + index * 16] = sym.name + "@plt"
        return out

    def symbols(self, pattern: Optional[str] = None) -> List[Symbol]:
        if pattern is None:
            return list(self._symbols)
        rx = re.compile(pattern)
        return [s for s in self._symbols if rx.search(s.name)]

    def find_symbol(self, name: str) -> Optional[Symbol]:
        return self._symbols_by_name.get(name)

    def describe_addr(self, addr: int) -> str:
        if addr in self._plt_symbols:
            return self._plt_symbols[addr]
        if addr in self._symbols_by_addr:
            return self._symbols_by_addr[addr].name
        best: Optional[Symbol] = None
        for sym in self._symbols:
            if sym.addr and sym.size and sym.addr <= addr < sym.addr + sym.size:
                if best is None or sym.addr > best.addr:
                    best = sym
        if best is not None:
            return f"{best.name}+0x{addr - best.addr:x}"
        return f"0x{addr:x}"

    def vaddr_to_offset(self, addr: int) -> int:
        for seg in self.elf.iter_segments():
            if seg["p_type"] != "PT_LOAD":
                continue
            start = int(seg["p_vaddr"])
            size = int(seg["p_memsz"])
            if start <= addr < start + size:
                return int(seg["p_offset"]) + (addr - start)
        raise ValueError(f"virtual address 0x{addr:x} is not in a PT_LOAD segment")

    def read_vaddr(self, addr: int, size: int) -> bytes:
        off = self.vaddr_to_offset(addr)
        self._fp.seek(off)
        return self._fp.read(size)

    def disassemble_symbol(self, name: str, max_bytes: Optional[int] = None) -> List[str]:
        sym = self.find_symbol(name)
        if sym is None:
            raise KeyError(name)
        if sym.size <= 0:
            raise ValueError(f"symbol {name} has no size")
        size = min(sym.size, max_bytes) if max_bytes else sym.size
        code = self.read_vaddr(sym.addr, size)
        lines = [f"{name} @ 0x{sym.addr:x} size=0x{sym.size:x}"]
        for insn in self._cs.disasm(code, sym.addr):
            comment = ""
            if insn.mnemonic in {"b", "bl", "cbz", "cbnz", "tbz", "tbnz"}:
                imm = self._first_imm(insn)
                if imm is not None:
                    comment = f" ; {self.describe_addr(imm)}"
            lines.append(f"0x{insn.address:08x}:\t{insn.mnemonic}\t{insn.op_str}{comment}")
        return lines

    @staticmethod
    def _first_imm(insn) -> Optional[int]:
        for op in insn.operands:
            if op.type == ARM64_OP_IMM:
                return int(op.imm)
        return None


def _default_so() -> Path:
    return (
        Path.cwd()
        / "IR_Camera_SDK-1.0.1"
        / "android"
        / "android"
        / "lib"
        / "arm64-v8a"
        / "libcoresdk.so"
    )


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Small ARM64 ELF disassembler for libcoresdk.so")
    parser.add_argument("--elf", default=str(_default_so()))
    parser.add_argument("--symbol", action="append", default=[])
    parser.add_argument("--grep-symbols", help="Regex for listing matching symbols")
    parser.add_argument("--max-bytes", type=lambda s: int(s, 0), default=None)
    args = parser.parse_args(argv)

    analyzer = ElfAnalyzer(args.elf)
    try:
        if args.grep_symbols:
            for sym in analyzer.symbols(args.grep_symbols):
                print(f"0x{sym.addr:08x}\t0x{sym.size:x}\t{sym.typ}\t{sym.name}")
        for name in args.symbol:
            for line in analyzer.disassemble_symbol(name, args.max_bytes):
                print(line)
            print()
    finally:
        analyzer.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())