#!/usr/bin/env python3

from __future__ import annotations

import csv
import re
from pathlib import Path

import ida_bytes
import ida_funcs
import ida_gdl
import ida_idaapi
import idautils
import idc

# 根据项目填写
ROOT = Path("/")
OUT_DIR = ROOT / "outputs"
TABLE_ADDR = 0x119A68
TABLE_END = 0x151EE0
TABLE_GETTER = 0x56F58
CACHE_GETTER = 0x56F6C


def qname(ea: int) -> str:
    name = idc.get_name(ea, idc.GN_VISIBLE)
    return name or f"sub_{ea:X}"


def iter_heads(start: int, end: int):
    ea = start
    while ea != ida_idaapi.BADADDR and ea < end:
        yield ea
        ea = idc.next_head(ea, end)


def norm_reg(s: str) -> str:
    s = s.upper().strip()
    s = s.replace(" ", "")
    if re.fullmatch(r"[WX]\d+", s):
        return "X" + s[1:]
    if s in {"WZR", "XZR"}:
        return "ZR"
    return s


def reg_width(s: str) -> int:
    s = s.upper().strip()
    if s.startswith("W"):
        return 32
    if s.startswith("X"):
        return 64
    return 64


def mask_for_width(width: int) -> int:
    return (1 << width) - 1


def dest_reg(ea: int) -> str | None:
    if idc.get_operand_type(ea, 0) != idc.o_reg:
        return None
    return norm_reg(idc.print_operand(ea, 0))


def set_reg(regs: dict[str, int | None], op_text: str, val: int | None, width: int | None = None):
    raw = op_text.upper().strip().replace(" ", "")
    reg = norm_reg(raw)
    if not re.fullmatch(r"X\d+", reg):
        return
    if val is not None:
        if width is None:
            width = reg_width(raw)
        val &= mask_for_width(width)
    regs[reg] = val


def get_reg(regs: dict[str, int | None], op_text: str) -> int | None:
    reg = norm_reg(op_text)
    if reg == "ZR":
        return 0
    return regs.get(reg)


def invalid_written_reg(regs: dict[str, int | None], ea: int):
    dst = dest_reg(ea)
    if dst and re.fullmatch(r"X\d+", dst):
        regs[dst] = None


def sim_one(regs: dict[str, int | None], ea: int):
    m = idc.print_insn_mnem(ea).upper()
    op0 = idc.print_operand(ea, 0)
    op1 = idc.print_operand(ea, 1)
    op2 = idc.print_operand(ea, 2)
    op3 = idc.print_operand(ea, 3)
    dst = dest_reg(ea)

    if ida_idaapi.BADADDR == ea:
        return

    if m in {"BL", "BLR"}:
        # A normal call clobbers argument/result registers.
        for i in range(8):
            regs[f"X{i}"] = None
        return

    if m in {"STR", "STP", "CBZ", "CBNZ", "TBZ", "TBNZ", "CMP", "TST", "B", "B.EQ", "B.NE", "B.CC", "B.CS", "B.LO", "B.HS", "B.LT", "B.GE", "B.LE", "B.GT", "RET"}:
        return

    if dst is None:
        return

    width = reg_width(op0)
    mask = mask_for_width(width)

    if m in {"MOV", "MOVZ"}:
        if idc.get_operand_type(ea, 1) in {idc.o_imm, idc.o_near}:
            val = idc.get_operand_value(ea, 1)
            set_reg(regs, op0, val, width)
            return
        if idc.get_operand_type(ea, 1) == idc.o_reg:
            val = get_reg(regs, op1)
            set_reg(regs, op0, val, width)
            return
    elif m == "MOVN":
        if idc.get_operand_type(ea, 1) == idc.o_imm:
            val = (~idc.get_operand_value(ea, 1)) & mask
            set_reg(regs, op0, val, width)
            return
    elif m == "MOVK":
        if idc.get_operand_type(ea, 1) == idc.o_imm:
            old = get_reg(regs, op0)
            if old is None:
                set_reg(regs, op0, None, width)
                return
            shift = 0
            tail = (op2 or "").upper().replace(" ", "")
            m_shift = re.search(r"LSL#?0X?([0-9A-F]+)", tail)
            if m_shift:
                shift = int(m_shift.group(1), 16)
            else:
                m_shift = re.search(r"LSL#([0-9]+)", tail)
                if m_shift:
                    shift = int(m_shift.group(1), 10)
            imm = idc.get_operand_value(ea, 1) & 0xFFFF
            val = (old & ~(0xFFFF << shift)) | (imm << shift)
            set_reg(regs, op0, val, width)
            return
    elif m in {"ADD", "SUB", "EOR", "ORR", "AND"}:
        left = get_reg(regs, op1)
        right = None
        if idc.get_operand_type(ea, 2) == idc.o_imm:
            right = idc.get_operand_value(ea, 2)
        elif idc.get_operand_type(ea, 2) == idc.o_reg:
            right = get_reg(regs, op2)
        if left is not None and right is not None:
            if m == "ADD":
                val = left + right
            elif m == "SUB":
                val = left - right
            elif m == "EOR":
                val = left ^ right
            elif m == "ORR":
                val = left | right
            else:
                val = left & right
            set_reg(regs, op0, val, width)
            return
    elif m in {"LDR", "LDRB", "LDRH", "LDRSW", "ADRP", "ADR", "SXTW", "SXTB", "UXTB", "UBFX", "BFI", "CSEL", "CSET", "MUL", "LSL", "LSR", "ASR"}:
        invalid_written_reg(regs, ea)
        return

    invalid_written_reg(regs, ea)


def find_block(func, ea: int):
    for block in ida_gdl.FlowChart(func):
        if block.start_ea <= ea < block.end_ea:
            return block
    return None


def recover_w0_in_block(call_ea: int) -> int | None:
    func = ida_funcs.get_func(call_ea)
    if not func:
        return None
    block = find_block(func, call_ea)
    if not block:
        return None
    regs: dict[str, int | None] = {}
    ea = block.start_ea
    while ea != ida_idaapi.BADADDR and ea < call_ea:
        sim_one(regs, ea)
        ea = idc.next_head(ea, block.end_ea)
    val = regs.get("X0")
    if val is None:
        return None
    return val & 0xFFFFFFFF


def recover_w0_backward(call_ea: int, limit: int = 80) -> int | None:
    """Fallback for simple immediate moves immediately before the call."""
    func = ida_funcs.get_func(call_ea)
    if not func:
        return None
    wanted = "X0"
    ea = idc.prev_head(call_ea, func.start_ea)
    seen = 0
    while ea != ida_idaapi.BADADDR and ea >= func.start_ea and seen < limit:
        seen += 1
        m = idc.print_insn_mnem(ea).upper()
        if m in {"BL", "BLR", "RET"}:
            return None
        dst = dest_reg(ea)
        if dst == wanted:
            op0 = idc.print_operand(ea, 0)
            op1 = idc.print_operand(ea, 1)
            if m in {"MOV", "MOVZ"} and idc.get_operand_type(ea, 1) in {idc.o_imm, idc.o_near}:
                return idc.get_operand_value(ea, 1) & 0xFFFFFFFF
            if m == "MOV" and idc.get_operand_type(ea, 1) == idc.o_reg:
                wanted = norm_reg(op1)
            else:
                return None
        ea = idc.prev_head(ea, func.start_ea)
    return None


def recover_w0(call_ea: int) -> int | None:
    val = recover_w0_in_block(call_ea)
    if val is not None:
        return val
    return recover_w0_backward(call_ea)


def decoder_constants(func_start: int, func_end: int) -> tuple[int, int] | None:
    known: dict[str, int | None] = {}
    for ea in iter_heads(func_start, func_end):
        m = idc.print_insn_mnem(ea).upper()
        op0 = idc.print_operand(ea, 0)
        op1 = idc.print_operand(ea, 1)
        op2 = idc.print_operand(ea, 2)

        if m in {"MOV", "MOVZ"} and dest_reg(ea):
            if idc.get_operand_type(ea, 1) == idc.o_imm:
                known[norm_reg(op0)] = idc.get_operand_value(ea, 1)
            elif idc.get_operand_type(ea, 1) == idc.o_reg:
                known[norm_reg(op0)] = known.get(norm_reg(op1))
            else:
                known[norm_reg(op0)] = None
        elif m == "MOVK" and dest_reg(ea):
            reg = norm_reg(op0)
            old = known.get(reg)
            if old is None:
                known[reg] = None
            elif idc.get_operand_type(ea, 1) == idc.o_imm:
                shift = 0
                tail = (op2 or "").upper().replace(" ", "")
                m_shift = re.search(r"LSL#?0X?([0-9A-F]+)", tail)
                if m_shift:
                    shift = int(m_shift.group(1), 16)
                else:
                    m_shift = re.search(r"LSL#([0-9]+)", tail)
                    if m_shift:
                        shift = int(m_shift.group(1), 10)
                imm = idc.get_operand_value(ea, 1) & 0xFFFF
                known[reg] = (old & ~(0xFFFF << shift)) | (imm << shift)
        elif dest_reg(ea) and m not in {"ADD", "EOR"}:
            known[norm_reg(op0)] = None

        if m != "ADD":
            continue
        acc = norm_reg(op0)
        if acc != norm_reg(op1):
            continue
        if idc.get_operand_type(ea, 2) != idc.o_reg:
            continue
        n1 = idc.next_head(ea, func_end)
        n2 = idc.next_head(n1, func_end)
        m1 = idc.print_insn_mnem(n1).upper()
        m2 = idc.print_insn_mnem(n2).upper()
        if (
            m1 == "EOR"
            and m2 == "ADD"
            and norm_reg(idc.print_operand(n1, 0)) == acc
            and norm_reg(idc.print_operand(n1, 1)) == acc
            and norm_reg(idc.print_operand(n2, 0)) == acc
            and norm_reg(idc.print_operand(n2, 1)) == acc
            and idc.get_operand_type(n2, 2) == idc.o_imm
        ):
            if idc.get_operand_type(n1, 2) == idc.o_imm:
                xor_const = idc.get_operand_value(n1, 2)
            elif idc.get_operand_type(n1, 2) == idc.o_reg:
                xor_reg = norm_reg(idc.print_operand(n1, 2))
                xor_const = known.get(xor_reg)
                if xor_const is None:
                    continue
            else:
                continue
            add_const = idc.get_operand_value(n2, 2)
            return xor_const & 0xFF, add_const & 0xFF
        if (
            m1 == "ADD"
            and norm_reg(idc.print_operand(n1, 0)) == acc
            and norm_reg(idc.print_operand(n1, 1)) == acc
            and idc.get_operand_type(n1, 2) == idc.o_imm
        ):
            add_const = idc.get_operand_value(n1, 2)
            return 0, add_const & 0xFF
    return None


def decode_at(table: bytes, off: int, xor_const: int, add_const: int) -> str:
    if off < 0 or off + 2 >= len(table):
        raise ValueError("offset outside table")
    key = table[off]
    size = table[off + 1] ^ key
    if off + 2 + size >= len(table):
        raise ValueError("record outside table")
    out = bytearray()
    rolling = key
    for j in range(size):
        out.append(table[off + 2 + j] ^ rolling)
        rolling = ((rolling + j) ^ xor_const) + add_const
        rolling &= 0xFF
    chk = 0xFF
    for c in out:
        chk ^= c
    expected = (~chk) & 0xFF
    actual = key ^ table[off + 2 + size]
    if actual != expected:
        raise ValueError(f"checksum mismatch actual=0x{actual:02x} expected=0x{expected:02x}")
    return out.decode("utf-8", errors="replace")


def cref_from_func_to(target: int, func_start: int, func_end: int) -> bool:
    for ea in iter_heads(func_start, func_end):
        for ref in idautils.CodeRefsFrom(ea, False):
            if ref == target:
                return True
    return False


def discover_decoders():
    decoders = []
    skipped = []
    for call_ea in sorted(idautils.CodeRefsTo(TABLE_GETTER, False)):
        func = ida_funcs.get_func(call_ea)
        if not func:
            skipped.append({"call_ea": call_ea, "reason": "no function"})
            continue
        if not cref_from_func_to(CACHE_GETTER, func.start_ea, func.end_ea):
            skipped.append({
                "call_ea": call_ea,
                "func": func.start_ea,
                "name": qname(func.start_ea),
                "reason": "uses table getter but not cache getter; not a string decoder wrapper",
            })
            continue
        consts = decoder_constants(func.start_ea, func.end_ea)
        if consts is None:
            skipped.append({
                "call_ea": call_ea,
                "func": func.start_ea,
                "name": qname(func.start_ea),
                "reason": "decoder constants not found",
            })
            continue
        decoders.append({
            "func": func.start_ea,
            "name": qname(func.start_ea),
            "call_to_table_getter": call_ea,
            "xor_const": consts[0],
            "add_const": consts[1],
        })
    uniq = {}
    for d in decoders:
        uniq[d["func"]] = d
    return [uniq[k] for k in sorted(uniq)], skipped


def find_callsites(decoders, table: bytes):
    rows = []
    for d in decoders:
        target = d["func"]
        for call_ea in sorted(idautils.CodeRefsTo(target, False)):
            caller = ida_funcs.get_func(call_ea)
            if not caller:
                continue
            mnem = idc.print_insn_mnem(call_ea).upper()
            if mnem not in {"BL", "B"}:
                continue
            off = recover_w0(call_ea)
            caller_name = qname(caller.start_ea)
            row_base = {
                "call_ea": call_ea,
                "caller_func": caller.start_ea,
                "caller_name": caller_name,
                "decoder_func": d["func"],
                "decoder_name": d["name"],
                "xor_const": d["xor_const"],
                "add_const": d["add_const"],
            }
            if off is None:
                continue
            try:
                text = decode_at(table, off, d["xor_const"], d["add_const"])
                rows.append({
                    **row_base,
                    "offset": off,
                    "table_va": TABLE_ADDR + off,
                    "string": text,
                })
            except Exception:
                continue
    rows.sort(key=lambda r: (r["caller_func"], r["call_ea"], r["decoder_func"], r.get("offset", -1)))
    return rows


def write_csv(rows):
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    csv_path = OUT_DIR / "anogs_obf_string_locations.csv"
    fields = [
        "call_ea",
        "caller_func",
        "caller_name",
        "decoder_func",
        "decoder_name",
        "xor_const",
        "add_const",
        "offset",
        "table_va",
        "string",
    ]
    with csv_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for row in rows:
            w.writerow({k: row.get(k, "") for k in fields})


def main():
    table = ida_bytes.get_bytes(TABLE_ADDR, TABLE_END - TABLE_ADDR)
    if table is None:
        raise RuntimeError(f"failed to read table bytes at 0x{TABLE_ADDR:X}")

    decoders, _ = discover_decoders()
    rows = find_callsites(decoders, table)
    write_csv(rows)


if __name__ == "__main__":
    main()