#!/usr/bin/env python3
from __future__ import annotations

import argparse
import struct
from pathlib import Path


MAGIC_TTAB = 0x42415454  # 'TTAB'
HDR_FMT = "<6I"


def rol32(x: int, r: int) -> int:
    r &= 31
    return ((x << r) | (x >> (32 - r))) & 0xFFFFFFFF


def prng_step(state: int, tweak: int) -> int:
    x = (state ^ tweak) & 0xFFFFFFFF
    x ^= (x << 13) & 0xFFFFFFFF
    x ^= (x >> 17) & 0xFFFFFFFF
    x ^= (x << 5) & 0xFFFFFFFF
    return x & 0xFFFFFFFF


def alu_op_int(a: int, b: int) -> int:
    return (a + b) & 0xFFFFFFFF


def hash_op_int(a: int, b: int) -> int:
    return (a + b + 0x9E3779B9) & 0xFFFFFFFF


def load_firmware(path: Path):
    blob = path.read_bytes()
    if len(blob) < struct.calcsize(HDR_FMT):
        raise ValueError("firmware too small")
    magic, version, code_len, data_len, flag_len, target = struct.unpack_from(HDR_FMT, blob, 0)
    if magic != MAGIC_TTAB or version != 1:
        raise ValueError("bad firmware header")

    off = struct.calcsize(HDR_FMT)
    code_sz = code_len * 4
    data_sz = data_len * 4
    if len(blob) < off + code_sz + data_sz:
        raise ValueError("truncated firmware")

    code = list(struct.unpack_from(f"<{code_len}I", blob, off))
    off += code_sz
    data = list(struct.unpack_from(f"<{data_len}I", blob, off))
    return {
        "code": code,
        "data": data,
        "flag_len": flag_len,
        "target": target,
    }


def solve(
    fw,
    prefix: str | None,
    suffix: str | None,
    alphabet: str | None,
    no_format: bool,
) -> bytes:
    data = fw["data"]
    flag_len = fw["flag_len"]
    target = fw["target"]
    if flag_len < 2 or len(data) < 4:
        raise ValueError("bad firmware constants")

    seed_r14 = data[0]
    seed_r15 = data[1]
    seed_prng = data[2]

    output_base = 3 + 3 * flag_len
    if len(data) < output_base + flag_len + 1:
        raise ValueError("bad firmware layout")

    consts = data[3:output_base]
    c_mul = consts[0::3]
    c_rot = consts[1::3]
    c_tw = consts[2::3]

    outputs = data[output_base:output_base + flag_len]
    target2 = data[output_base + flag_len]

    recovered = []
    r15 = seed_r15
    for out in outputs:
        b = (out - r15 - 0x9E3779B9) & 0xFFFFFFFF
        if b > 0xFF:
            raise RuntimeError("invalid output stream: byte out of range")
        recovered.append(b)
        r15 = out

    flag = bytes(recovered)

    # format constraints
    if not no_format:
        if prefix is None or suffix is None or alphabet is None:
            raise ValueError("format constraints require prefix/suffix/alphabet")
        if not flag.startswith(prefix.encode("ascii")):
            raise RuntimeError("prefix mismatch")
        if not flag.endswith(suffix.encode("ascii")):
            raise RuntimeError("suffix mismatch")
        mid = flag[len(prefix): len(flag) - len(suffix)]
        alpha = set(alphabet.encode("ascii"))
        if any(ch not in alpha for ch in mid):
            raise RuntimeError("alphabet mismatch")

    # verify accumulators match firmware targets
    prng = seed_prng
    r14 = seed_r14
    r15 = seed_r15
    mul_out = 0
    for i, ch in enumerate(flag):
        stale_mul = mul_out
        stale_prng = prng
        mul_out = (ch + c_mul[i]) & 0xFFFFFFFF
        rot = rol32(r14, c_rot[i])
        alu = alu_op_int(rot, stale_prng)
        r14 = hash_op_int(stale_mul, alu)
        prng = prng_step(prng, c_tw[i])
        r15 = hash_op_int(r15, ch)

    if r14 != target or r15 != target2:
        raise RuntimeError("verification failed against firmware targets")

    return flag


def main() -> int:
    ap = argparse.ArgumentParser(description="Recover flag from firmware.bin.")
    ap.add_argument("firmware", nargs="?", default="public/firmware.bin")
    ap.add_argument("--prefix", default="caplag{", help="known flag prefix")
    ap.add_argument("--suffix", default="}", help="known flag suffix (1 char)")
    ap.add_argument("--alphabet", default="0123456789ABCDEF", help="alphabet for middle bytes")
    ap.add_argument("--no-format", action="store_true", help="disable prefix/suffix/alphabet constraints")
    args = ap.parse_args()

    fw = load_firmware(Path(args.firmware))
    flag = solve(
        fw,
        prefix=None if args.no_format else args.prefix,
        suffix=None if args.no_format else args.suffix,
        alphabet=None if args.no_format else args.alphabet,
        no_format=args.no_format,
    )
    try:
        print(flag.decode("ascii"))
    except UnicodeDecodeError:
        print(flag)
    return 0


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