"""Private original-code oracle for current EntityId hash membership.

No process handle is opened and no live game function is called.
"""
import importlib.util
import json
import random
import struct
import sys
from pathlib import Path

sys.path[:0] = ['src/loader/runtime/profile-tools/dev/tools/enshrouded']
from discovery.pe import PE
from unicorn.x86_const import (UC_X86_REG_RAX, UC_X86_REG_RCX, UC_X86_REG_RDX,
                               UC_X86_REG_R8, UC_X86_REG_R9, UC_X86_REG_RSP, UC_X86_REG_RIP)
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

spec = importlib.util.spec_from_file_location(
    'attribute_oracle', 'src/loader/runtime/profile-tools/dev/tools/enshrouded/verify-attribute-vm.py')
vm = importlib.util.module_from_spec(spec)
spec.loader.exec_module(vm)
pe = PE(r'C:\Program Files (x86)\Steam\steamapps\common\Enshrouded\enshrouded.exe')
assert pe.sha256 == '4d6d43ea4a91e71c82d60b045dde932179c99e32e40657efe374688b2933589a'
o = vm.Oracle(pe)
u = o.uc
manager, bits, keys, pointers = 0x70000000, 0x70001000, 0x70002000, 0x70003000
argument, entity, layout, storage, component, output = (0x70004000, 0x70005000,
    0x70006000, 0x70008000, 0x70009000, 0x7000a000)
table = manager + 0x148
u.mem_map(manager, 0x10000)


def write(address, fmt, *values):
    u.mem_write(address, struct.pack(fmt, *values))


def read(address, fmt):
    return struct.unpack(fmt, u.mem_read(address, struct.calcsize(fmt)))[0]


def hash_id(value):
    value = (((value >> 16) ^ value) * 0x45d9f3b) & 0xffffffff
    value = (((value >> 16) ^ value) * 0x45d9f3b) & 0xffffffff
    return (value >> 16) ^ value


def setup(values, capacity=16):
    u.mem_write(manager, bytes(0x10000))
    write(table + 8, '<Q', bits)
    write(table + 0x10, '<Q', capacity)
    write(table + 0x28, '<Q', keys)
    write(table + 0x40, '<Q', pointers)
    write(table + 0x54, '<I', len(values))
    occupied = set()
    locations = {}
    for value in values:
        slot = hash_id(value) & (capacity - 1)
        while slot in occupied:
            slot = (slot + 1) & (capacity - 1)
        occupied.add(slot)
        locations[value] = slot
        write(keys + slot * 4, '<I', value)
        write(pointers + slot * 8, '<Q', entity)
    for block in range((capacity + 63) // 64):
        write(bits + block * 8, '<Q', sum(1 << (slot % 64) for slot in occupied if slot // 64 == block))
    return locations


def call(rva, rcx, rdx, r8=0, r9=0):
    rsp = o.STACK + 0xff08
    write(rsp, '<Q', o.STOP)
    for reg, value in ((UC_X86_REG_RCX, rcx), (UC_X86_REG_RDX, rdx),
                       (UC_X86_REG_R8, r8), (UC_X86_REG_R9, r9), (UC_X86_REG_RSP, rsp)):
        u.reg_write(reg, value)
    u.emu_start(pe.base + rva, o.STOP, timeout=100000, count=100000)
    assert u.reg_read(UC_X86_REG_RIP) == o.STOP
    return u.reg_read(UC_X86_REG_RAX)


cases = []
for value in [0, 1, 0x7fffffff, 0x80000000, 0xdeadbeef, 0xffffffff]:
    locations = setup([value])
    write(argument, '<I', value)
    actual = call(0x9a7ab0, table, argument)
    assert actual == locations[value]
    cases.append(dict(name='u32-boundary', id=value, expected=locations[value], actual=actual))

colliding = [v for v in range(1, 10000) if hash_id(v) & 15 == 15][:5]
locations = setup(colliding)
assert list(locations.values()) == [15, 0, 1, 2, 3]
for value in colliding:
    write(argument, '<I', value)
    actual = call(0x9a7ab0, table, argument)
    assert actual == locations[value]
    cases.append(dict(name='linear-probe-wrap', id=value, expected=locations[value], actual=actual))

rng = random.Random(25966634)
for _ in range(12):
    values = rng.sample(range(1, 0xffffffff), 37)
    locations = setup(values, capacity=64)
    for value in values + rng.sample(range(1, 0xffffffff), 8):
        write(argument, '<I', value)
        actual = call(0x9a7ab0, table, argument)
        expected = locations.get(value, 64)
        assert actual == expected
        cases.append(dict(name='random-membership-or-miss', id=value, expected=expected, actual=actual))

value = 0xfedcba98
locations = setup([value])
slot = locations[value]
write(entity + 0x10, '<I', value)
write(entity + 0x18, '<Q', layout)
write(entity + 0x20, '<Q', storage)
write(entity + 0x30, '<I', 2)
write(layout, '<Q', 1 << 5)
write(layout + 0x84 + 5 * 2, '<H', 0x100)
write(layout + 0xa84 + 5 * 2, '<H', 4)
write(component, '<H', 5)
call(0x9a5900, output, manager, component, value)
assert read(output, '<Q') == storage + 0x108 and read(output + 8, '<Q') == 4
cases.append(dict(name='component-helper-resolves-occupied-ID', result=[read(output, '<Q'), read(output + 8, '<Q')]))
write(bits + (slot // 64) * 8, '<Q', 0)
call(0x9a5900, output, manager, component, value)
assert read(output, '<Q') == 0 and read(output + 8, '<Q') == 0
cases.append(dict(name='component-helper-rejects-nonnull-erased-pointer', result=[read(output, '<Q'), read(output + 8, '<Q')]))

c = Cs(CS_ARCH_X86, CS_MODE_64)
functions = {}
for name, rva in [('lookupEntityId', 0x9a7ab0), ('lookupComponentByEntityId', 0x9a5900)]:
    fragments = [f for f in pe.functions if f['primaryBeginRva'] == rva]
    functions[name] = dict(rva=rva, fragments=[dict(beginRva=f['beginRva'], endRva=f['endRva'],
        instructions=[dict(rva=i.address, mnemonic=i.mnemonic, operands=i.op_str)
            for i in c.disasm(pe.read(f['beginRva'], f['endRva']-f['beginRva']), f['beginRva'])]) for f in fragments])

out = Path('src/loader/runtime/profile-tools/devdata/creative-mode-146/current-entity-membership-oracle-20261016.json')
out.write_text(json.dumps(dict(executable=pe.identity(), processOpened=False, processWrites=False,
    nativeGameCalls=False, cases=cases, functions=functions,
    findings=dict(hashWidth=32, multiplier=0x45d9f3b, occupancyOffset=0x150,
        capacityOffset=0x158, keysOffset=0x170, pointerTableOffset=0x188,
        logicalCountOffset=0x19c, linearProbe=True, wrapAtCapacity=True,
        emptyBucketStopsLookup=True, nonnullErasedPointersAreNotMembership=True),
    limitations=['Private original-code execution, not concurrent mutation coverage']), indent=2))
print(json.dumps(dict(cases=len(cases), passed=True, path=str(out))))
