"""Locate named Client method token checks in a supplied PE (read only).

The correlation was observed for GridFightModule.GetSeason and _SyncSeason in
this Client: the native check loads the low 24 bits of MetadataToken into ECX
before a direct call. Report candidate sites; do not infer function boundaries
or claim that every candidate is a method implementation.
"""
import argparse
import hashlib
import json
import pathlib
import re
import struct


def executable_sections(data):
    if data[:2] != b"MZ":
        raise ValueError("Expected a PE image")
    pe = struct.unpack_from("<I", data, 0x3c)[0]
    if data[pe:pe + 4] != b"PE\0\0":
        raise ValueError("Invalid PE signature")
    count = struct.unpack_from("<H", data, pe + 6)[0]
    optional_size = struct.unpack_from("<H", data, pe + 20)[0]
    sections = []
    for index in range(count):
        offset = pe + 24 + optional_size + index * 40
        virtual_size, rva, raw_size, raw = struct.unpack_from("<IIII", data, offset + 8)
        characteristics = struct.unpack_from("<I", data, offset + 36)[0]
        if characteristics & 0x20000000:  # IMAGE_SCN_MEM_EXECUTE in the PE schema.
            if raw + raw_size > len(data):
                raise ValueError("Truncated executable section")
            sections.append((rva, data[raw:raw + raw_size]))
    return sections


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--assembly", type=pathlib.Path, required=True)
    parser.add_argument("--tokens", type=pathlib.Path, required=True)
    parser.add_argument("--output", type=pathlib.Path, required=True)
    args = parser.parse_args()
    data = args.assembly.read_bytes()
    sections = executable_sections(data)
    methods, by_check = {}, {}
    for line in args.tokens.read_text(encoding="utf-8-sig").splitlines():
        match = re.search(r"\[GAC_CW_METHOD_TOKEN\] METHOD (\S+) (\S+) (\d+)$", line)
        if not match:
            continue
        token = int(match[3])
        name = match[1] + "." + match[2] + "#" + hex(token)
        if token & 0xff000000 != 0x06000000:  # ECMA-335 MethodDef token table.
            raise SystemExit("Not a MethodDef token: " + name)
        methods[name] = {"metadata_token": hex(token), "check_id": hex(token & 0x00ffffff), "candidate_sites": []}
        by_check.setdefault(token & 0x00ffffff, []).append(name)
    if not methods:
        raise SystemExit("No captured method tokens found")
    # One section pass, rather than rereading a large executable for every
    # captured method. The byte pattern is deliberately only a candidate.
    pattern = re.compile(rb"(?=\xb9(.{4})\xe8(.{4}))", re.S)
    for section_rva, section in sections:
        for match in pattern.finditer(section):
            check_id = struct.unpack("<I", match[1])[0]
            if check_id not in by_check:
                continue
            call_rva = section_rva + match.start() + 5
            target = call_rva + 5 + struct.unpack("<i", match[2])[0]
            for name in by_check[check_id]:
                methods[name]["candidate_sites"].append({"check_rva": hex(section_rva + match.start()), "call_target_rva": hex(target)})
    output = {
        "assembly_sha256": hashlib.sha256(data).hexdigest(),
        "token_capture_sha256": hashlib.sha256(args.tokens.read_bytes()).hexdigest(),
        "limitation": "Candidate token-check sites only; no inferred function start or gameplay semantics",
        "methods": methods,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2) + "\n", encoding="utf-8")
    print(f"{len(methods)} named methods; {sum(bool(v['candidate_sites']) for v in methods.values())} with candidate native sites")


if __name__ == "__main__":
    main()