#!/usr/bin/env python3
"""hiby_player MIPS32 static analysis helpers (capstone-based)."""
import sys, struct, subprocess

CAP_VER = None
try:
    from capstone import Cs, CS_ARCH_MIPS, CS_MODE_MIPS32, CS_MODE_LITTLE_ENDIAN, CS_OP_MEM, CS_OP_IMM, CS_OP_REG
    CAP_VER = "capstone"
except ImportError:
    CAP_VER = None

BIN = "/home/Aska/Project/R3ProII/unpacked_v14/squashfs-root/usr/bin/hiby_player"

def load_segments():
    """Return list of (vaddr, size, bytes, name) for executable/rodata segments."""
    with open(BIN, "rb") as f:
        data = f.read()
    segs = []
    out = subprocess.run(["readelf", "-lW", BIN], capture_output=True, text=True).stdout
    for line in out.splitlines():
        parts = line.split()
        if len(parts) >= 9 and parts[0] == "LOAD":
            # Type Off VAddr PAddr FileSiz MemSiz Flg Align  (Flg may split into "R E")
            off, vaddr = int(parts[1], 16), int(parts[2], 16)
            filesz = int(parts[4], 16)
            flags = "".join(parts[6:8])
            segs.append((vaddr, filesz, data[off:off + filesz], flags))
    return segs

def text_ranges():
    """Executable section ranges: [(vaddr, size)] for .init/.text/.MIPS.stubs/.fini."""
    out = []
    res = subprocess.run(["readelf", "-SW", BIN], capture_output=True, text=True).stdout
    for line in res.splitlines():
        parts = line.split()
        if len(parts) >= 7 and parts[1] in (".init", ".text", ".MIPS.stubs", ".fini"):
            vaddr, size = int(parts[3], 16), int(parts[5], 16)
            out.append((vaddr, size))
    return out

def disasm_all():
    """Full disassembly of executable sections -> list of (addr, mnem, op_str).
    Skips invalid bytes by stepping 4 bytes (MIPS instructions are 4-byte aligned)."""
    insns = []
    md = Cs(CS_ARCH_MIPS, CS_MODE_MIPS32 + CS_MODE_LITTLE_ENDIAN)
    md.detail = True
    with open(BIN, "rb") as f:
        data = f.read()
    for vaddr, size in text_ranges():
        off = vaddr - 0x400000
        blob = data[off:off + size]
        addr = vaddr
        end = vaddr + size
        while addr < end:
            try:
                got = list(md.disasm(blob[addr - vaddr:addr - vaddr + 64], addr, count=1))
            except Exception:
                got = []
            if not got:
                addr += 4
                continue
            ins = got[0]
            insns.append((ins.address, ins.mnemonic, ins.op_str))
            addr += ins.size
    return insns

def strings_with_addr(minlen=5):
    """Extract ASCII strings with their virtual addresses from all segments."""
    import re
    out = []
    for vaddr, size, blob, flags in load_segments():
        for m in re.finditer(rb"[\x20-\x7e]{%d,}" % minlen, blob):
            out.append((vaddr + m.start(), m.group().decode()))
    return out

def find_xrefs(target_vaddrs, insns=None, window=12):
    """Find instructions that reference any of target_vaddrs via lui+addiu or jalr patterns.
    Returns list of (ref_insn_addr, ref_insn_str, target_str, disasm_window)."""
    if insns is None:
        insns = disasm_all()
    results = []
    targets = {v: s for v, s in target_vaddrs}
    # MIPS address load: lui reg, hi16 ; addiu/ori reg, reg, lo16
    # Build map: reg -> (insn_addr, value_hi) after lui
    # We'll scan for lui whose hi16 matches, then look ahead for addiu with lo16.
    for i, (addr, mnem, op) in enumerate(insns):
        if mnem != "lui":
            continue
        try:
            reg, imm = op.replace(" ", "").split(",")
            hi = int(imm, 16) << 16
        except Exception:
            continue
        for j in range(i + 1, min(i + window, len(insns))):
            a2, m2, o2 = insns[j]
            if m2 in ("addiu", "ori") and o2.startswith(reg + ","):
                try:
                    raw = o2.split(",")[-1].strip()
                    lo = int(raw, 16) if raw.startswith(("0x", "0X")) else int(raw, 0)
                except Exception:
                    continue
                full = (hi + (lo if lo < 0x8000 else lo - 0x10000)) if m2 == "addiu" else hi + lo
                if full in targets:
                    ctx = " | ".join(f"{x:08x}: {mn} {op}" for x, mn, op in insns[max(0, i - 3):i + 3])
                    results.append((addr, f"lui {reg}, 0x{hi >> 16:x}", targets[full], ctx))
    return results

def dump(insns, start, count=40):
    for a, m, o in insns:
        if a >= start:
            print(f"{a:08x}: {m:8s} {o}")
            count -= 1
            if count <= 0:
                break

if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "info"
    if cmd == "xref":
        # usage: xref <string-substr> [n]
        substr = sys.argv[2]
        n = int(sys.argv[3]) if len(sys.argv) > 3 else 30
        sw = strings_with_addr(5)
        hits = [(v, s) for v, s in sw if substr in s]
        print(f"strings matching '{substr}': {len(hits)}")
        for v, s in hits[:20]:
            print(f"  {v:08x}: {s[:90]}")
        insns = disasm_all()
        print(f"\nxrefs ({len(insns)} instructions):")
        for ref_addr, ref_str, tstr, ctx in find_xrefs(hits, insns)[:n]:
            print(f"\n@ {ref_addr:08x}  [{ref_str}]  -> '{tstr[:70]}'")
            print(f"   ctx: {ctx}")
    elif cmd == "dump":
        start = int(sys.argv[2], 16)
        insns = disasm_all()
        dump(insns, start, int(sys.argv[3]) if len(sys.argv) > 3 else 40)
    elif cmd == "info":
        segs = load_segments()
        for v, s, b, f in segs:
            print(f"LOAD vaddr={v:#x} filesz={s:#x} flags={f}")