#!/usr/bin/env python3
"""90CN.dll self-integrity CRC32 patcher.

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.

IMPORTANT — the magic is NOT unique in the image. MSVC encodes the comparison
`magic != 0xC0DEC0DEu` as an immediate inside .text, so a naive "find the first
0xC0DEC0DE" scan lands on machine code and overwrites instructions, which shows
up at runtime as an access violation right after the slot read. This tool
therefore resolves the slot through the PE section table and only accepts a
candidate that lives in a data section (.data/.rdata) and whose offset/expected
fields are either zero (fresh build) or already self-consistent.

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 info   <dll>     # show slot + section layout
    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
import zlib
from pathlib import Path

MAGIC = 0xC0DEC0DE
MAGIC_BYTES = struct.pack("<I", MAGIC)
SLOT_SIZE = 12
OFFSET_FIELD = 4
EXPECTED_FIELD = 8

# Section names whose raw data may hold the slot. MSVC puts a zero-initialised
# (but non-const) static object in .data; if it is truly constant it can end up
# in .rdata. Anything executable is rejected outright.
SLOT_SECTIONS = (".data", ".rdata")


class SlotError(RuntimeError):
    pass


def parse_sections(blob: bytes) -> list[dict]:
    if len(blob) < 0x40 or blob[:2] != b"MZ":
        raise SlotError("not a PE image (missing MZ header)")
    pe = struct.unpack_from("<I", blob, 0x3C)[0]
    if blob[pe : pe + 4] != b"PE\0\0":
        raise SlotError("not a PE image (missing PE signature)")
    nsec = struct.unpack_from("<H", blob, pe + 6)[0]
    opt_size = struct.unpack_from("<H", blob, pe + 20)[0]
    sh = pe + 24 + opt_size
    sections = []
    for i in range(nsec):
        base = sh + i * 40
        if base + 40 > len(blob):
            break
        name = blob[base : base + 8].rstrip(b"\0").decode("latin1")
        vsize, va, rsize, roff = struct.unpack_from("<IIII", blob, base + 8)
        chars = struct.unpack_from("<I", blob, base + 36)[0]
        sections.append(
            {
                "name": name,
                "vsize": vsize,
                "rva": va,
                "rawsize": rsize,
                "rawoff": roff,
                "executable": bool(chars & 0x20000000),
                "writable": bool(chars & 0x80000000),
            }
        )
    if not sections:
        raise SlotError("PE has no section table")
    return sections


def section_of(sections: list[dict], file_offset: int) -> dict | None:
    for sec in sections:
        if sec["rawoff"] <= file_offset < sec["rawoff"] + sec["rawsize"]:
            return sec
    return None


def find_slot(blob: bytes) -> tuple[int, dict]:
    """Return (file offset, section) of the SelfCrcSlot struct."""
    sections = parse_sections(blob)
    candidates: list[tuple[int, dict]] = []
    pos = blob.find(MAGIC_BYTES)
    while pos >= 0:
        sec = section_of(sections, pos)
        if sec is not None and sec["name"] in SLOT_SECTIONS:
            candidates.append((pos, sec))
        pos = blob.find(MAGIC_BYTES, pos + 1)
    if not candidates:
        # Distinguish "magic only in code" from "magic nowhere": give a useful hint.
        if blob.find(MAGIC_BYTES) >= 0:
            raise SlotError(
                "0xC0DEC0DE only found in an executable section — refusing to "
                "patch machine code (the real slot lives in .data/.rdata)"
            )
        raise SlotError("0xC0DEC0DE slot magic not found")

    def score(item: tuple[int, dict]) -> tuple[int, int]:
        p, sec = item
        offset = struct.unpack_from("<I", blob, p + OFFSET_FIELD)[0]
        fresh = offset == 0
        patched = offset == p + EXPECTED_FIELD
        # Prefer a patched slot, then a fresh one, then .data over .rdata.
        return (2 if patched else 1 if fresh else 0, 1 if sec["name"] == ".data" else 0)

    candidates.sort(key=score, reverse=True)
    best = candidates[0]
    if score(best)[0] == 0:
        raise SlotError(
            "slot magic found in data but offset field is neither 0 nor self-referential"
        )
    return best


def compute_crc(blob: bytes, expected_offset: int) -> int:
    """CRC32 of the whole file, skipping the 4 `expected` bytes.

    Must match 90CN.cpp VerifySelfIntegrity() exactly:

        unsigned int crc = 0xFFFFFFFF;
        crc = Crc32Update(crc, image, slotOffset);
        crc = Crc32Update(crc, image + slotOffset + 4, size - slotOffset - 4);
        crc ^= 0xFFFFFFFF;

    Crc32Update is the reflected CRC-32 (poly 0xEDB88320) with no internal
    conditioning, so "init 0xFFFFFFFF ... final XOR 0xFFFFFFFF" is exactly
    zlib.crc32()/binascii.crc32() with the DEFAULT seed 0.

    Do NOT seed binascii.crc32 with 0xFFFFFFFF: that argument is the running
    CRC *after* the final XOR, so seeding it re-applies the conditioning and
    yields a different, non-matching variant (this bug shipped once and made
    the client reject its own DLL).
    """
    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 = zlib.crc32(blob[:expected_offset])
    crc = zlib.crc32(blob[expected_offset + 4 :], crc)
    return crc & 0xFFFFFFFF


def read_slot(blob: bytes) -> tuple[int, dict, int, int]:
    pos, sec = find_slot(blob)
    magic, offset, expected = struct.unpack_from("<III", blob, pos)
    if magic != MAGIC:
        raise SlotError("slot magic mismatch after read")
    return pos, sec, offset, expected


def cmd_info(path: Path) -> int:
    blob = path.read_bytes()
    sections = parse_sections(blob)
    print(f"file   : {path}  ({len(blob)} bytes)")
    for sec in sections:
        print(
            f"  {sec['name']:9s} rva=0x{sec['rva']:08X} vsize={sec['vsize']:8d} "
            f"raw=0x{sec['rawoff']:08X} rawsize={sec['rawsize']:8d}"
            f"{'  EXEC' if sec['executable'] else ''}"
            f"{'  WRITABLE' if sec['writable'] else ''}"
        )
    try:
        pos, sec, offset, expected = read_slot(blob)
    except SlotError as exc:
        print(f"slot   : NOT FOUND ({exc})")
        return 1
    print(
        f"slot   : file 0x{pos:08X} in {sec['name']} "
        f"offset=0x{offset:08X} expected=0x{expected:08X}"
    )
    return 0


def cmd_verify(path: Path) -> int:
    blob = path.read_bytes()
    try:
        pos, sec, offset, expected = read_slot(blob)
    except SlotError as exc:
        print(f"FAIL {path.name}: {exc}")
        return 1
    print(
        f"slot   : file 0x{pos:08X} in {sec['name']}  "
        f"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, sec = 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} in {sec['name']}  "
        f"offset=0x{expected_offset:08X}  expected=0x{crc:08X}"
    )
    print(f"OK   {out.name}: patched ({len(blob)} bytes)")
    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=["info", "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 == "info":
        return cmd_info(args.dll)
    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:]))