pm3py/pm3py/sim/iso14443a.py

"""ISO 14443-A transponder and reader state machines."""
from __future__ import annotations

from enum import IntEnum

from .frame import RFFrame
from .medium import Medium
from .trace_fmt import decode_14443a
from .transponder import Transponder
from .reader import Reader

# ---- Constants ----

REQA = 0x26
WUPA = 0x52
HLTA = 0x50
CL1 = 0x93
CL2 = 0x95
CL3 = 0x97
CT = 0x88  # Cascade Tag
RATS_CMD = 0xE0

NVB_ANTICOL = 0x20  # 2 bytes valid (SEL + NVB only)
NVB_SELECT = 0x70   # 7 bytes valid (full SELECT)

CL_MAP = {CL1: 0, CL2: 1, CL3: 2}


class State14443A(IntEnum):
    IDLE = 0
    READY = 1
    ACTIVE = 2
    HALT = 3
    PROTOCOL = 4  # Layer 4 (ISO-DEP) active


def _compute_bcc(data: bytes) -> int:
    r = 0
    for b in data:
        r ^= b
    return r


def _split_uid_cascades(uid: bytes) -> list[bytes]:
    """Split UID into cascade-level chunks of 4 bytes each."""
    if len(uid) == 4:
        return [uid]
    elif len(uid) == 7:
        return [bytes([CT]) + uid[0:3], uid[3:7]]
    elif len(uid) == 10:
        return [bytes([CT]) + uid[0:3], bytes([CT]) + uid[3:6], uid[6:10]]
    else:
        raise ValueError(f"UID must be 4, 7, or 10 bytes, got {len(uid)}")


class Tag14443A(Transponder):
    """ISO 14443-A Part 2+3 base: anticollision + selection."""

    def __init__(self, uid: bytes | str, atqa: bytes = b"\x04\x00", sak: int = 0x08):
        super().__init__()
        uid = self._parse_uid(uid)
        if len(uid) not in (4, 7, 10):
            raise ValueError(f"UID must be 4, 7, or 10 bytes, got {len(uid)}")
        self._uid = uid
        self._atqa = atqa
        self._sak = sak
        self._state = State14443A.IDLE
        self._cascades = _split_uid_cascades(uid)
        self._current_cl = 0

    def decode_trace(self, direction: int, payload: bytes) -> str | None:
        return decode_14443a(direction, payload)

    @property
    def atqa(self) -> bytes:
        return self._atqa

    @property
    def sak(self) -> int:
        return self._sak

    def _cascade_entries(self) -> list[tuple[int, bytes]]:
        """Return (sel_byte, uid_chunk) for each cascade level."""
        sels = [CL1, CL2, CL3]
        return list(zip(sels[:len(self._cascades)], self._cascades))

    async def power_on(self) -> None:
        self._state = State14443A.IDLE
        self._current_cl = 0

    async def power_off(self) -> None:
        self._state = State14443A.IDLE
        self._current_cl = 0

    @property
    def state(self) -> str:
        return self._state.name

    async def handle_frame(self, frame: RFFrame) -> RFFrame | None:
        if not frame.data:
            return None
        cmd = frame.data[0]

        match self._state:
            case State14443A.IDLE:
                if cmd == REQA or cmd == WUPA:
                    self._state = State14443A.READY
                    self._current_cl = 0
                    return RFFrame.from_bytes(self._atqa)
                return None

            case State14443A.READY:
                if cmd == REQA or cmd == WUPA:
                    self._current_cl = 0
                    return RFFrame.from_bytes(self._atqa)
                return self._handle_anticollision(frame)

            case State14443A.ACTIVE:
                # Application handler gets first chance (e.g., MIFARE auth response)
                app_result = self._handle_application(frame)
                if app_result is not None:
                    return app_result
                if cmd == HLTA and len(frame.data) >= 2:
                    self._state = State14443A.HALT
                    return None
                if cmd == REQA or cmd == WUPA:
                    # Re-selection: reset to READY
                    self._state = State14443A.READY
                    self._current_cl = 0
                    return RFFrame.from_bytes(self._atqa)
                if cmd == RATS_CMD:
                    return self._handle_rats(frame)
                return None

            case State14443A.HALT:
                if cmd == WUPA:
                    self._state = State14443A.READY
                    self._current_cl = 0
                    return RFFrame.from_bytes(self._atqa)
                return None

            case State14443A.PROTOCOL:
                if cmd == REQA or cmd == WUPA:
                    self._state = State14443A.READY
                    self._current_cl = 0
                    return RFFrame.from_bytes(self._atqa)
                return self._handle_layer4(frame)

        return None

    def _handle_anticollision(self, frame: RFFrame) -> RFFrame | None:
        if len(frame.data) < 2:
            return None

        sel = frame.data[0]
        nvb = frame.data[1]

        if sel not in CL_MAP:
            return None

        cl_index = CL_MAP[sel]
        if cl_index != self._current_cl:
            return None

        cascade_uid = self._cascades[cl_index]
        uid_bcc = cascade_uid + bytes([_compute_bcc(cascade_uid)])

        if nvb == NVB_SELECT and len(frame.data) >= 7:
            # Full SELECT
            incoming_uid = frame.data[2:6]
            if incoming_uid == cascade_uid:
                self._current_cl += 1
                if self._current_cl >= len(self._cascades):
                    self._state = State14443A.ACTIVE
                    return RFFrame.from_bytes(bytes([self._sak]))
                else:
                    return RFFrame.from_bytes(bytes([0x04]))
            return None

        # ANTICOLLISION with known bytes
        known_bytes = (nvb >> 4) - 2  # subtract SEL + NVB bytes
        if known_bytes < 0:
            known_bytes = 0

        # Check if our UID matches the known bytes
        if known_bytes > 0 and len(frame.data) > 2:
            known_data = frame.data[2:2 + known_bytes]
            for i in range(min(known_bytes, len(uid_bcc))):
                if i < len(known_data) and known_data[i] != uid_bcc[i]:
                    return None  # doesn't match

        return RFFrame.from_bytes(uid_bcc)

    def _handle_rats(self, frame: RFFrame) -> RFFrame | None:
        return None

    def _handle_application(self, frame: RFFrame) -> RFFrame | None:
        return None

    def _handle_layer4(self, frame: RFFrame) -> RFFrame | None:
        return None


class Tag14443A_3(Tag14443A):
    """ISO 14443-A Part 3 only — no ISO-DEP. SAK bit 5 = 0."""

    def __init__(self, uid: bytes, atqa: bytes = b"\x04\x00", sak: int = 0x08):
        super().__init__(uid=uid, atqa=atqa, sak=sak & ~0x20)


class Tag14443A_4(Tag14443A):
    """ISO 14443-A Part 3+4 — ISO-DEP capable. SAK bit 5 = 1."""

    def __init__(self, uid: bytes, atqa: bytes = b"\x04\x00", sak: int = 0x20,
                 ats: bytes = b"\x05\x78\x80\x70\x02"):
        super().__init__(uid=uid, atqa=atqa, sak=sak | 0x20)
        self._ats = ats
        self._block_number = 0

    def _handle_rats(self, frame: RFFrame) -> RFFrame | None:
        self._state = State14443A.PROTOCOL
        self._block_number = 0
        return RFFrame.from_bytes(self._ats)

    def _handle_layer4(self, frame: RFFrame) -> RFFrame | None:
        if not frame.data:
            return None
        pcb = frame.data[0]
        if pcb & 0xC0 == 0x00:  # I-block
            payload = frame.data[1:]
            response_payload = self._handle_apdu(payload)
            resp_pcb = 0x02 if (pcb & 0x01) == 0 else 0x03
            self._block_number = resp_pcb & 0x01
            return RFFrame.from_bytes(bytes([resp_pcb]) + response_payload)
        return None

    def _handle_apdu(self, apdu: bytes) -> bytes:
        """Override in subclasses. Default: 6A82 (file not found)."""
        return b"\x6A\x82"


# ---- Reader14443A ----

class Reader14443A(Reader):
    """ISO 14443-A reader with full anticollision tree walk."""

    async def run(self) -> dict:
        uids = await self.inventory()
        return {"uids": uids}

    async def inventory(self) -> list[bytes]:
        """Perform full anticollision, return all UIDs found."""
        uids: list[bytes] = []
        while True:
            uid = await self._select_one_tag()
            if uid is None:
                break
            uids.append(uid)
            # HALT so this tag won't respond to REQA
            await self._medium.transmit_reader(RFFrame.from_hex("5000"))
            await self._medium.receive_reader()
        return uids

    async def _select_one_tag(self) -> bytes | None:
        """Wake tags with REQA, resolve collisions, return one full UID."""
        await self._medium.transmit_reader(RFFrame.from_bytes(bytes([REQA])))
        atqa = await self._medium.receive_reader()
        if atqa is None:
            return None

        uid_acc = b""
        for sel in (CL1, CL2, CL3):
            result = await self._anticol_select(sel)
            if result is None:
                return None
            sak = result["sak"]
            cl_uid = result["uid"]

            if sak & 0x04:
                uid_acc += cl_uid[1:4] if cl_uid[0] == CT else cl_uid
            else:
                return uid_acc + cl_uid

        return None

    async def _anticol_select(self, sel: int,
                              known_bytes: bytes = b"") -> dict | None:
        """Anticollision + SELECT for one cascade level.

        Uses byte-level collision resolution: when collision is detected
        at byte N, try all 256 values for that byte position until one
        produces a clean (no-collision) response.
        """
        # Build ANTICOL command with known bytes
        nvb_byte_count = len(known_bytes) + 2  # +2 for SEL + NVB
        nvb = nvb_byte_count << 4  # no extra bits
        anticol = bytes([sel, nvb]) + known_bytes
        await self._medium.transmit_reader(RFFrame.from_bytes(anticol))
        resp = await self._medium.receive_reader()
        if resp is None:
            return None

        if resp.has_collision and resp.collision_positions:
            # Find which response byte has the first collision
            first_col = resp.collision_positions[0]
            col_byte_idx = first_col // 8

            if col_byte_idx >= 4:
                # Collision only in BCC byte — use data bytes as-is
                pass
            else:
                # Try each possible value for the collision byte
                # Use values from the collision response as starting hints
                resp_byte = resp.data[col_byte_idx] if col_byte_idx < len(resp.data) else 0
                # Try the collision byte with bit set to 0 and 1 first, then others
                from bitarray import bitarray as ba
                resp_bits = ba(resp.bits[:resp.bit_count])
                resp_bits[first_col] = 0
                hint0 = resp_bits[col_byte_idx*8:(col_byte_idx+1)*8].tobytes()[0]
                resp_bits[first_col] = 1
                hint1 = resp_bits[col_byte_idx*8:(col_byte_idx+1)*8].tobytes()[0]

                candidates = [hint0, hint1]
                for candidate in candidates:
                    prefix = known_bytes[:col_byte_idx] + bytes([candidate])
                    if len(prefix) <= col_byte_idx:
                        # Need intermediate bytes from response
                        prefix = resp.data[:col_byte_idx] + bytes([candidate])
                    result = await self._anticol_select(sel, prefix)
                    if result is not None:
                        return result
                return None

        # No collision — we have a clean UID
        if len(resp.data) < 5:
            return None
        uid_bytes = resp.data[0:4]
        bcc = _compute_bcc(uid_bytes)

        # SELECT
        select_cmd = bytes([sel, NVB_SELECT]) + uid_bytes + bytes([bcc])
        await self._medium.transmit_reader(RFFrame.from_bytes(select_cmd))
        sak_resp = await self._medium.receive_reader()
        if sak_resp is None:
            return None

        return {"uid": uid_bytes, "sak": sak_resp.data[0]}

    async def select_tag(self, uid: bytes) -> dict:
        """Full SELECT sequence for a known UID."""
        cascades = _split_uid_cascades(uid)
        sels = [CL1, CL2, CL3]
        sak = 0

        await self._medium.transmit_reader(RFFrame.from_bytes(bytes([WUPA])))
        await self._medium.receive_reader()

        for i, (sel, cascade_uid) in enumerate(zip(sels, cascades)):
            bcc = _compute_bcc(cascade_uid)
            select_cmd = bytes([sel, NVB_SELECT]) + cascade_uid + bytes([bcc])
            await self._medium.transmit_reader(RFFrame.from_bytes(select_cmd))
            resp = await self._medium.receive_reader()
            if resp is None:
                raise RuntimeError(f"No response to SELECT CL{i+1}")
            sak = resp.data[0]

        return {"uid": uid, "sak": sak}

    async def rats(self, cid: int = 0) -> dict:
        """Send RATS, enter Layer 4. Returns ATS."""
        fsd_cid = 0x50 | (cid & 0x0F)
        rats_frame = RFFrame.from_bytes(bytes([RATS_CMD, fsd_cid]))
        await self._medium.transmit_reader(rats_frame)
        resp = await self._medium.receive_reader()
        if resp is None:
            raise RuntimeError("No response to RATS")
        return {"ats": resp.data}

    async def transceive_apdu(self, apdu: bytes, block_number: int = 0) -> bytes:
        """Send APDU via I-block, return response payload."""
        pcb = 0x02 | (block_number & 0x01)
        iblock = bytes([pcb]) + apdu
        await self._medium.transmit_reader(RFFrame.from_bytes(iblock))
        resp = await self._medium.receive_reader()
        if resp is None:
            raise RuntimeError("No response to I-block")
        return resp.data[1:]