#!/usr/bin/env python3
"""Interpret FIFA 17's atom -> field-id dispatch functions instead of pattern-scanning them.

WHY THIS EXISTS
---------------
CardsDLL turns a JSON key into an "atom index" (a position in the string-pointer
table at .data 0x1802d2760), then a per-response-family mapper converts that index
into an internal field id with a chain of integer compares and jump tables.

A previous attempt to recover each mapper's accepted atoms by scanning for
`sub ecx,K` / `cmp ecx,L` / `ja` patterns produced a confidently wrong answer: it
reported that no mapper accepts atom 424 (`manager`), while a live client plainly
holds a resident manager record. Pattern scanning cannot see control flow, so it
cannot tell which compares are actually reachable.

This module executes the mappers instead. The modelled subset is exactly what these
functions use: the resolver call, integer cmp/sub/add/dec, conditional and computed
jumps, jump-table loads out of the image, lea, movsxd, and `mov eax,imm; ret`.
Anything outside that subset raises Unsupported, so a wrong field id is never
returned silently.

TWO DECODER TRAPS THIS MODULE IS REQUIRED TO HANDLE
---------------------------------------------------
1. ModRM rm==5 with mod!=0 is [rbp+disp], NOT RIP-relative. Only mod==0 with rm==5
   is RIP-relative. Treating all rm==5 as RIP-relative hides rbp-based DTO accesses.
   Covered by test_rbp_relative_is_not_rip_relative.
2. A constant frequently arrives in a register (`mov r8d,0x4` ... later stored), so
   searching for an immediate-to-memory store misses it. The interpreter tracks
   register values, so propagated constants are followed.
   Covered by test_constant_propagated_through_register.

Run `--selftest` to execute the positive controls. Negative results from this tool
are only admissible when the selftest passes.
"""
from __future__ import annotations

import argparse
import bisect
import struct
import sys
from pathlib import Path

REGS = ("rax", "rcx", "rdx", "rbx", "rsp", "rbp", "rsi", "rdi",
        "r8", "r9", "r10", "r11", "r12", "r13", "r14", "r15")

ATOM_TABLE_BASE = 0x1802D2760      # validated against 6 known anchors, see anchors()
ATOM_RESOLVER = 0x180180D00        # key string -> atom index, returns in eax
ITEM_MAPPER = 0x18012FD40          # the DTO/item mapper: atom 568 'players' -> 1


class Unsupported(Exception):
    """The mapper used an instruction or address outside the modelled subset."""


def s32(v: int) -> int:
    v &= 0xFFFFFFFF
    return v - 0x100000000 if v & 0x80000000 else v


class Image:
    """A parsed PE, with VA<->file mapping and .pdata function bounds."""

    def __init__(self, path: Path):
        self.buf = path.read_bytes()
        b = self.buf
        pe = struct.unpack_from("<I", b, 0x3C)[0]
        if b[pe:pe + 4] != b"PE\0\0":
            raise ValueError(f"{path} is not a PE image")
        nsec = struct.unpack_from("<H", b, pe + 6)[0]
        optsz = struct.unpack_from("<H", b, pe + 20)[0]
        self.base = struct.unpack_from("<Q", b, pe + 24 + 24)[0]
        self.sections = []
        for i in range(nsec):
            o = pe + 24 + optsz + 40 * i
            name = b[o:o + 8].rstrip(b"\0").decode(errors="replace")
            vsz, va, rsz, raw = struct.unpack_from("<IIII", b, o + 8)
            self.sections.append((name, va, vsz, raw, rsz))
        self._funcs = None

    def va2off(self, va: int):
        rva = va - self.base
        for _name, sva, vsz, raw, rsz in self.sections:
            if sva <= rva < sva + max(vsz, rsz):
                off = raw + (rva - sva)
                if off < len(self.buf):
                    return off
        return None

    def rd8(self, va: int) -> int:
        o = self.va2off(va)
        if o is None:
            raise Unsupported(f"unmapped byte read 0x{va:x}")
        return self.buf[o]

    def rd32(self, va: int) -> int:
        o = self.va2off(va)
        if o is None:
            raise Unsupported(f"unmapped dword read 0x{va:x}")
        return struct.unpack_from("<I", self.buf, o)[0]

    def cstr(self, va: int, maxlen: int = 96):
        o = self.va2off(va)
        if o is None:
            return None
        end = self.buf.find(b"\0", o, o + maxlen)
        if end < 0:
            return None
        try:
            return self.buf[o:end].decode("ascii")
        except UnicodeDecodeError:
            return None

    # ---- .pdata gives exact function bounds; never guess a prologue ----
    def functions(self):
        if self._funcs is None:
            sec = next(s for s in self.sections if s[0] == ".pdata")
            _n, _va, vsz, raw, _rsz = sec
            out = []
            for i in range(vsz // 12):
                beg, end, _unw = struct.unpack_from("<III", self.buf, raw + 12 * i)
                if beg or end:
                    out.append((self.base + beg, self.base + end))
            out.sort()
            self._funcs = out
        return self._funcs

    def function_of(self, va: int):
        fs = self.functions()
        starts = [f[0] for f in fs]
        i = bisect.bisect_right(starts, va) - 1
        if i >= 0 and fs[i][0] <= va < fs[i][1]:
            return fs[i]
        return None

    def atom(self, index: int):
        ptr = struct.unpack_from("<Q", self.buf, self.va2off(ATOM_TABLE_BASE) + 8 * index)[0]
        return self.cstr(ptr)

    def atom_index(self, name: str):
        off = self.va2off(ATOM_TABLE_BASE)
        for i in range(4096):
            ptr = struct.unpack_from("<Q", self.buf, off + 8 * i)[0]
            if self.cstr(ptr) == name:
                return i
        return None


class Mapper:
    """Executes one dispatch function for a given atom index."""

    def __init__(self, image: Image, resolver: int = ATOM_RESOLVER):
        self.img = image
        self.resolver = resolver

    def _ea(self, k: int, rex: int, r: dict):
        """Decode ModRM[+SIB][+disp].

        Returns (nbytes, dst_reg, addr, src_reg). addr is an int, or the marker
        ("rip", disp) which the caller resolves once it knows the instruction
        length, or None for a register-form operand.

        TRAP 1: rm==5 is RIP-relative ONLY when mod==0. With mod 1 or 2 it is
        [rbp+disp] and must be resolved from rbp.
        """
        b = self.img.buf
        modrm = b[k]
        mod, rm = modrm >> 6, modrm & 7
        dst = REGS[(((modrm >> 3) & 7) | ((rex & 4) << 1)) & 15]
        n = 1
        if mod == 3:
            return n, dst, None, REGS[(rm | ((rex & 1) << 3)) & 15]
        base_v = idx_v = disp = 0
        if rm == 4:
            sib = b[k + 1]
            n += 1
            scale = 1 << (sib >> 6)
            ir = ((sib >> 3) & 7) | ((rex & 2) << 2)
            br = (sib & 7) | ((rex & 1) << 3)
            if (ir & 15) != 4:
                idx_v = r[REGS[ir & 15]] * scale
            if (sib & 7) == 5 and mod == 0:
                disp = struct.unpack_from("<i", b, k + n)[0]
                n += 4
            else:
                base_v = r[REGS[br & 15]]
        elif rm == 5 and mod == 0:
            disp = struct.unpack_from("<i", b, k + 1)[0]
            return n + 4, dst, ("rip", disp), None
        else:
            base_v = r[REGS[(rm | ((rex & 1) << 3)) & 15]]
        if mod == 1:
            disp = struct.unpack_from("<b", b, k + n)[0]
            n += 1
        elif mod == 2:
            disp = struct.unpack_from("<i", b, k + n)[0]
            n += 4
        return n, dst, (base_v + idx_v + disp) & 0xFFFFFFFFFFFFFFFF, None

    @staticmethod
    def _cond(cc: int, last) -> bool:
        a, b = last
        sa, sb = s32(a), s32(b)
        ua, ub = a & 0xFFFFFFFF, b & 0xFFFFFFFF
        if cc == 0x4: return sa == sb
        if cc == 0x5: return sa != sb
        if cc == 0xF: return sa > sb
        if cc == 0xD: return sa >= sb
        if cc == 0xC: return sa < sb
        if cc == 0xE: return sa <= sb
        if cc == 0x7: return ua > ub
        if cc == 0x3: return ua >= ub
        if cc == 0x2: return ua < ub
        if cc == 0x6: return ua <= ub
        if cc == 0x8: return sa < sb
        if cc == 0x9: return sa >= sb
        raise Unsupported(f"condition code 0x{cc:x}")

    def run(self, start: int, atom: int, limit: int = 5000) -> int:
        b = self.img.buf
        r = {k: 0 for k in REGS}
        last = (0, 0)
        va = start
        for _ in range(limit):
            i0 = self.img.va2off(va)
            if i0 is None:
                raise Unsupported(f"pc unmapped 0x{va:x}")
            j = i0
            while b[j] in (0x66, 0x67, 0xF2, 0xF3):
                j += 1
            rex = 0
            if 0x40 <= b[j] <= 0x4F:
                rex = b[j]
                j += 1
            op = b[j]
            pre = j - i0

            if op == 0xC3:
                return r["rax"] & 0xFFFFFFFF
            if op == 0xCC:
                raise Unsupported(f"int3 at 0x{va:x}: ran off the end of the function")
            if op == 0xE8:
                tgt = va + pre + 5 + struct.unpack_from("<i", b, j + 1)[0]
                if tgt != self.resolver:
                    raise Unsupported(f"call to non-resolver 0x{tgt:x} at 0x{va:x}")
                r["rax"] = atom & 0xFFFFFFFF          # resolver returns the atom index
                va += pre + 5
                continue
            if op == 0xE9:
                va += pre + 5 + struct.unpack_from("<i", b, j + 1)[0]
                continue
            if op == 0xEB:
                va += pre + 2 + struct.unpack_from("<b", b, j + 1)[0]
                continue
            if 0x70 <= op <= 0x7F:
                nxt = va + pre + 2
                rel = struct.unpack_from("<b", b, j + 1)[0]
                va = nxt + rel if self._cond(op & 0xF, last) else nxt
                continue
            if op == 0x0F and 0x80 <= b[j + 1] <= 0x8F:
                nxt = va + pre + 6
                rel = struct.unpack_from("<i", b, j + 2)[0]
                va = nxt + rel if self._cond(b[j + 1] & 0xF, last) else nxt
                continue
            if 0xB8 <= op <= 0xBF:
                r[REGS[((op - 0xB8) | ((rex & 1) << 3)) & 15]] = struct.unpack_from("<I", b, j + 1)[0]
                va += pre + 5
                continue
            if op in (0x05, 0x2D, 0x3D):
                # accumulator short forms: add/sub/cmp eax, imm32
                imm = struct.unpack_from("<i", b, j + 1)[0]
                cur = r["rax"] & 0xFFFFFFFF
                if op == 0x3D:
                    last = (cur, imm & 0xFFFFFFFF)
                elif op == 0x2D:
                    r["rax"] = (cur - imm) & 0xFFFFFFFF
                    last = (r["rax"], 0)
                else:
                    r["rax"] = (cur + imm) & 0xFFFFFFFF
                    last = (r["rax"], 0)
                va += pre + 5
                continue
            if op in (0x81, 0x83):
                w = 4 if op == 0x81 else 1
                modrm = b[j + 1]
                if modrm >> 6 != 3:
                    raise Unsupported(f"{op:02x} memory form at 0x{va:x}")
                reg = REGS[((modrm & 7) | ((rex & 1) << 3)) & 15]
                imm = struct.unpack_from("<i" if w == 4 else "<b", b, j + 2)[0]
                ext = (modrm >> 3) & 7
                cur = r[reg] & 0xFFFFFFFF
                if ext == 7:
                    last = (cur, imm & 0xFFFFFFFF)
                elif ext == 5:
                    r[reg] = (cur - imm) & 0xFFFFFFFF
                    last = (r[reg], 0)
                elif ext == 0:
                    r[reg] = (cur + imm) & 0xFFFFFFFF
                    last = (r[reg], 0)
                else:
                    raise Unsupported(f"{op:02x} /{ext} at 0x{va:x}")
                va += pre + 2 + w
                continue
            if op == 0xFF and b[j + 1] >> 6 == 3:
                ext = (b[j + 1] >> 3) & 7
                reg = REGS[((b[j + 1] & 7) | ((rex & 1) << 3)) & 15]
                if ext == 1:
                    r[reg] = (r[reg] - 1) & 0xFFFFFFFF
                    last = (r[reg], 0)
                    va += pre + 2
                    continue
                if ext == 4:
                    va = r[reg]
                    continue
                raise Unsupported(f"ff /{ext} at 0x{va:x}")
            if op == 0x0F and b[j + 1] == 0xB6:
                n, dst, addr, src = self._ea(j + 2, rex, r)
                end = va + pre + 2 + n
                if isinstance(addr, tuple):
                    addr = end + addr[1]
                r[dst] = self.img.rd8(addr) if addr is not None else r[src] & 0xFF
                va = end
                continue
            if op in (0x8B, 0x8D):
                n, dst, addr, src = self._ea(j + 1, rex, r)
                end = va + pre + 1 + n
                if isinstance(addr, tuple):
                    addr = end + addr[1]
                if op == 0x8D:
                    if addr is None:
                        raise Unsupported(f"lea with register operand at 0x{va:x}")
                    r[dst] = addr
                else:
                    if addr is None:
                        # register form: mov r32, r32 (e.g. 8b c8 = mov ecx,eax)
                        r[dst] = r[src] if rex & 8 else r[src] & 0xFFFFFFFF
                    else:
                        r[dst] = self.img.rd32(addr)
                va = end
                continue
            if op == 0x89:
                modrm = b[j + 1]
                if modrm >> 6 != 3:
                    raise Unsupported(f"89 memory store at 0x{va:x}")
                src = REGS[((((modrm >> 3) & 7) | ((rex & 4) << 1))) & 15]
                dst = REGS[((modrm & 7) | ((rex & 1) << 3)) & 15]
                r[dst] = r[src] if rex & 8 else r[src] & 0xFFFFFFFF
                va += pre + 2
                continue
            if op == 0x63:
                modrm = b[j + 1]
                if modrm >> 6 != 3:
                    raise Unsupported(f"63 memory form at 0x{va:x}")
                src = REGS[((modrm & 7) | ((rex & 1) << 3)) & 15]
                dst = REGS[((((modrm >> 3) & 7) | ((rex & 4) << 1))) & 15]
                r[dst] = s32(r[src]) & 0xFFFFFFFFFFFFFFFF
                va += pre + 2
                continue
            if op in (0x01, 0x03, 0x29, 0x2B, 0x39, 0x3B,
                      0x09, 0x0B, 0x21, 0x23, 0x31, 0x33, 0x85):
                modrm = b[j + 1]
                if modrm >> 6 != 3:
                    raise Unsupported(f"{op:02x} memory form at 0x{va:x}")
                a = REGS[((modrm & 7) | ((rex & 1) << 3)) & 15]
                c = REGS[((((modrm >> 3) & 7) | ((rex & 4) << 1))) & 15]
                m = 0xFFFFFFFFFFFFFFFF if rex & 8 else 0xFFFFFFFF
                if op == 0x01:
                    r[a] = (r[a] + r[c]) & m
                elif op == 0x03:
                    r[c] = (r[c] + r[a]) & m
                elif op == 0x29:
                    r[a] = (r[a] - r[c]) & m
                    last = (r[a] & 0xFFFFFFFF, 0)
                elif op == 0x2B:
                    r[c] = (r[c] - r[a]) & m
                    last = (r[c] & 0xFFFFFFFF, 0)
                elif op in (0x09, 0x0B, 0x21, 0x23, 0x31, 0x33):
                    fn = {0x09: lambda x, y: x | y, 0x0B: lambda x, y: x | y,
                          0x21: lambda x, y: x & y, 0x23: lambda x, y: x & y,
                          0x31: lambda x, y: x ^ y, 0x33: lambda x, y: x ^ y}[op]
                    if op in (0x09, 0x21, 0x31):
                        r[a] = fn(r[a], r[c]) & m
                        last = (r[a] & 0xFFFFFFFF, 0)
                    else:
                        r[c] = fn(r[c], r[a]) & m
                        last = (r[c] & 0xFFFFFFFF, 0)
                elif op == 0x85:
                    last = ((r[a] & r[c]) & 0xFFFFFFFF, 0)
                elif op == 0x39:
                    last = (r[a] & 0xFFFFFFFF, r[c] & 0xFFFFFFFF)
                else:
                    last = (r[c] & 0xFFFFFFFF, r[a] & 0xFFFFFFFF)
                va += pre + 2
                continue
            if op == 0x90:
                va += pre + 1
                continue
            if op == 0x0F and b[j + 1] == 0x1F:
                n, _d, _a, _s = self._ea(j + 2, rex, r)
                va += pre + 2 + n
                continue
            raise Unsupported(f"opcode {op:02x} at 0x{va:x}")
        raise Unsupported("instruction limit reached")


def find_mappers(img: Image, resolver: int = ATOM_RESOLVER):
    """Every function containing a direct call to the atom resolver."""
    sec = next(s for s in img.sections if s[0] == ".text")
    _n, tva, _vsz, traw, trsz = sec
    out = {}
    for i in range(traw, traw + trsz - 5):
        if img.buf[i] != 0xE8:
            continue
        va = img.base + tva + (i - traw)
        if va + 5 + struct.unpack_from("<i", img.buf, i + 1)[0] == resolver:
            f = img.function_of(va)
            if f:
                out.setdefault(f[0], []).append(va)
    return out


# --------------------------------------------------------------------------
# selftest: the two decoder traps plus the live-verified positive controls
# --------------------------------------------------------------------------
def test_atom_anchors(img: Image) -> list:
    """The atom table base must reproduce known anchors, or every index is wrong."""
    anchors = {11: "actives", 363: "itemData", 376: "kicktakers",
               424: "manager", 568: "players", 718: "squadActives"}
    fails = []
    for idx, want in anchors.items():
        got = img.atom(idx)
        if got != want:
            fails.append(f"atom[{idx}] = {got!r}, expected {want!r}")
    return fails


def test_rbp_relative_is_not_rip_relative(img: Image) -> list:
    """TRAP 1. mod!=0 with rm==5 must resolve as [rbp+disp], not RIP-relative.

    Encoding under test: 8b 4d 20 == mov ecx,[rbp+0x20] (mod=01, rm=101).
    A decoder that treats rm==5 as RIP-relative computes a wildly different
    address and silently reads the wrong memory.
    """
    m = Mapper(img)
    r = {k: 0 for k in REGS}
    r["rbp"] = 0x140000000
    saved = img.buf
    try:
        img.buf = bytes.fromhex("8b4d20")
        n, dst, addr, _src = m._ea(1, 0, r)
    finally:
        img.buf = saved
    fails = []
    if isinstance(addr, tuple):
        fails.append("mod=01 rm=101 decoded as RIP-relative; must be [rbp+disp]")
    elif addr != 0x140000020:
        fails.append(f"[rbp+0x20] resolved to 0x{addr:x}, expected 0x140000020")
    if dst != "rcx":
        fails.append(f"destination decoded as {dst}, expected rcx")
    if n != 2:
        fails.append(f"modrm+disp8 consumed {n} bytes, expected 2")
    return fails


def test_constant_propagated_through_register(img: Image) -> list:
    """TRAP 2. A constant reaching a use through a register must be followed.

    Program: mov eax,0; mov r8d,4; mov eax,r8d; ret  -> must yield 4, which is
    only observable if register values propagate. Scanning for an immediate
    store would see nothing.
    """
    m = Mapper(img)
    saved = img.buf
    prog = bytes.fromhex("b800000000" "41b804000000" "4489c0" "c3")
    try:
        img.buf = prog
        img_va2off = img.va2off
        img.va2off = lambda va: va if 0 <= va < len(prog) else None
        got = m.run(0, 0)
    finally:
        img.buf = saved
        img.va2off = img_va2off
    return [] if got == 4 else [f"register-propagated constant yielded {got}, expected 4"]


def test_item_mapper_controls(img: Image) -> list:
    """Live/disassembly-verified behaviour of the item mapper."""
    m = Mapper(img)
    fails = []
    got = m.run(ITEM_MAPPER, 568)
    if got != 1:
        fails.append(f"item mapper atom 568 'players' -> {got}, expected 1")
    got = m.run(ITEM_MAPPER, 11)
    if got != 0:
        fails.append(f"item mapper atom 11 'actives' -> {got}, expected 0")
    return fails


def test_manager_424_is_accepted_somewhere(img: Image) -> list:
    """MANDATORY control. A live client holds a resident manager record, so some
    mapper must map atom 424 to a non-zero field id. The previous pattern-scan
    method failed exactly here, and any replacement must not."""
    m = Mapper(img)
    accepting = []
    for start in find_mappers(img):
        try:
            if m.run(start, 424):
                accepting.append(start)
        except Unsupported:
            continue
    if not accepting:
        return ["no mapper maps atom 424 'manager' to a non-zero field id, "
                "which contradicts the live resident manager record"]
    return []


def selftest(img: Image) -> int:
    checks = [
        ("atom table anchors", test_atom_anchors),
        ("trap 1: rbp-relative modrm", test_rbp_relative_is_not_rip_relative),
        ("trap 2: constant via register", test_constant_propagated_through_register),
        ("item mapper positive controls", test_item_mapper_controls),
        ("mandatory: manager atom 424 accepted", test_manager_424_is_accepted_somewhere),
    ]
    bad = 0
    for name, fn in checks:
        try:
            fails = fn(img)
        except Exception as exc:                       # noqa: BLE001 - report, don't mask
            fails = [f"raised {type(exc).__name__}: {exc}"]
        if fails:
            bad += 1
            print(f"  FAIL  {name}")
            for f in fails:
                print(f"          {f}")
        else:
            print(f"  ok    {name}")
    print("\n  ALL PASS" if not bad else f"\n  {bad} CHECK(S) FAILED - negative results are NOT admissible")
    return 1 if bad else 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("image", type=Path, help="CardsDLL_Win64_retail.dll")
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--atom", type=int, action="append", default=[],
                    help="atom index to resolve through every mapper")
    ap.add_argument("--name", action="append", default=[],
                    help="atom name to resolve through every mapper")
    args = ap.parse_args()
    img = Image(args.image)

    if args.selftest:
        return selftest(img)

    atoms = list(args.atom)
    for nm in args.name:
        idx = img.atom_index(nm)
        if idx is None:
            print(f"  atom {nm!r} not found in the table")
            return 2
        atoms.append(idx)
    if not atoms:
        ap.error("give --atom/--name, or --selftest")

    m = Mapper(img)
    mappers = find_mappers(img)
    print(f"  {len(mappers)} mapper function(s) found\n")
    for a in atoms:
        print(f"  === atom {a} ({img.atom(a)!r}) ===")
        rows, unsup = [], 0
        for start in sorted(mappers):
            try:
                fid = m.run(start, a)
            except Unsupported:
                unsup += 1
                continue
            if fid:
                rows.append((start, fid))
        for start, fid in rows:
            print(f"    mapper 0x{start:x} -> field id {fid} (0x{fid:x})")
        print(f"    {len(rows)} mapper(s) accept it; {unsup} not modelled\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())