pm3py/pm3py/core/transport.py

"""Proxmark3 NG wire protocol frame encoding/decoding and async serial I/O."""
import asyncio
import struct
from dataclasses import dataclass
from typing import Callable, Awaitable

from .protocol import (
    Cmd, PM3Status, crc16_a,
    CMD_PREAMBLE_MAGIC, RESP_PREAMBLE_MAGIC,
    CMD_POSTAMBLE_NOCRC, RESP_POSTAMBLE_NOCRC,
    CMD_PREAMBLE_SIZE, CMD_POSTAMBLE_SIZE,
    RESP_PREAMBLE_SIZE, RESP_POSTAMBLE_SIZE,
    PM3_CMD_DATA_SIZE, MIX_ARGS_SIZE,
)


def _crc_to_wire(crc: int) -> int:
    """Byte-swap CRC to match PM3 wire format: (low << 8) | high."""
    return ((crc & 0xFF) << 8) | ((crc >> 8) & 0xFF)


def encode_ng_frame(cmd: int, payload: bytes = b"", use_crc: bool = False) -> bytes:
    """Encode an NG command frame. use_crc=False for USB (default), True for FPC UART."""
    if len(payload) > PM3_CMD_DATA_SIZE:
        raise ValueError(f"Payload too large: {len(payload)} > {PM3_CMD_DATA_SIZE}")
    length_ng = len(payload) | 0x8000  # set ng bit
    preamble = struct.pack("<IHH", CMD_PREAMBLE_MAGIC, length_ng, cmd)
    body = preamble + payload
    if use_crc:
        crc = _crc_to_wire(crc16_a(body))
    else:
        crc = CMD_POSTAMBLE_NOCRC
    postamble = struct.pack("<H", crc)
    return body + postamble


def encode_mix_frame(cmd: int, arg0: int = 0, arg1: int = 0, arg2: int = 0,
                     payload: bytes = b"", use_crc: bool = False) -> bytes:
    """Encode a MIX command frame (NG wire format with OLD-style args). use_crc=False for USB."""
    args = struct.pack("<QQQ", arg0, arg1, arg2)
    full_payload = args + payload
    if len(full_payload) > PM3_CMD_DATA_SIZE:
        raise ValueError(f"MIX payload too large: {len(full_payload)} > {PM3_CMD_DATA_SIZE}")
    length_ng = len(full_payload) & 0x7FFF  # ng bit NOT set
    preamble = struct.pack("<IHH", CMD_PREAMBLE_MAGIC, length_ng, cmd)
    body = preamble + full_payload
    if use_crc:
        crc = _crc_to_wire(crc16_a(body))
    else:
        crc = CMD_POSTAMBLE_NOCRC
    postamble = struct.pack("<H", crc)
    return body + postamble


def decode_response_frame(data: bytes) -> dict:
    """Decode a response frame into a dict."""
    if len(data) < RESP_PREAMBLE_SIZE + RESP_POSTAMBLE_SIZE:
        raise ValueError(f"Frame too short: {len(data)}")

    magic, length_ng, status, reason, cmd = struct.unpack_from("<IHbbH", data, 0)
    if magic != RESP_PREAMBLE_MAGIC:
        raise ValueError(f"Bad response magic: 0x{magic:08X}")

    ng = bool(length_ng & 0x8000)
    payload_len = length_ng & 0x7FFF
    expected_len = RESP_PREAMBLE_SIZE + payload_len + RESP_POSTAMBLE_SIZE
    if len(data) < expected_len:
        raise ValueError(f"Frame truncated: got {len(data)}, expected {expected_len}")

    payload = data[RESP_PREAMBLE_SIZE:RESP_PREAMBLE_SIZE + payload_len]
    crc_received = struct.unpack_from("<H", data, RESP_PREAMBLE_SIZE + payload_len)[0]

    # Verify CRC unless it's the no-CRC magic
    if crc_received != RESP_POSTAMBLE_NOCRC:
        crc_expected = _crc_to_wire(crc16_a(data[:RESP_PREAMBLE_SIZE + payload_len]))
        if crc_received != crc_expected:
            raise ValueError(f"CRC mismatch: got 0x{crc_received:04X}, expected 0x{crc_expected:04X}")

    result = {
        "cmd": cmd,
        "status": status,
        "reason": reason,
        "ng": ng,
        "data": payload,
    }

    # For MIX frames, also extract the legacy args
    if not ng and payload_len >= MIX_ARGS_SIZE:
        arg0, arg1, arg2 = struct.unpack_from("<QQQ", payload, 0)
        result["oldarg"] = [arg0, arg1, arg2]
        result["data"] = payload[MIX_ARGS_SIZE:]

    return result


class PM3Error(Exception):
    """Proxmark3 communication error."""
    def __init__(self, message: str, status: int = 0):
        super().__init__(message)
        self.status = status


@dataclass
class PM3Response:
    """Parsed response from device."""
    cmd: int
    status: int
    reason: int
    ng: bool
    data: bytes
    oldarg: list[int] | None = None

    @classmethod
    def from_dict(cls, d: dict) -> "PM3Response":
        return cls(
            cmd=d["cmd"], status=d["status"], reason=d["reason"],
            ng=d["ng"], data=d["data"], oldarg=d.get("oldarg"),
        )


ProgressCallback = Callable[[dict], Awaitable[None]] | Callable[[dict], None] | None


class PM3Transport:
    """Async serial transport for Proxmark3 NG protocol."""

    def __init__(self, port: str, baudrate: int = 115200):
        self.port = port
        self.baudrate = baudrate
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._lock = asyncio.Lock()

    async def connect(self) -> None:
        import serial_asyncio
        self._reader, self._writer = await serial_asyncio.open_serial_connection(
            url=self.port,
            baudrate=self.baudrate,
            bytesize=8,
            parity="N",
            stopbits=1,
            xonxoff=False,
            rtscts=False,
            dsrdtr=False,
        )
        # Match C client: no flow control, flush stale data
        transport = self._writer.transport
        if hasattr(transport, 'serial'):
            transport.serial.reset_input_buffer()
            transport.serial.reset_output_buffer()

    async def disconnect(self) -> None:
        if self._writer:
            self._writer.close()
            await self._writer.wait_closed()
            self._writer = None
            self._reader = None

    async def send_frame(self, frame: bytes) -> None:
        if not self._writer:
            raise PM3Error("Not connected")
        self._writer.write(frame)
        await self._writer.drain()

    async def read_response(self, timeout: float = 2.0) -> PM3Response:
        """Read and decode one response frame from the device."""
        if not self._reader:
            raise PM3Error("Not connected")

        try:
            # Read until we find response magic
            raw = await asyncio.wait_for(
                self._read_response_frame(), timeout=timeout,
            )
            d = decode_response_frame(raw)
            return PM3Response.from_dict(d)
        except asyncio.TimeoutError:
            raise PM3Error("Response timeout", status=-4)

    async def _read_response_frame(self) -> bytes:
        """Read a complete response frame, scanning for magic bytes."""
        assert self._reader is not None
        buf = bytearray()

        # Scan for response magic
        while True:
            b = await self._reader.read(1)
            if not b:
                raise PM3Error("Connection closed")
            buf.append(b[0])
            if len(buf) >= 4:
                magic = struct.unpack_from("<I", buf, len(buf) - 4)[0]
                if magic == RESP_PREAMBLE_MAGIC:
                    # Found magic — trim everything before it
                    buf = bytearray(buf[len(buf) - 4:])
                    break

        # Read rest of preamble (6 more bytes: length_ng + status + reason + cmd)
        remaining_preamble = RESP_PREAMBLE_SIZE - 4
        data = await self._reader.readexactly(remaining_preamble)
        buf.extend(data)

        # Parse length
        length_ng = struct.unpack_from("<H", buf, 4)[0]
        payload_len = length_ng & 0x7FFF

        # Read payload + postamble
        remaining = payload_len + RESP_POSTAMBLE_SIZE
        data = await self._reader.readexactly(remaining)
        buf.extend(data)

        return bytes(buf)

    async def send_ng(self, cmd: int, payload: bytes = b"",
                      timeout: float = 2.0) -> PM3Response:
        """Send NG command and wait for response."""
        async with self._lock:
            frame = encode_ng_frame(cmd, payload)
            await self.send_frame(frame)
            return await self.read_response(timeout=timeout)

    async def send_mix(self, cmd: int, arg0: int = 0, arg1: int = 0,
                       arg2: int = 0, payload: bytes = b"",
                       timeout: float = 2.0) -> PM3Response:
        """Send MIX command and wait for response."""
        async with self._lock:
            frame = encode_mix_frame(cmd, arg0, arg1, arg2, payload)
            await self.send_frame(frame)
            return await self.read_response(timeout=timeout)

    async def send_ng_no_response(self, cmd: int, payload: bytes = b"") -> None:
        """Send NG command without waiting for response (e.g. BREAK_LOOP)."""
        async with self._lock:
            frame = encode_ng_frame(cmd, payload)
            await self.send_frame(frame)