#!/usr/bin/env python3
"""90CN.dll self-integrity CRC32 patcher (reimplementation of patch_dll_crc.py).

The client patch embeds this struct (see client-patch/90CN.cpp, VerifySelfIntegrity):

    struct SelfCrcSlot {
        unsigned int magic;     // 0xC0DEC0DE, used to locate the slot in the file
        unsigned int offset;    // absolute file offset of `expected`
        unsigned int expected;  // CRC32 of the whole file, skipping [offset, offset+4)
    };

At runtime the DLL re-reads itself, recomputes that CRC and compares. Any byte
change (e.g. patching out the check) makes the client refuse to start.

The CRC is the standard reflected CRC-32 (poly 0xEDB88320, init 0xFFFFFFFF,
final XOR 0xFFFFFFFF) == zlib.crc32 / binascii.crc32.

Usage:
    python patch_dll_crc.py verify <dll>     # check the current stored CRC
    python patch_dll_crc.py patch  <dll>     # rewrite offset+expected in place
    python patch_dll_crc.py patch  <dll> -o <out.dll>   # write to a copy instead

Exit codes: 0 = ok, 1 = verification failure / slot not found, 2 = usage error.
"""

from __future__ import annotations

import argparse
import binascii
import struct
import sys
from pathlib import Path

MAGIC = 0xC0DEC0DE
MAGIC_BYTES = struct.pack("<I", MAGIC)
# struct layout is three consecutive little-endian uint32
SLOT_SIZE = 12
OFFSET_FIELD = 4
EXPECTED_FIELD = 8


class SlotError(RuntimeError):
    pass


def find_slot(blob: bytes) -> int:
    """Return the file offset of the SelfCrcSlot struct.

    The magic must be followed by a plausible `offset` field pointing at the
    `expected` field it lives in; that rejects accidental magic collisions in
    code or data.
    """
    start = 0
    candidates: list[int] = []
    while True:
        pos = blob.find(MAGIC_BYTES, start)
        if pos < 0:
            break
        candidates.append(pos)
        start = pos + 1
    if not candidates:
        raise SlotError("0xC0DEC0DE slot magic not found")

    for pos in candidates:
        if pos + SLOT_SIZE > len(blob):
            continue
        offset = struct.unpack_from("<I", blob, pos + OFFSET_FIELD)[0]
        if offset == pos + EXPECTED_FIELD and offset + 4 <= len(blob):
            return pos
    # Fall back to the first candidate so `patch` can still repair an unpatched
    # or corrupted slot (fresh build has offset == 0).
    return candidates[0]


def compute_crc(blob: bytes, expected_offset: int) -> int:
    """CRC32 of the whole file, skipping the 4 `expected` bytes."""
    if expected_offset < 16 or expected_offset + 4 > len(blob):
        raise SlotError(
            f"slot offset 0x{expected_offset:X} out of range (size 0x{len(blob):X})"
        )
    crc = 0xFFFFFFFF
    crc = binascii.crc32(blob[:expected_offset], crc)
    crc = binascii.crc32(blob[expected_offset + 4 :], crc)
    return crc & 0xFFFFFFFF


def read_slot(blob: bytes) -> tuple[int, int, int]:
    pos = find_slot(blob)
    magic, offset, expected = struct.unpack_from("<III", blob, pos)
    return pos, offset, expected


def cmd_verify(path: Path) -> int:
    blob = path.read_bytes()
    try:
        pos, offset, expected = read_slot(blob)
    except SlotError as exc:
        print(f"FAIL {path.name}: {exc}")
        return 1
    print(f"slot   : file 0x{pos:08X}  offset=0x{offset:08X}  expected=0x{expected:08X}")
    if offset != pos + EXPECTED_FIELD:
        print(f"FAIL {path.name}: slot offset does not point at itself")
        return 1
    actual = compute_crc(blob, offset)
    print(f"crc32  : computed=0x{actual:08X}  stored=0x{expected:08X}")
    if actual != expected:
        print(f"FAIL {path.name}: CRC mismatch (build not patched, or file modified)")
        return 1
    print(f"OK   {path.name}: self-integrity CRC verified")
    return 0


def cmd_patch(src: Path, dst: Path | None) -> int:
    blob = bytearray(src.read_bytes())
    try:
        pos = find_slot(blob)
    except SlotError as exc:
        print(f"FAIL {src.name}: {exc}")
        return 1

    expected_offset = pos + EXPECTED_FIELD
    # offset must be stored before computing: it is part of the hashed range.
    struct.pack_into("<I", blob, pos + OFFSET_FIELD, expected_offset)
    crc = compute_crc(bytes(blob), expected_offset)
    struct.pack_into("<I", blob, pos + EXPECTED_FIELD, crc)

    out = dst or src
    out.write_bytes(bytes(blob))
    print(f"slot   : file 0x{pos:08X}  offset=0x{expected_offset:08X}  expected=0x{crc:08X}")
    print(f"OK   {out.name}: patched ({len(blob)} bytes)")

    # Re-read and verify what we just wrote.
    return cmd_verify(out)


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description="90CN.dll self-integrity CRC patcher")
    parser.add_argument("action", choices=["verify", "patch"])
    parser.add_argument("dll", type=Path)
    parser.add_argument("-o", "--out", type=Path, help="output path for patch (default: in place)")
    args = parser.parse_args(argv)

    if not args.dll.is_file():
        print(f"FAIL {args.dll}: not a file")
        return 2
    if args.action == "verify":
        return cmd_verify(args.dll)
    return cmd_patch(args.dll, args.out)


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