pm3py/pm3py/transponders/hf/iso14443a/nxp/mifare_classic.py

"""MIFARE Classic transponder and reader models with Crypto-1 authentication."""
from __future__ import annotations

import struct

from pm3py.sim.frame import RFFrame
from pm3py.sim.medium import Medium
from pm3py.transponders.bitfield import BitField, Register
from ..base import Tag14443A_3, Reader14443A, _compute_bcc
from .crypto1 import Crypto1

# MIFARE Classic commands
AUTH_A = 0x60
AUTH_B = 0x61
READ_CMD = 0x30
WRITE_CMD = 0xA0
ACK = 0x0A
NACK = 0x00

BLOCK_SIZE = 16

# Per-block access-condition (C1 C2 C3) meanings. Data-block table matches Proxmark3
# mifare4.c; the trailer table is written to the NXP MF1S50 datasheet (r=read, w=write,
# A/B = which key, AB = either key, - = never).
_MFAC_DATA = {
    0: "r=AB w=AB incr=AB dec/tr/rst=AB  (transport)",
    1: "r=AB dec/tr/rst=AB  (value block)",
    2: "r=AB  (read-only)",
    3: "r=B w=B",
    4: "r=AB w=B",
    5: "r=B  (read-only, key B)",
    6: "r=AB w=B incr=B dec/tr/rst=AB  (value block)",
    7: "no access",
}
_MFAC_TRAILER = {
    0: "KeyA w=A; AC r=A; KeyB r=A w=A  (KeyB is readable data)",
    1: "KeyA w=A; AC r=A w=A; KeyB r=A w=A  (transport/default)",
    2: "KeyA w=-; AC r=A; KeyB r=A w=-",
    3: "KeyA w=B; AC r=AB w=B; KeyB w=B",
    4: "KeyA w=B; AC r=AB; KeyB w=B",
    5: "AC r=AB w=B  (keys locked)",
    6: "KeyA w=-; AC r=AB  (keys locked)",
    7: "AC r=AB  (all keys locked)",
}


class MifareAccessConditions(Register):
    """MIFARE Classic sector-trailer access bits — trailer bytes 6-9 (the three access bytes
    plus the general-purpose byte), decoded as a 32-bit word.

    Each of the sector's four blocks (three data blocks + the trailer) has a 3-bit access
    condition ``C1 C2 C3``. The three access bytes store every condition *twice* — once true,
    once inverted — so one flipped bit permanently bricks the sector. This spells out each
    block in plain English, flags an inconsistent inverted copy (``.valid``), and, when you
    :meth:`build` one, computes the inverted copy for you so it is always self-consistent::

        MifareAccessConditions(0xFF078069)             # decode the factory-default trailer
        ac = MifareAccessConditions.build(trailer=3)   # KeyB-controlled trailer
        bytes(ac)                                       # -> the 4 trailer bytes 6..9 to write
    """

    _width_bits = 32
    _label = "MIFARE access bits (trailer 6-9)"

    gpb = BitField(7, 0, doc="general-purpose byte — free user data (e.g. the MAD GPB)")

    _BLOCKS = ("block0", "block1", "block2", "trailer")

    def _cond(self, blockn: int) -> int:
        """The 3-bit access condition (C1<<2 | C2<<1 | C3) for one block of the sector."""
        v = self.value
        c1 = (v >> (20 + blockn)) & 1   # C1 = byte7 high nibble
        c2 = (v >> (8 + blockn)) & 1    # C2 = byte8 low nibble
        c3 = (v >> (12 + blockn)) & 1   # C3 = byte8 high nibble
        return (c1 << 2) | (c2 << 1) | c3

    @property
    def block0(self) -> int:
        return self._cond(0)

    @property
    def block1(self) -> int:
        return self._cond(1)

    @property
    def block2(self) -> int:
        return self._cond(2)

    @property
    def trailer(self) -> int:
        return self._cond(3)

    def describe(self, blockn: int) -> str:
        table = _MFAC_TRAILER if blockn == 3 else _MFAC_DATA
        return table.get(self._cond(blockn), "?")

    @property
    def valid(self) -> bool:
        """True iff the redundant inverted copy is the nibble-complement of the true copy."""
        v = self.value
        return (((v >> 24) & 0x0F) == (((v >> 20) & 0x0F) ^ 0x0F) and   # ~C1 vs C1
                ((v >> 28) & 0x0F) == (((v >> 8) & 0x0F) ^ 0x0F) and    # ~C2 vs C2
                ((v >> 16) & 0x0F) == (((v >> 12) & 0x0F) ^ 0x0F))      # ~C3 vs C3

    @classmethod
    def build(cls, block0: int = 0, block1: int = 0, block2: int = 0,
              trailer: int = 1, gpb: int = 0x00) -> "MifareAccessConditions":
        """Compose access bytes from the four per-block conditions (0..7), filling in the
        inverted copy so the result always passes ``.valid``. ``trailer`` defaults to 1
        (the factory/transport trailer). ``gpb`` is the general-purpose byte."""
        c1 = c2 = c3 = 0
        for n, cond in enumerate((block0, block1, block2, trailer)):
            if not 0 <= cond <= 7:
                raise ValueError(f"block {n} access condition must be 0..7, got {cond}")
            c1 |= ((cond >> 2) & 1) << n
            c2 |= ((cond >> 1) & 1) << n
            c3 |= (cond & 1) << n
        byte6 = ((c2 ^ 0x0F) << 4) | (c1 ^ 0x0F)
        byte7 = (c1 << 4) | (c3 ^ 0x0F)
        byte8 = (c3 << 4) | c2
        return cls((byte6 << 24) | (byte7 << 16) | (byte8 << 8) | (gpb & 0xFF))

    def _rows(self) -> list[dict]:
        rows = [
            {"name": name, "value": f"{self._cond(n)}  {self.describe(n)}", "bits": "", "doc": ""}
            for n, name in enumerate(self._BLOCKS)
        ]
        rows.append({"name": "gpb", "value": f"0x{self.gpb:02X}", "bits": "[7:0]", "doc": ""})
        if not self.valid:
            rows.append({"name": "(!)", "value": "inverted copy inconsistent — INVALID",
                         "bits": "", "doc": ""})
        return rows


class MifareClassicTag(Tag14443A_3):
    """MIFARE Classic 1K/4K transponder with Crypto-1 authentication."""

    _SIZE_CONFIG = {
        "1k": {"atqa": b"\x04\x00", "sak": 0x08, "blocks": 64},
        "4k": {"atqa": b"\x02\x00", "sak": 0x18, "blocks": 256},
    }

    def __init__(self, uid: bytes, size: str = "1k",
                 keys_a: dict[int, bytes] | None = None,
                 keys_b: dict[int, bytes] | None = None,
                 data: bytearray | None = None):
        cfg = self._SIZE_CONFIG[size]
        super().__init__(uid=uid, atqa=cfg["atqa"], sak=cfg["sak"])
        self._size = size
        self._num_blocks = cfg["blocks"]
        self._num_sectors = 16 if size == "1k" else 40
        self._keys_a = keys_a or {s: b"\xFF" * 6 for s in range(self._num_sectors)}
        self._keys_b = keys_b or {s: b"\xFF" * 6 for s in range(self._num_sectors)}
        self._data = data if data is not None else bytearray(self._num_blocks * BLOCK_SIZE)
        # Region for block data (14a EML not yet supported, so eml_offset=-1)
        from pm3py.sim.memory import DirtyByteArray, MemoryRegion
        self.regions["blocks"] = MemoryRegion(
            name="blocks", data=DirtyByteArray(self._data),
            block_size=BLOCK_SIZE, eml_offset=-1)
        self._crypto: Crypto1 | None = None
        self._auth_sector: int | None = None
        self._auth_state: str = "NONE"  # NONE, NONCE_SENT, AUTHENTICATED
        self._tag_nonce: bytes = b""

    def sector_for_block(self, block: int) -> int:
        """Map block number to sector number."""
        if block < 128:
            return block // 4
        else:
            return 32 + (block - 128) // 16

    def _sector_trailer_block(self, sector: int) -> int:
        """Get the trailer block number for a sector."""
        if sector < 32:
            return sector * 4 + 3
        else:
            return 128 + (sector - 32) * 16 + 15

    def access(self, sector: int) -> MifareAccessConditions:
        """Decode the sector's trailer access bits (trailer bytes 6-9) as a self-describing
        :class:`MifareAccessConditions` register."""
        trailer = self.read_block_raw(self._sector_trailer_block(sector))
        return MifareAccessConditions(int.from_bytes(trailer[6:10], "big"))

    def read_block_raw(self, block: int) -> bytes:
        """Direct read of block data (bypassing auth)."""
        offset = block * BLOCK_SIZE
        return bytes(self._data[offset:offset + BLOCK_SIZE])

    def write_block_raw(self, block: int, data: bytes) -> None:
        """Direct write of block data (bypassing auth)."""
        if len(data) != BLOCK_SIZE:
            raise ValueError(f"Block data must be {BLOCK_SIZE} bytes")
        offset = block * BLOCK_SIZE
        self._data[offset:offset + BLOCK_SIZE] = data

    def _handle_application(self, frame: RFFrame) -> RFFrame | None:
        """Handle MIFARE Classic commands after SELECT."""
        if not frame.data:
            return None

        # Auth response comes as raw encrypted data (not a command)
        if self._auth_state == "NONCE_SENT":
            return self._handle_auth_response(frame)

        # Write phase 2: receive data after ACK (before command dispatch)
        if self._auth_state == "WRITE_PENDING":
            return self._handle_write_phase2(frame)

        cmd = frame.data[0]

        # AUTH commands work in any post-SELECT state
        if cmd in (AUTH_A, AUTH_B):
            return self._handle_auth_start(frame)

        if self._auth_state != "AUTHENTICATED":
            return None  # not authenticated

        # Authenticated commands (plaintext in simulation)
        if cmd == READ_CMD:
            return self._handle_read(frame)
        if cmd == WRITE_CMD:
            return self._handle_write_phase1(frame)

        return None

    def _handle_auth_start(self, frame: RFFrame) -> RFFrame | None:
        """Handle AUTH_A/AUTH_B command: return tag nonce."""
        if len(frame.data) < 2:
            return None

        cmd = frame.data[0]
        block = frame.data[1]
        sector = self.sector_for_block(block)
        key_type = "a" if cmd == AUTH_A else "b"
        key = self._keys_a[sector] if key_type == "a" else self._keys_b[sector]

        # Store the key for verification during auth response
        self._auth_key = key
        self._auth_sector = sector

        # Generate tag nonce
        self._crypto = Crypto1(key)
        self._tag_nonce = self._crypto.generate_nonce()
        self._auth_state = "NONCE_SENT"

        # Initialize cipher with uid ^ nt
        uid_int = struct.unpack(">I", self._uid[:4])[0]
        nt_int = struct.unpack(">I", self._tag_nonce)[0]

        # Re-init crypto with key and auth init
        self._crypto = Crypto1(key)
        self._crypto.init_auth(uid_int, nt_int)

        return RFFrame.from_bytes(self._tag_nonce)

    def _handle_auth_response(self, frame: RFFrame) -> RFFrame | None:
        """Handle reader's auth response (nr_enc + ar_enc).

        Verifies the reader used the correct key by re-deriving the expected
        encrypted response and comparing.
        """
        if len(frame.data) < 8 or self._crypto is None:
            self._auth_state = "NONE"
            return None

        nr_enc = frame.data[0:4]
        ar_enc = frame.data[4:8]

        # Verify: build a reader-side Crypto1 with the tag's key and same
        # init_auth parameters. If the reader used the correct key, its
        # encrypted output will match what we compute.
        uid_int = struct.unpack(">I", self._uid[:4])[0]
        nt_int = struct.unpack(">I", self._tag_nonce)[0]

        verify = Crypto1(self._auth_key)
        verify.init_auth(uid_int, nt_int)

        # The reader encrypted `nr` with its keystream. If keys match,
        # decrypting with our matching keystream gives the real `nr`.
        # Then re-encrypting `nr` with our verify keystream should match nr_enc.
        expected_ks = bytearray(8)
        for i in range(8):
            expected_ks[i] = verify.generate_byte()

        # If keys match: nr_enc = nr XOR ks. Both sides have same ks.
        # We can verify by checking that nr_enc XOR ks produces the same
        # nr that the reader intended, and then checking ar consistency.
        # Simpler: just verify the reader's 8 bytes match what our key produces.
        # The reader sends nr_enc = nr XOR reader_ks, ar_enc = ar XOR reader_ks.
        # If reader has same key, reader_ks == our_ks, so:
        #   nr_enc XOR our_ks[0:4] == nr (the reader's chosen nonce)
        #   ar_enc XOR our_ks[4:8] == ar (should be a specific value)
        # Since nr is random (reader-chosen), we can't verify it.
        # But we CAN verify the relationship: if keys differ, ks differs,
        # so the decrypted values will be inconsistent.

        # For simulation simplicity: accept if keys match, reject otherwise.
        # We detect key mismatch by checking that both sides produce
        # identical keystreams after init_auth with the same parameters.
        tag_ks = bytearray(8)
        # Reset tag's crypto to same state
        tag_verify = Crypto1(self._auth_key)
        tag_verify.init_auth(uid_int, nt_int)
        for i in range(8):
            tag_ks[i] = tag_verify.generate_byte()

        # Now check if nr_enc was produced by the same keystream
        # nr_enc = nr XOR ks_reader. If ks_reader == ks_tag, then
        # nr_enc XOR ks_tag gives us the real nr.
        # We don't know nr, but we can verify that the READER also used
        # this same ks by checking if their encrypted output is valid.
        # The reader's key produces a different ks if wrong.
        # Detect by: reader computes ks_r = Crypto1(wrong_key).init_auth(uid,nt)
        # nr_enc = nr XOR ks_r. ks_r != ks_tag. So nr_enc XOR ks_tag != nr.
        # This is fine, but we can't verify nr is "correct" since it's random.

        # ACTUAL SIMPLE APPROACH: The reader embeds its key hash in the protocol.
        # For simulation, we'll pass the reader's key through the frame metadata.
        # But that breaks the protocol abstraction.

        # PRAGMATIC APPROACH: Use deterministic nr and verify the full 8-byte
        # encrypted block matches what a correct-key reader would produce.
        reader_crypto_check = Crypto1(self._auth_key)
        reader_crypto_check.init_auth(uid_int, nt_int)
        # Generate the same nr the reader uses (we know it's deterministic: 0x01234567)
        # and compute what the correct output should be
        check_nr = b"\x01\x23\x45\x67"  # reader's known nr
        expected_nr_enc = reader_crypto_check.encrypt_bytes(check_nr)
        expected_ar_enc = reader_crypto_check.encrypt_bytes(b"\x00\x00\x00\x00")

        if nr_enc + ar_enc != expected_nr_enc + expected_ar_enc:
            self._auth_state = "NONE"
            self._crypto = None
            return None

        # Consume our crypto's keystream to stay in sync
        self._crypto.encrypt_bytes(nr_enc)
        self._crypto.encrypt_bytes(ar_enc)

        self._auth_state = "AUTHENTICATED"
        at = self._crypto.encrypt_bytes(b"\x00\x00\x00\x00")
        return RFFrame.from_bytes(at)

    def _handle_read(self, frame: RFFrame) -> RFFrame | None:
        """Handle READ command — return 16 bytes of block data."""
        if len(frame.data) < 2:
            return None
        block = frame.data[1]
        if block >= self._num_blocks:
            return None

        # Check block is in authenticated sector
        if self.sector_for_block(block) != self._auth_sector:
            return None

        data = self.read_block_raw(block)
        return RFFrame.from_bytes(data)

    def _handle_write_phase1(self, frame: RFFrame) -> RFFrame | None:
        """Handle WRITE command phase 1 — ACK, wait for data."""
        if len(frame.data) < 2:
            return None
        block = frame.data[1]
        if block >= self._num_blocks:
            return None
        if self.sector_for_block(block) != self._auth_sector:
            return None

        self._write_target_block = block
        self._auth_state = "WRITE_PENDING"
        return RFFrame.from_bytes(bytes([ACK]))

    def _handle_write_phase2(self, frame: RFFrame) -> RFFrame | None:
        """Handle WRITE data (16 bytes after ACK)."""
        if len(frame.data) < BLOCK_SIZE:
            self._auth_state = "AUTHENTICATED"
            return None
        self.write_block_raw(self._write_target_block, frame.data[:BLOCK_SIZE])
        self._auth_state = "AUTHENTICATED"
        return RFFrame.from_bytes(bytes([ACK]))


class MifareClassicReader:
    """MIFARE Classic reader with Crypto-1 authentication."""

    def __init__(self, medium: Medium):
        self._medium = medium
        self._reader = Reader14443A(medium)

    async def read_block(self, uid: bytes, block: int,
                         key: bytes, key_type: str = "a") -> dict:
        """Authenticate and read a block."""
        try:
            await self._reader.select_tag(uid)
        except RuntimeError:
            return {"success": False, "data": None}

        # AUTH
        auth_cmd = AUTH_A if key_type == "a" else AUTH_B
        auth_frame = RFFrame.from_bytes(bytes([auth_cmd, block]))
        await self._medium.transmit_reader(auth_frame)
        nt_resp = await self._medium.receive_reader()
        if nt_resp is None or len(nt_resp.data) < 4:
            return {"success": False, "data": None}

        nt = nt_resp.data[:4]

        # Initialize crypto
        crypto = Crypto1(key)
        uid_int = struct.unpack(">I", uid[:4])[0]
        nt_int = struct.unpack(">I", nt)[0]
        crypto.init_auth(uid_int, nt_int)

        # Send nr_enc + ar_enc
        nr = b"\x01\x23\x45\x67"
        nr_enc = crypto.encrypt_bytes(nr)
        ar_enc = crypto.encrypt_bytes(b"\x00\x00\x00\x00")
        auth_resp_frame = RFFrame.from_bytes(nr_enc + ar_enc)
        await self._medium.transmit_reader(auth_resp_frame)
        at_resp = await self._medium.receive_reader()
        if at_resp is None:
            return {"success": False, "data": None}

        # Authenticated — send READ
        read_frame = RFFrame.from_bytes(bytes([READ_CMD, block]))
        await self._medium.transmit_reader(read_frame)
        data_resp = await self._medium.receive_reader()
        if data_resp is None:
            return {"success": False, "data": None}

        return {"success": True, "data": data_resp.data[:16]}

    async def write_block(self, uid: bytes, block: int, data: bytes,
                          key: bytes, key_type: str = "a") -> dict:
        """Authenticate and write a block."""
        try:
            await self._reader.select_tag(uid)
        except RuntimeError:
            return {"success": False}

        # AUTH
        auth_cmd = AUTH_A if key_type == "a" else AUTH_B
        auth_frame = RFFrame.from_bytes(bytes([auth_cmd, block]))
        await self._medium.transmit_reader(auth_frame)
        nt_resp = await self._medium.receive_reader()
        if nt_resp is None or len(nt_resp.data) < 4:
            return {"success": False}

        nt = nt_resp.data[:4]

        # Initialize crypto
        crypto = Crypto1(key)
        uid_int = struct.unpack(">I", uid[:4])[0]
        nt_int = struct.unpack(">I", nt)[0]
        crypto.init_auth(uid_int, nt_int)

        # Send nr_enc + ar_enc
        nr = b"\x01\x23\x45\x67"
        nr_enc = crypto.encrypt_bytes(nr)
        ar_enc = crypto.encrypt_bytes(b"\x00\x00\x00\x00")
        await self._medium.transmit_reader(RFFrame.from_bytes(nr_enc + ar_enc))
        at_resp = await self._medium.receive_reader()
        if at_resp is None:
            return {"success": False}

        # WRITE phase 1
        write_frame = RFFrame.from_bytes(bytes([WRITE_CMD, block]))
        await self._medium.transmit_reader(write_frame)
        ack_resp = await self._medium.receive_reader()
        if ack_resp is None or ack_resp.data[0] != ACK:
            return {"success": False}

        # WRITE phase 2: send data (plaintext in simulation)
        await self._medium.transmit_reader(RFFrame.from_bytes(data))
        await self._medium.receive_reader()  # consume any response
        return {"success": True}