"""Core IL2CPP reverse-engineering helpers for SAND's GameAssembly.dll.

dump.cs (Il2CppDumper output) gives every method's RVA, file Offset, VA and
signature. GameAssembly.dll holds the actual x86-64 bodies. This module:

  - maps VA <-> file offset via the PE section table,
  - builds a sorted method index from dump.cs (cached to a pickle),
  - attributes any VA to its containing method,
  - finds call-rel32 / jmp-rel32 sites that target a given VA (xrefs),
  - disassembles a method body and pulls out float constants it loads.

No method bodies live in dump.cs — only signatures. The values we want
(weapon damage floats) are operands inside the bodies, recovered here.
"""
import os, re, pickle, bisect, struct

GAME = "/mnt/d/SteamLibrary/steamapps/common/Sand Playtest"
DLL = os.path.join(GAME, "GameAssembly.dll")
DUMP = "/home/downloadpizza/sand_tools/il2cpp/dump.cs"
CACHE = "/home/downloadpizza/sand_tools/reverse/_method_index.pkl"

IMAGE_BASE = 0x180000000  # from PE optional header (VA = IMAGE_BASE + RVA)


class PE:
    """Minimal PE map: VA<->file-offset, plus raw byte access."""
    def __init__(self, path=DLL):
        import pefile
        self.pe = pefile.PE(path, fast_load=True)
        self.base = self.pe.OPTIONAL_HEADER.ImageBase
        self.data = self.pe.__data__  # bytes-like of whole file
        self.sections = []
        for s in self.pe.sections:
            self.sections.append((
                s.VirtualAddress, s.Misc_VirtualSize,
                s.PointerToRawData, s.SizeOfRawData,
                s.Name.decode(errors="replace").strip("\x00"),
            ))

    def rva_to_off(self, rva):
        for va, vsz, praw, rsz, name in self.sections:
            if va <= rva < va + max(vsz, rsz):
                if rva - va < rsz:
                    return praw + (rva - va)
                return None  # in bss / uninitialized
        return None

    def off_to_rva(self, off):
        for va, vsz, praw, rsz, name in self.sections:
            if praw <= off < praw + rsz:
                return va + (off - praw)
        return None

    def va_to_off(self, va):
        return self.rva_to_off(va - self.base)

    def read_off(self, off, n):
        return bytes(self.data[off:off + n])

    def read_va(self, va, n):
        off = self.va_to_off(va)
        if off is None:
            return None
        return self.read_off(off, n)

    def text_range(self):
        for va, vsz, praw, rsz, name in self.sections:
            if name == ".text":
                return va, vsz, praw, rsz
        return None

    def code_ranges(self):
        """Executable code sections. IL2CPP method bodies live in the
        'il2cpp' section; runtime/engine glue in '.text'."""
        out = []
        for va, vsz, praw, rsz, name in self.sections:
            if name in (".text", "il2cpp"):
                out.append((va, vsz, praw, rsz, name))
        return out


_RVA_RE = re.compile(r"// RVA: 0x([0-9A-Fa-f]+) Offset: 0x([0-9A-Fa-f]+) VA: 0x([0-9A-Fa-f]+)")


def build_index(force=False):
    """Parse dump.cs -> list of methods sorted by VA.

    Each entry: dict(va, rva, off, sig). sig is the C#-ish declaration line.
    Cached to CACHE.
    """
    if not force and os.path.exists(CACHE):
        with open(CACHE, "rb") as f:
            return pickle.load(f)
    methods = []
    with open(DUMP, "r", encoding="utf-8", errors="replace") as f:
        prev = None
        for line in f:
            m = _RVA_RE.search(line)
            if m:
                prev = (int(m.group(1), 16), int(m.group(2), 16), int(m.group(3), 16))
                continue
            if prev is not None:
                sig = line.strip()
                rva, off, va = prev
                if rva != 0:
                    methods.append({"va": va, "rva": rva, "off": off, "sig": sig})
                prev = None
    methods.sort(key=lambda d: d["va"])
    with open(CACHE, "wb") as f:
        pickle.dump(methods, f)
    return methods


class Index:
    def __init__(self):
        self.methods = build_index()
        self.vas = [m["va"] for m in self.methods]

    def method_at(self, va):
        """Return the method whose body contains `va` (largest va <= target)."""
        i = bisect.bisect_right(self.vas, va) - 1
        if i < 0:
            return None
        return self.methods[i]

    def next_va(self, va):
        i = bisect.bisect_right(self.vas, va)
        return self.vas[i] if i < len(self.vas) else None

    def find_sig(self, needle):
        return [m for m in self.methods if needle in m["sig"]]


def find_xrefs(pe, target_vas, call_only=True):
    """Scan all code sections for E8 (call) / E9 (jmp) rel32 sites whose
    target is in `target_vas` (int or set/iterable of ints).

    Returns list of (site_va, opcode_byte, target_va). Linear byte scan with
    displacement check; a few false positives are possible but validated by
    the disassembler downstream.
    """
    if isinstance(target_vas, int):
        target_vas = {target_vas}
    else:
        target_vas = set(target_vas)
    out = []
    base = pe.base
    data = pe.data
    for sva, svsz, spraw, srsz, name in pe.code_ranges():
        raw = bytes(data[spraw:spraw + srsz])
        n = len(raw)
        j = 0
        while True:
            j = raw.find(b"\xe8", j) if call_only else _find_either(raw, j)
            if j < 0 or j + 5 > n:
                break
            disp = struct.unpack_from("<i", raw, j + 1)[0]
            site_va = base + sva + j  # off_to_rva is identity here: rva = sva + j
            tgt = site_va + 5 + disp
            if tgt in target_vas:
                out.append((site_va, raw[j], tgt))
            j += 1
    return out


def find_rip_refs(pe, target_va):
    """Find RIP-relative instructions whose disp32 (assumed to be the final 4
    bytes of the instruction) resolves to `target_va`. Returns list of
    (disp_field_va, instr_start_guess_va). Vectorized over all code sections.

    For a load like `mov reg,[rip+disp]`, target = (va_after_instr) + disp and
    va_after_instr = disp_field_va + 4. So target_va = disp_field_va + 4 + disp
    => disp + i == (target_va - base - sva - 4)  where i is the in-section offset
    of the disp field. We solve for i.
    """
    import numpy as np
    out = []
    base = pe.base
    data = pe.data
    for sva, svsz, spraw, srsz, name in pe.code_ranges():
        raw = np.frombuffer(bytes(data[spraw:spraw + srsz]), dtype=np.uint8)
        n = len(raw)
        # little-endian int32 at every byte offset 0..n-4
        disp = (raw[0:n-3].astype(np.int64)
                | (raw[1:n-2].astype(np.int64) << 8)
                | (raw[2:n-1].astype(np.int64) << 16)
                | (raw[3:n].astype(np.int64) << 24))
        # sign-extend 32-bit
        disp = disp - ((disp & 0x80000000) << 1)
        idx = np.arange(n - 3, dtype=np.int64)
        const = target_va - base - sva - 4
        hits = np.nonzero(disp + idx == const)[0]
        for i in hits.tolist():
            out.append(base + sva + int(i))  # va of the disp field
    return out


def find_rip_refs_batch(pe, targets):
    """One pass over code: for a set of target VAs, return {target_va: [disp_field_va,...]}.
    Assumes disp32 is the final 4 bytes of the instruction (movss/mov/lea loads)."""
    import numpy as np
    targets = np.array(sorted(set(int(t) for t in targets)), dtype=np.int64)
    out = {int(t): [] for t in targets}
    base = pe.base
    data = pe.data
    for sva, svsz, spraw, srsz, name in pe.code_ranges():
        raw = np.frombuffer(bytes(data[spraw:spraw + srsz]), dtype=np.uint8)
        n = len(raw)
        disp = (raw[0:n-3].astype(np.int64) | (raw[1:n-2].astype(np.int64) << 8)
                | (raw[2:n-1].astype(np.int64) << 16) | (raw[3:n].astype(np.int64) << 24))
        disp = disp - ((disp & 0x80000000) << 1)
        idx = np.arange(n - 3, dtype=np.int64)
        tgt = base + sva + idx + 4 + disp  # absolute target of a disp-final instruction
        mask = np.isin(tgt, targets)
        for i in np.nonzero(mask)[0]:
            t = int(tgt[i])
            if t in out:
                out[t].append(base + sva + int(i))
    return out


def disasm_around(pe, va, back=24, total=64):
    """Disassemble a window [va-back, va-back+total) to recover the instruction
    that contains/precedes `va` (e.g. a RIP-ref disp field)."""
    from capstone import Cs, CS_ARCH_X86, CS_MODE_64
    start = va - back
    off = pe.va_to_off(start)
    code = pe.read_off(off, total)
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True
    return list(md.disasm(code, start))


def _find_either(raw, start):
    a = raw.find(b"\xe8", start)
    b = raw.find(b"\xe9", start)
    if a < 0:
        return b
    if b < 0:
        return a
    return min(a, b)


def disasm_method(pe, idx, va, max_bytes=0x4000):
    """Disassemble one method body (until next method VA or max_bytes)."""
    from capstone import Cs, CS_ARCH_X86, CS_MODE_64
    nxt = idx.next_va(va)
    size = max_bytes
    if nxt:
        size = min(size, nxt - va)
    off = pe.va_to_off(va)
    code = pe.read_off(off, size)
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True
    return list(md.disasm(code, va))


def rip_target(insn):
    """If an instruction has a RIP-relative memory operand, return its absolute VA."""
    from capstone import x86
    for op in insn.operands:
        if op.type == x86.X86_OP_MEM and op.mem.base == x86.X86_REG_RIP:
            return insn.address + insn.size + op.mem.disp
    return None


def read_f32(pe, va):
    b = pe.read_va(va, 4)
    return None if b is None else struct.unpack("<f", b)[0]


def read_f64(pe, va):
    b = pe.read_va(va, 8)
    return None if b is None else struct.unpack("<d", b)[0]


def scan_movss_consts(pe, lo=None, hi=None):
    """Linear-scan code sections for `movss xmm,[rip+disp32]` (F3 0F 10 /r, mod=00 rm=101)
    and `movsd` (F2 0F 10). Returns list of (insn_va, const_value, kind). Optional
    [lo,hi] filters the constant's absolute value. Anchor-free way to find where
    float constants (e.g. damage numbers) are loaded, regardless of Entitas plumbing."""
    import numpy as np, struct as _s
    out = []
    base = pe.base
    data = pe.data
    rm_ok = set(0x05 | (r << 3) for r in range(8))  # mod=00, rm=101, reg=0..7
    for sva, svsz, spraw, srsz, name in pe.code_ranges():
        raw = bytes(data[spraw:spraw + srsz])
        n = len(raw)
        # find F3 0F 10 and F2 0F 10
        for pref, kind in ((b"\xf3\x0f\x10", "f32"), (b"\xf2\x0f\x10", "f64")):
            j = raw.find(pref)
            while j >= 0:
                if j + 8 <= n:
                    modrm = raw[j + 3]
                    if modrm in rm_ok:
                        disp = _s.unpack_from("<i", raw, j + 4)[0]
                        insn_va = base + sva + j
                        tgt = insn_va + 8 + disp  # instruction length = 8 (F3 0F 10 modrm disp32)
                        if kind == "f32":
                            b = pe.read_va(tgt, 4); v = _s.unpack("<f", b)[0] if b else None
                        else:
                            b = pe.read_va(tgt, 8); v = _s.unpack("<d", b)[0] if b else None
                        if v is not None and v == v:  # not NaN
                            av = abs(v)
                            if (lo is None or av >= lo) and (hi is None or av <= hi):
                                out.append((insn_va, v, kind))
                j = raw.find(pref, j + 1)
    return out


def analyze(pe, idx, va, max_bytes=0x6000, show=True):
    """Disassemble a method; resolve call targets to method sigs and read any
    RIP-relative float/double constants loaded (movss/movsd/comiss/addss/...).
    Returns (insns, calls, floats)."""
    from capstone import x86
    ins = disasm_method(pe, idx, va, max_bytes)
    calls, floats, lines = [], [], []
    for i in ins:
        tgt = rip_target(i)
        note = ""
        mn = i.mnemonic
        if mn in ("call", "jmp") and i.operands and i.operands[0].type == x86.X86_OP_IMM:
            d = i.operands[0].imm
            m = idx.method_at(d)
            nm = m["sig"] if (m and abs(m["va"] - d) < 0x6000) else "?"
            if mn == "call":
                calls.append((i.address, d, nm))
            note = " -> %x %s" % (d, nm[:78])
        elif tgt is not None and ("ss" in mn or "sd" in mn or mn == "movss" or mn == "movsd"):
            # float/double constant load from .rdata
            if "sd" in mn:
                v = read_f64(pe, tgt)
                floats.append((i.address, tgt, v, "f64"))
                note = "  ; =%s (f64 @%x)" % (v, tgt)
            else:
                v = read_f32(pe, tgt)
                floats.append((i.address, tgt, v, "f32"))
                note = "  ; =%s (f32 @%x)" % (v, tgt)
        elif tgt is not None:
            note = "  ; [data %x]" % tgt
        lines.append("  %x  %-9s %s%s" % (i.address, mn, i.op_str, note))
    if show:
        print("\n".join(lines))
    return ins, calls, floats


if __name__ == "__main__":
    pe = PE()
    idx = Index()
    print("methods indexed:", len(idx.methods))
    print("text section:", [hex(x) for x in pe.text_range()])
    # sanity: the known DamagePhysical float setter
    m = idx.method_at(0x1849EC5B0)
    print("method at 0x1849EC5B0:", m["sig"] if m else None, hex(m["va"]) if m else "")