#!/usr/bin/env python3
"""Evaluate the native Math.min hook's per-lap term F(i, b) of neon_outrun (static
reimplementation by emulating only the hook's pure arithmetic with Unicorn).

The hook (0x67eb83) for lap i with integer lap time b (trunc of lapSecs, 0..255):
    key  = flattened constant (0x67ed5d..0x67eec7)
    L,s  = FUN_0068003e(FUN_0067f9d8(i, key))          (rax, dl)
    t    = C1 / X * (L + (-b if s&1 else b))
    t    = t - round(t / C2) * C2
    F    = sum of N Taylor terms of sin(t)             (0x687cff / 0x687d80)
    acc += weight_i * F                                (weights queued by the sort hook)
replaceAll only reveals the lap-time bytes if acc == sum(weights), i.e. F(i, b_i) == 1
for every lap, so b_i is the unique byte with sin(t) == 1.
No OS, no syscalls: only the ELF's PT_LOAD bytes are mapped; the lock/push call is skipped.
"""
import struct
import sys

from elftools.elf.elffile import ELFFile
from elftools.elf.relocation import RelocationSection

from unicorn import Uc, UC_ARCH_X86, UC_MODE_64, UC_HOOK_CODE
from unicorn.x86_const import (UC_X86_REG_RSP, UC_X86_REG_RIP, UC_X86_REG_R14, UC_X86_REG_RAX,
                               UC_X86_REG_RDX, UC_X86_REG_XMM0)

STACK = 0x7f0000000000


class Model:
    def __init__(self, path):
        raw = open(path, "rb").read()
        self.mu = mu = Uc(UC_ARCH_X86, UC_MODE_64)
        # map every PT_LOAD at its VA (file bytes, zero-filled tail) and apply RELATIVE relocs
        elf = ELFFile(open(path, "rb"))
        for seg in elf.iter_segments():
            if seg["p_type"] != "PT_LOAD":
                continue
            va, off, fs, ms = seg["p_vaddr"], seg["p_offset"], seg["p_filesz"], seg["p_memsz"]
            lo, hi = va & ~0xfff, (va + ms + 0xfff) & ~0xfff
            try:
                mu.mem_map(lo, hi - lo)
            except Exception:  # adjacent segments sharing a page
                mu.mem_map(lo + 0x1000, hi - lo - 0x1000)
            mu.mem_write(va, raw[off:off + fs])
        for sec in elf.iter_sections():
            if isinstance(sec, RelocationSection):
                for r in sec.iter_relocations():
                    if r["r_info_type"] == 8:
                        mu.mem_write(r["r_offset"], struct.pack("<Q", r["r_addend"]))
        mu.mem_map(STACK, 0x100000)
        self.skip = {}
        mu.hook_add(UC_HOOK_CODE, self._hook)

    def _hook(self, mu, addr, size, _):
        if addr in self.skip:
            mu.reg_write(UC_X86_REG_RIP, self.skip[addr])

    def run(self, start, end, regs=None, stack=None):
        mu = self.mu
        rsp = STACK + 0x80000
        mu.mem_write(STACK, b"\0" * 0x100000)
        mu.reg_write(UC_X86_REG_RSP, rsp)
        for r, v in (regs or {}).items():
            mu.reg_write(r, v)
        for off, data in (stack or {}).items():
            mu.mem_write(rsp + off, data)
        mu.emu_start(start, end, count=50_000_000)
        return rsp

    def key(self):
        self.run(0x67ed5d, 0x67eec7)
        return self.mu.reg_read(UC_X86_REG_RAX) & 0xff

    def lap_const(self, i, key):
        self.run(0x67eec7, 0x67eeda, {UC_X86_REG_R14: i, UC_X86_REG_RAX: key})
        return self.mu.reg_read(UC_X86_REG_RAX), self.mu.reg_read(UC_X86_REG_RDX) & 0xff

    def F(self, L, s, b):
        # skip the lock+push+unlock of the byte buffer (0x67ef42 -> 0x67ef77)
        self.skip = {0x67ef42: 0x67ef77}
        stack = {0x8: struct.pack("<d", float(b)), 0x28: struct.pack("<Q", L), 0x7: bytes([s])}
        self.run(0x67eee3, 0x67f958, stack=stack)
        self.skip = {}
        xmm0 = self.mu.reg_read(UC_X86_REG_XMM0) & ((1 << 64) - 1)
        return struct.unpack("<d", struct.pack("<Q", xmm0))[0]


def solve(path, laps=28):
    m = Model(path)
    key = m.key()
    out = []
    for i in range(laps):
        L, s = m.lap_const(i, key)
        best = max(range(256), key=lambda b: m.F(L, s, b))
        out.append((i, L, s, best, m.F(L, s, best)))
    return key, out


if __name__ == "__main__":
    key, res = solve(sys.argv[1])
    print(f"key byte {key:#x}")
    for i, L, s, b, f in res:
        print(f"lap {i + 1:2}: L={L:#x} s={s} best b={b:3} ({chr(b)!r}) F={f:.12f}")
    print("bytes:", bytes(r[3] for r in res))