#!/usr/bin/env python3
"""Oracle for the dwt_prf transform on the live MAME terminal.

Each query opens a fresh connection, sends opcode 7, reads the 16-byte
challenge, sends our own 16-byte reply, and then reads the reply buffer that
the firmware dumped at the memcmp (dwt/evidence/run-live/cmp_a11.bin).
The transform is per 8-byte record, so the first 8 output bytes answer the
first 8 input bytes.
"""
import socket, struct, sys, time, re

HOST, PORT = "192.168.99.2", 771
DUMP = "/home/arti/code/k-space/vt1200/dwt/evidence/run-live/cmp_a11.bin"


def recv_frame(s):
    h = s.recv(4)
    if len(h) < 4:
        raise EOFError
    n = struct.unpack(">I", h)[0]
    b = b""
    while len(b) < n:
        c = s.recv(n - len(b))
        if not c:
            break
        b += c
    return b


def parse_dump(path):
    words = []
    try:
        text = open(path).read()
    except FileNotFoundError:
        return b""
    for line in text.splitlines():
        m = re.match(r"^[0-9A-F]+:\s+(.*)$", line.strip())
        if not m:
            continue
        for w in m.group(1).split():
            if re.fullmatch(r"[0-9A-F]{4}", w):
                v = int(w, 16)
                words.append(v & 0xFF)
                words.append(v >> 8)
    return bytes(words)


def query(reply: bytes, settle=0.7) -> bytes:
    assert len(reply) == 16
    s = socket.create_connection((HOST, PORT), 5)
    s.settimeout(5)
    s.sendall(struct.pack(">I", 7))
    ch = recv_frame(s)
    if ch[0] != 0x7F:
        raise RuntimeError(f"no challenge: {ch.hex()}")
    s.sendall(reply)
    time.sleep(0.2)
    s.close()
    time.sleep(settle)
    return parse_dump(DUMP)


def main():
    cases = {}
    if len(sys.argv) > 1:
        for arg in sys.argv[1:]:
            name, hexs = arg.split("=", 1)
            cases[name] = bytes.fromhex(hexs)
    else:
        cases = {
            "zero": bytes(16),
            "p0_01": bytes.fromhex("01" + "00" * 15),
            "p0_02": bytes.fromhex("02" + "00" * 15),
            "p0_03": bytes.fromhex("03" + "00" * 15),
            "p0_10": bytes.fromhex("10" + "00" * 15),
            "p0_20": bytes.fromhex("20" + "00" * 15),
            "p0_30": bytes.fromhex("30" + "00" * 15),
            "p0_80": bytes.fromhex("80" + "00" * 15),
            "p0_ff": bytes.fromhex("ff" + "00" * 15),
            "p1_01": bytes.fromhex("00" + "01" + "00" * 14),
            "p7_01": bytes.fromhex("00" * 7 + "01" + "00" * 8),
        }
    for name, reply in cases.items():
        out = query(reply)
        print(f"{name}\tin={reply[:8].hex(' ')}\tout={out[:8].hex(' ')}")
        sys.stdout.flush()


if __name__ == "__main__":
    main()