#!/usr/bin/env python3
"""Locate metadata-usage tokens / pointers inside an il2cpp binary.

On metadata v27 the compiled code does not carry literal indices: a string
literal is reached through a metadata-usage value `(type << 29) | (index << 1) | 1`
(type 5 = string literal). That value shows up in the image either as an
immediate operand in code or as the pre-loader content of a usage slot, so a
plain byte scan is enough to find the referencing function - parse #29705's
token with tools/sep-literals.py, scan here, then read the enclosing method
from the runtime address map.

Usage:
    python tools/sep-scan.py --binary <GameAssembly.dll> --token 0xA000E813[,..]
    python tools/sep-scan.py --binary <GameAssembly.dll> --pointer 0x1891B46E8[,..]
    python tools/sep-scan.py --binary <GameAssembly.dll> --map <radicade-addr-map.txt> --token ...

Options:
    --section .text   restrict the scan (repeatable; default: every section)
    --context N       extra bytes of hex context to print (default 16)
"""

import argparse
import os
import struct
import sys

METADATA_USAGE_STRING_LITERAL = 5


def literal_token(index):
    return (METADATA_USAGE_STRING_LITERAL << 29) | (index << 1) | 1


def parse_number(text):
    return int(text, 0)


def read_sections(blob):
    pe_offset = struct.unpack_from("<I", blob, 0x3C)[0]
    if blob[pe_offset : pe_offset + 4] != b"PE\0\0":
        raise SystemExit("not a PE image")
    machine, section_count = struct.unpack_from("<HH", blob, pe_offset + 4 + 2)
    optional_size = struct.unpack_from("<H", blob, pe_offset + 4 + 16)[0]
    optional = pe_offset + 4 + 20
    magic = struct.unpack_from("<H", blob, optional)[0]
    if magic == 0x20B:
        image_base = struct.unpack_from("<Q", blob, optional + 24)[0]
    elif magic == 0x10B:
        image_base = struct.unpack_from("<I", blob, optional + 28)[0]
    else:
        raise SystemExit(f"unknown optional header magic {magic:#x}")
    sections = []
    header = optional + optional_size
    for i in range(section_count):
        entry = header + 40 * i
        name = blob[entry : entry + 8].split(b"\0")[0].decode("ascii", "replace")
        virtual_size, virtual_address, raw_size, raw_offset = struct.unpack_from("<IIII", blob, entry + 8)
        sections.append((name, virtual_address, virtual_size, raw_offset, raw_size))
    return image_base, machine, sections


def load_map(path, image_base):
    entries = []
    with open(path, "r", encoding="utf-8", errors="replace") as handle:
        for line in handle:
            if line.startswith("#"):
                continue
            parts = line.rstrip("\r\n").split("\t")
            if len(parts) < 3:
                continue
            try:
                rva = int(parts[0], 16)
            except ValueError:
                continue
            entries.append((rva, parts[1], parts[2]))
    entries.sort()
    return entries


def enclosing(entries, rva, image_base):
    lo, hi = 0, len(entries) - 1
    best = None
    while lo <= hi:
        mid = (lo + hi) // 2
        if entries[mid][0] <= rva:
            best = entries[mid]
            lo = mid + 1
        else:
            hi = mid - 1
    if best is None:
        return None
    return best, rva - best[0]


def main(argv):
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument("--binary", required=True)
    parser.add_argument("--map")
    parser.add_argument("--token", action="append", default=[])
    parser.add_argument("--pointer", action="append", default=[])
    parser.add_argument("--literal", action="append", type=int, default=[])
    parser.add_argument("--section", action="append", default=[])
    parser.add_argument("--context", type=int, default=16)
    args = parser.parse_args(argv[1:])

    patterns = []
    labels = []
    for value in args.literal:
        patterns.append(struct.pack("<I", literal_token(value)))
        labels.append(f"literal #{value} token 0x{literal_token(value):X} (4-byte)")
    for group in args.token:
        for value in group.split(","):
            number = parse_number(value)
            patterns.append(struct.pack("<I", number))
            patterns.append(struct.pack("<Q", number))
            labels.append(f"token {number:#x}")
            labels.append(f"token {number:#x} (8-byte)")
    for group in args.pointer:
        for value in group.split(","):
            number = parse_number(value)
            patterns.append(struct.pack("<Q", number))
            labels.append(f"pointer {number:#x}")
    if not patterns:
        raise SystemExit("nothing to search for: pass --literal/--token/--pointer")

    blob = open(args.binary, "rb").read()
    image_base, machine, sections = read_sections(blob)
    print(f"binary {os.path.basename(args.binary)}: machine {machine:#x}, image base {image_base:#x}, {len(sections)} sections")

    entries = load_map(args.map, image_base) if args.map else None
    total = 0
    for index, pattern in enumerate(patterns):
        wanted = [s for s in sections if not args.section or s[0] in args.section]
        for name, virtual_address, virtual_size, raw_offset, raw_size in wanted:
            size = min(raw_size, max(0, len(blob) - raw_offset))
            start = raw_offset
            end = raw_offset + size
            offset = start
            while True:
                offset = blob.find(pattern, offset, end)
                if offset < 0:
                    break
                va = image_base + virtual_address + (offset - raw_offset)
                line = f"  {labels[index]} @ RVA {virtual_address + (offset - raw_offset):#x} VA {va:#x} ({name})"
                if entries:
                    found = enclosing(entries, virtual_address + (offset - raw_offset), image_base)
                    if found:
                        (rva, module, method), delta = found
                        line += f"  -> {module} {method} +{delta:#x}"
                print(line)
                context = blob[max(offset - 4, start) : offset + len(pattern) + args.context]
                print(f"     ctx {context.hex(' ')}")
                total += 1
                offset += 1
    print(f"{total} occurrence(s)")

    if entries:
        print("map loaded: %d methods" % len(entries))
    return 0


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