#!/usr/bin/env python3
"""Overstrike: recover the required CarriedMark and decrypt SealedRecord.

Stdlib only. Reads the byte blob straight out of Overstrike.dll (PE RVA ->
file offset), runs GameState.Mix backwards from the hardcoded target seal to
the mark that produces it, then rebuilds the SHA-256 keystream and XORs.

Usage:
    python3 unseal.py Overstrike.dll
"""

from __future__ import annotations

import argparse
import hashlib
import logging
import struct
import sys
from pathlib import Path

LOG = logging.getLogger("unseal")

CONFIG = {
    # from the .NET metadata: FieldRVA of the SealedRecord array-init blob,
    # and __StaticArrayInitTypeSize=56
    "record_rva": 0x2070,
    "record_len": 56,
    # GameState.get_WorldIsAligned: ldc.i8 -2764723033133996666
    "target_seal": (-2764723033133996666) & 0xFFFFFFFFFFFFFFFF,
    # GameState.Mix constants (splitmix64 finalizer)
    "gamma": 0x9E3779B97F4A7C15,
    "m1": 0xBF58476D1CE4E5B9,
    "m2": 0x94D049BB133111EB,
}

MASK64 = 0xFFFFFFFFFFFFFFFF


def rva_to_offset(data: bytes, rva: int) -> int:
    """Translate a PE relative virtual address to a file offset."""
    if data[:2] != b"MZ":
        raise ValueError("not a PE file (missing MZ header)")
    pe_off = struct.unpack_from("<I", data, 0x3C)[0]
    if data[pe_off:pe_off + 4] != b"PE\0\0":
        raise ValueError("not a PE file (missing PE signature)")
    n_sections = struct.unpack_from("<H", data, pe_off + 6)[0]
    opt_size = struct.unpack_from("<H", data, pe_off + 20)[0]
    sec_table = pe_off + 24 + opt_size
    for i in range(n_sections):
        base = sec_table + i * 40
        va, raw_size, raw_ptr = struct.unpack_from("<III", data, base + 12)
        if va <= rva < va + max(raw_size, 1):
            return raw_ptr + (rva - va)
    raise ValueError(f"RVA {rva:#x} is not inside any section")


def mix(x: int) -> int:
    """GameState.Mix - the splitmix64 finalizer, forwards."""
    x = (x + CONFIG["gamma"]) & MASK64
    x ^= x >> 30
    x = (x * CONFIG["m1"]) & MASK64
    x ^= x >> 27
    x = (x * CONFIG["m2"]) & MASK64
    return x ^ (x >> 31)


def unxorshift(y: int, shift: int) -> int:
    """Undo x ^= x >> shift by rebuilding the value from the top bits down."""
    x = y
    while True:
        candidate = y ^ (x >> shift)
        if candidate == x:
            return x & MASK64
        x = candidate


def unmix(y: int) -> int:
    """GameState.Mix, backwards. Every step is invertible, so no search."""
    x = unxorshift(y, 31)
    x = (x * pow(CONFIG["m2"], -1, 1 << 64)) & MASK64
    x = unxorshift(x, 27)
    x = (x * pow(CONFIG["m1"], -1, 1 << 64)) & MASK64
    x = unxorshift(x, 30)
    return (x - CONFIG["gamma"]) & MASK64


def unseal(record: bytes, mark: int) -> bytes:
    """Rebuild the SHA-256 counter keystream for `mark` and XOR the record."""
    seed = hashlib.sha256(mark.to_bytes(8, "little")).digest()
    stream = bytearray()
    counter = 0
    while len(stream) < len(record):
        stream += hashlib.sha256(seed + struct.pack("<i", counter)).digest()
        counter += 1
    return bytes(a ^ b for a, b in zip(record, stream))


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Recover the Overstrike registry text from Overstrike.dll.",
        epilog="example: python3 unseal.py ext/assets/.godot/mono/publish/x86_64/Overstrike.dll",
    )
    parser.add_argument("dll", type=Path, help="path to Overstrike.dll")
    parser.add_argument("--rva", type=lambda s: int(s, 0),
                        default=CONFIG["record_rva"],
                        help="FieldRVA of SealedRecord (default 0x2070)")
    parser.add_argument("--length", type=int, default=CONFIG["record_len"],
                        help="length of SealedRecord in bytes (default 56)")
    parser.add_argument("-v", "--verbose", action="store_true", help="debug output")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)-7s %(message)s",
        datefmt="%H:%M:%S",
    )

    try:
        data = args.dll.read_bytes()
    except OSError as exc:
        LOG.error("cannot read %s: %s", args.dll, exc)
        return 1

    try:
        offset = rva_to_offset(data, args.rva)
    except ValueError as exc:
        LOG.error("%s", exc)
        return 1

    record = data[offset:offset + args.length]
    if len(record) != args.length:
        LOG.error("file ends before the record does (got %d bytes)", len(record))
        return 1
    LOG.debug("record @ file offset %#x: %s", offset, record.hex())

    mark = unmix(CONFIG["target_seal"])
    if mix(mark) != CONFIG["target_seal"]:
        LOG.error("inversion failed - Mix constants do not match this binary")
        return 1

    LOG.info("target WorldSeal : %#018x", CONFIG["target_seal"])
    LOG.info("required mark    : %#018x  (%d)", mark, mark)

    plain = unseal(record, mark)
    printable = all(32 <= b < 127 for b in plain)
    LOG.info("decrypted (%s): %s", "clean ASCII" if printable else "NOT ascii",
             plain.decode("ascii", "replace"))
    return 0 if printable else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\ninterrupted", file=sys.stderr)
        sys.exit(130)