pm3py/pm3py/sim/sim_session.py

"""SimSession — manages active card simulation with response table."""
from __future__ import annotations

import asyncio
import struct
import threading
import time

import serial

from .frame import RFFrame
from .table_compiler import ResponseTable, TableCompiler
from pm3py.trace.format import TraceFormatter
from .transponder import Transponder
from ..core.protocol import Cmd
from ..core.transport import encode_ng_frame, decode_response_frame, RESP_PREAMBLE_MAGIC, RESP_PREAMBLE_SIZE, RESP_POSTAMBLE_SIZE

# Command IDs (must match firmware pm3_cmd.h)
CMD_SIM_TABLE_UPLOAD = 0x0950
CMD_SIM_TABLE_CLEAR = 0x0951
CMD_SIM_TABLE_UPDATE = 0x0952
CMD_HF_ISO15693_SIM_TRACE = 0x0336

# Max NG payload per frame
MAX_PAYLOAD = 512

# Default serial port for Proxmark3
DEFAULT_PORT = "/dev/ttyACM0"


class SimSession:
    """Manages an active card simulation.

    For 15693 sims, use the synchronous API — no asyncio needed:

        session = SimSession.open()
        session.start_15693(tag)
        # ... modify tag in REPL ...
        tag.sync()

    The async API (via PM3Transport) is available for 14443-A
    WTX-based sims (future).
    """

    def __init__(self, transport=None, port: serial.Serial | None = None):
        self._t = transport
        self._port = port  # raw pyserial port for synchronous I/O
        self._active = False
        self._relay_task: asyncio.Task | None = None
        self._tag_model: Transponder | None = None
        self._formatter: TraceFormatter | None = None
        self.on_field_strength: callable | None = None  # callback(adc_mv: int)

    @classmethod
    def open(cls, port: str = DEFAULT_PORT, baudrate: int = 115200) -> "SimSession":
        """Open a synchronous SimSession directly over serial.

        No asyncio, no threads. For use in the REPL:
            session = SimSession.open()
            session.start_15693(tag)
            tag._memory[0:4] = b"\\xDE\\xAD\\xBE\\xEF"
            tag.sync()
        """
        ser = serial.Serial()
        ser.port = port
        ser.baudrate = baudrate
        ser.bytesize = 8
        ser.parity = "N"
        ser.stopbits = 1
        ser.timeout = 1
        ser.xonxoff = False
        ser.rtscts = False
        ser.dsrdtr = False
        ser.open()

        # Flush any stale data from previous session
        ser.reset_input_buffer()
        ser.reset_output_buffer()
        # Give firmware a moment to settle after port open
        time.sleep(0.5)
        ser.reset_input_buffer()

        return cls(port=ser)

    async def upload_table(self, table: ResponseTable) -> None:
        """Upload response table to firmware BigBuf."""
        data = table.serialize()
        for i in range(0, len(data), MAX_PAYLOAD):
            chunk = data[i:i + MAX_PAYLOAD]
            cmd = CMD_SIM_TABLE_UPLOAD if i == 0 else CMD_SIM_TABLE_UPDATE
            await self._t.send_ng(cmd, chunk)

    async def clear_table(self) -> None:
        """Clear the firmware response table."""
        await self._t.send_ng(CMD_SIM_TABLE_CLEAR, b"")

    async def start_14a(self, tag, compile: bool = True) -> None:
        """Compile table, upload, start 14443-A sim with WTX relay loop."""
        self._tag_model = tag
        if compile:
            if hasattr(tag, '_keys_a'):
                table = TableCompiler.compile_mifare(tag)
            else:
                table = TableCompiler.compile_14a(tag)
            await self.upload_table(table)

        # Start sim (fire-and-forget — response comes when sim ends)
        uid = tag._uid
        flag_val = 0x0001  # FLAG_INTERACTIVE
        if len(uid) == 4:
            flag_val |= 0x0010
        elif len(uid) == 7:
            flag_val |= 0x0020
        elif len(uid) == 10:
            flag_val |= 0x0030
        atqa = tag.atqa if hasattr(tag, 'atqa') else b"\x04\x00"
        sak = tag.sak if hasattr(tag, 'sak') else 0x08
        payload = atqa + bytes([sak]) + uid
        await self._t.send_ng_no_response(Cmd.HF_ISO14443A_SIMULATE)
        # TODO: use send_mix for 14a sim (MIX frame format)

        self._active = True
        self._relay_task = asyncio.create_task(self._relay_loop(tag))

    def start_15693(self, tag, compile: bool = True, trace: bool = False) -> None:
        """Start 15693 sim (synchronous).

        Firmware handles standard commands (inventory, read, write) autonomously.
        After starting, modify the tag in the REPL and call tag.sync() to push
        changes to firmware emulator memory.

        Args:
            trace: If True, print live reader↔tag communication to console.
        """
        self._tag_model = tag

        # Resolve the raw serial port
        if self._port:
            ser = self._port
        elif self._t and self._t._writer:
            ser = self._t._writer.transport.serial
        else:
            raise RuntimeError("No serial port available")

        # Bind serial port to tag for tag.sync()
        tag._serial = ser

        # Start sim with UID and block_size — same format as PM3 client
        # Client sends UID as-is (firmware reverses internally)
        flags = 0x02 if trace else 0x00
        block_size = tag._block_size if hasattr(tag, '_block_size') else 4
        payload = tag._uid + bytes([block_size, flags])
        frame = encode_ng_frame(Cmd.HF_ISO15693_SIMULATE, payload)
        ser.write(frame)

        # Give firmware time to start sim before any sync
        import time
        time.sleep(0.2)

        # Push tag memory to firmware EML (always — firmware EML is fresh on sim start)
        tag._uid_dirty = True
        for r in tag.regions.values():
            r.data._dirty = True
        tag.sync()

        self._active = True

        if compile:
            self._compile_and_upload_table(tag, ser)

        if trace:
            self._formatter = TraceFormatter(mode="sim", decoder=tag.decode_trace, crc_len=2)
            self._trace_thread = threading.Thread(
                target=self._trace_reader, args=(ser,), daemon=True
            )
            self._trace_thread.start()

        print(f"[SimSession] 15693 sim started, UID={tag._uid.hex()}"
              + (" (trace ON)" if trace else ""))
        print(f"[SimSession] Modify tag in REPL, then call tag.sync()")

    def _compile_and_upload_table(self, tag, ser) -> None:
        """Compile response table from tag model and upload to firmware."""
        table = ResponseTable(entries=[])

        # Check most specific IC first, fall back to base
        from pm3py.transponders.hf.iso15693.nxp.icode3 import Icode3Tag
        from pm3py.transponders.hf.iso15693.nxp.icode_dna import IcodeDnaTag
        from pm3py.transponders.hf.iso15693.nxp.icode_slix2 import IcodeSlix2Tag
        from pm3py.transponders.hf.iso15693.nxp.ntag5_platform import Ntag5PlatformTag
        if isinstance(tag, Ntag5PlatformTag):
            # NTAG 5 Switch/Link/Boost — use DNA compiler (same platform)
            table = TableCompiler.compile_icode_dna(tag)
        elif isinstance(tag, Icode3Tag):
            table = TableCompiler.compile_icode3(tag)
        elif isinstance(tag, IcodeDnaTag):
            table = TableCompiler.compile_icode_dna(tag)
        elif isinstance(tag, IcodeSlix2Tag):
            table = TableCompiler.compile_slix2(tag)
        elif hasattr(tag, '_passwords'):
            table = TableCompiler.compile_slix2(tag)
        elif hasattr(tag, '_signature'):
            table = TableCompiler.compile_nxp_icode(tag)

        if not table.entries:
            return

        # Upload: first chunk has 4-byte initial_groups header
        initial_groups = 0xFFFFFFFF
        data = table.serialize()
        entry_size = 120

        # First frame
        header = struct.pack("<I", initial_groups)
        entries_per_chunk = (MAX_PAYLOAD - len(header)) // entry_size
        first_data = data[:entries_per_chunk * entry_size]
        payload = header + first_data
        frame = encode_ng_frame(CMD_SIM_TABLE_UPLOAD, payload)
        ser.write(frame)
        time.sleep(0.05)
        if ser.in_waiting:
            ser.read(ser.in_waiting)

        # Remaining chunks via UPDATE
        remaining = data[entries_per_chunk * entry_size:]
        chunk_size = (MAX_PAYLOAD // entry_size) * entry_size
        for i in range(0, len(remaining), chunk_size):
            chunk = remaining[i:i + chunk_size]
            frame = encode_ng_frame(CMD_SIM_TABLE_UPDATE, chunk)
            ser.write(frame)
            time.sleep(0.05)
            if ser.in_waiting:
                ser.read(ser.in_waiting)

        print(f"[SimSession] Uploaded {len(table.entries)} table entries")

    def _trace_reader(self, ser: serial.Serial) -> None:
        """Background thread: read trace notifications from firmware and print them."""
        buf = bytearray()
        while self._active:
            try:
                chunk = ser.read(ser.in_waiting or 1)
                if not chunk:
                    continue
                buf.extend(chunk)

                # Scan for response frames
                while len(buf) >= RESP_PREAMBLE_SIZE + RESP_POSTAMBLE_SIZE:
                    # Find magic
                    idx = buf.find(struct.pack("<I", RESP_PREAMBLE_MAGIC))
                    if idx < 0:
                        buf = buf[-3:]  # keep last 3 bytes (partial magic)
                        break
                    if idx > 0:
                        buf = buf[idx:]

                    if len(buf) < RESP_PREAMBLE_SIZE:
                        break

                    length_ng = struct.unpack_from("<H", buf, 4)[0]
                    payload_len = length_ng & 0x7FFF
                    frame_len = RESP_PREAMBLE_SIZE + payload_len + RESP_POSTAMBLE_SIZE

                    if len(buf) < frame_len:
                        break  # wait for more data

                    try:
                        resp = decode_response_frame(bytes(buf[:frame_len]))
                        buf = buf[frame_len:]

                        if resp["cmd"] == CMD_HF_ISO15693_SIM_TRACE:
                            data = resp["data"]
                            if len(data) >= 2:
                                direction = data[0] & 0x0F
                                if direction == 0x03 and len(data) >= 3:
                                    # ADC field strength report
                                    adc_mv = (data[1] << 8) | data[2]
                                    if self.on_field_strength:
                                        self.on_field_strength(adc_mv)
                                else:
                                    crc_fail = bool(data[0] & 0x80)
                                    payload = data[1:]
                                    self._formatter.print(direction, payload, crc_fail=crc_fail)
                    except Exception:
                        buf = buf[4:]  # skip bad magic, try again

            except serial.SerialException:
                break
            except Exception:
                continue

    async def _relay_loop(self, tag: Transponder) -> None:
        """Background: listen for WTX relay events (14443-A Layer 4 only).

        15693 relay was removed — can't meet 311µs FDT over USB.
        This loop is for future 14443-A ISO-DEP WTX relay where
        S(WTX) gives seconds of response time.
        """
        while self._active:
            try:
                resp = await self._t.read_response(timeout=0.5)
            except Exception:
                continue

            if resp is None:
                continue

            elif resp.cmd == Cmd.HF_ISO15693_SIMULATE:
                # Sim ended (firmware sent final reply)
                print(f"[SimSession] Sim ended, status={resp.status}")
                self._active = False
                break

    async def sync(self) -> None:
        """Push tag model state to firmware emulator memory.

        Call this after modifying the tag in the REPL to update
        what the firmware serves for standard commands.
        """
        tag = self._tag_model
        if tag is None:
            return

        # Build iso15_tag_t header (offsets match firmware struct)
        uid_reversed = tag._uid[::-1]  # firmware stores UID reversed
        block_size = tag._block_size if hasattr(tag, '_block_size') else 4
        num_blocks = tag._num_blocks if hasattr(tag, '_num_blocks') else 64
        dsfid = tag._dsfid if hasattr(tag, '_dsfid') else 0
        afi = tag._afi if hasattr(tag, '_afi') else 0
        ic_ref = tag._ic_reference if hasattr(tag, '_ic_reference') else 0

        header = bytearray(15)
        header[0:8] = uid_reversed          # uid[8]
        header[8] = dsfid                    # dsfid
        header[9] = 0                        # dsfidLock
        header[10] = afi                     # afi
        header[11] = 0                       # afiLock
        header[12] = block_size              # bytesPerPage
        header[13] = num_blocks              # pagesCount
        header[14] = ic_ref                  # ic

        # Write header to EML offset 0
        await self._eml_set(0, header)

        # Write tag memory data at offset 175 (after uid+config+locks)
        # locks[160] at offset 15, data[2048] at offset 175
        memory = tag._memory if hasattr(tag, '_memory') else b""
        if memory:
            # Chunk memory writes (max ~500 bytes per EML_SETMEM)
            offset = 175
            chunk_size = 480
            for i in range(0, len(memory), chunk_size):
                chunk = bytes(memory[i:i + chunk_size])
                await self._eml_set(offset + i, chunk)

        print(f"[SimSession] Synced tag state to firmware (uid={tag._uid.hex()}, "
              f"{num_blocks} blocks of {block_size} bytes)")

    async def _eml_set(self, offset: int, data: bytes) -> None:
        """Write data to firmware emulator memory at offset."""
        payload = struct.pack("<IH", offset, len(data)) + data
        await self._t.send_ng_no_response(Cmd.HF_ISO15693_EML_SETMEM, payload)

    def stop(self) -> None:
        """Stop simulation."""
        self._active = False
        if self._port:
            frame = encode_ng_frame(Cmd.BREAK_LOOP)
            self._port.write(frame)
            print("[SimSession] Sim stopped")
        elif self._t:
            # Fall back to async transport
            asyncio.get_event_loop().run_until_complete(
                self._t.send_ng_no_response(Cmd.BREAK_LOOP)
            )
        if self._relay_task:
            self._relay_task.cancel()
            self._relay_task = None

    def close(self) -> None:
        """Stop sim and close serial port."""
        self.stop()
        if self._port:
            self._port.close()
            self._port = None
        if self._tag_model:
            self._tag_model._serial = None