pm3py/pm3py/sim/auth_aes.py

"""NXP AES-128 key management mixin.

Reusable mixin for ICODE DNA and NTAG 5 Link/Boost AES auth:
- 4 key slots with key headers and privilege bitmasks
- Global Crypto Header (GCH) for access enforcement
- Authentication limit with failure counting
- TAM1 (Tag Authentication Method 1) — one-way tag authentication
- MAM1/MAM2 (Mutual Authentication Method) — two-phase mutual auth
- Helper methods for key state and privilege checks
"""
from __future__ import annotations

import os

try:
    from Crypto.Cipher import AES as _AES
    _HAS_AES = True
except ImportError:
    _HAS_AES = False


# Key header values
KEY_HEADER_NOT_ACTIVE = 0x81
KEY_HEADER_ACTIVE_LOCKED = 0xE7
KEY_HEADER_DISABLED = 0xFF

# Key privilege bits (Table 5)
PRIV_READ = 0x01
PRIV_WRITE = 0x02
PRIV_PRIVACY = 0x04
PRIV_DESTROY = 0x08
PRIV_EAS_AFI = 0x10
PRIV_CRYPTO_CONFIG = 0x20
PRIV_AREA1_READ = 0x40
PRIV_AREA1_WRITE = 0x80

# NFC Global Crypto Header values
GCH_DEACTIVATED = 0x81
GCH_DEACTIVATED_PRIV_LOCKED = 0x87
GCH_ACTIVATED = 0xC1
GCH_ACTIVATED_PRIV_LOCKED = 0xC7
GCH_FINAL = 0xE7


class NxpAesAuth:
    """Mixin providing AES-128 key management for NXP ICODE/NTAG 5 tags.

    Call ``_init_aes_auth(aes_keys)`` from the concrete class's
    ``__init__`` to set up AES key state.
    """

    def _init_aes_auth(self, aes_keys: list[bytes | None] | None = None) -> None:
        """Initialize AES key management state.

        Args:
            aes_keys: optional list of up to 4 AES-128 keys (or None per slot).
        """
        self._aes_keys: list[bytes | None] = list(aes_keys) if aes_keys else [None] * 4
        self._aes_authenticated: set[int] = set()

        # Key headers: controls whether key is active for authentication
        self._key_headers: list[int] = [KEY_HEADER_NOT_ACTIVE] * 4

        # Key privileges: bitmask per key defining what auth grants
        self._key_privileges: list[int] = [0x00] * 4

        # Global Crypto Header — controls enforcement of access conditions
        self._nfc_gch: int = GCH_DEACTIVATED

        # Authentication limit (0 = unlimited)
        self._auth_limit: int = 0
        self._auth_fail_count: int = 0

    def is_key_active(self, key_id: int) -> bool:
        """Check if AES key slot is active (header == 0xE7)."""
        return self._key_headers[key_id] == KEY_HEADER_ACTIVE_LOCKED

    def has_privilege(self, key_id: int, privilege: int) -> bool:
        """Check if key_id is authenticated, active, and has the given privilege bit."""
        if key_id not in self._aes_authenticated:
            return False
        if not self.is_key_active(key_id):
            return False
        return bool(self._key_privileges[key_id] & privilege)

    def is_access_enforced(self) -> bool:
        """Check if GCH is activated (>= 0xC1), meaning access conditions are enforced."""
        return self._nfc_gch >= GCH_ACTIVATED

    # ----- TAM1 (Tag Authentication Method 1) -----

    # Crypto constants — first 2 bytes of plaintext block
    _C_TAM1 = b"\x96\xC5"
    _C_MAM1 = b"\xDA\x83"
    _C_MAM2_PURPOSE = b"\xDA\x80"

    def _handle_challenge(self, frame) -> None:
        """Handle CHALLENGE (0x39) — compute and store TResponse.

        Format: flags(1) cmd(1) CSI(1) AuthMethod(1) KeyID(1) IChallenge(10)
        Tag computes: AES-ECB-ENC(key, C_TAM1 || TRnd || IChallenge)
        Stores TResponse reversed (wire byte order).
        Returns None — CHALLENGE has no RF response.
        """
        self._tam1_response = None

        if not _HAS_AES:
            return

        data = frame.data
        if len(data) < 15:  # flags + cmd + CSI + AuthMethod + KeyID + 10 IChallenge
            return

        auth_method = data[3]
        if auth_method != 0x00:  # TAM1 only
            return

        key_id = data[4]
        if key_id > 3:
            return

        if not self.is_key_active(key_id):
            return

        key = self._aes_keys[key_id]
        if key is None:
            return

        ichallenge = data[5:15]
        if len(ichallenge) != 10:
            return

        # Build plaintext: C_TAM1(2) + TRnd(4) + IChallenge(10) = 16 bytes
        trnd = os.urandom(4)
        plaintext = self._C_TAM1 + trnd + ichallenge

        cipher = _AES.new(key, _AES.MODE_ECB)
        tresponse = cipher.encrypt(plaintext)

        # Store reversed for wire format
        self._tam1_response = tresponse[::-1]

    def _handle_readbuffer(self, frame):
        """Handle READBUFFER (0x3A) — return stored TResponse.

        Format: flags(1) cmd(1)
        Returns: flags(00) + TResponse_reversed(16 bytes)
        Returns None if no TResponse is stored.
        """
        from .frame import RFFrame

        resp = getattr(self, "_tam1_response", None)
        if resp is None:
            return None

        result = RFFrame.from_bytes(bytes([0x00]) + resp)
        # Clear after read (one-shot)
        self._tam1_response = None
        return result

    # ----- MAM (Mutual Authentication Method) -----

    def _handle_authenticate(self, frame):
        """Handle AUTHENTICATE (0x35) — dispatch MAM1 or MAM2.

        Format: flags(1) cmd(1) CSI(1) AuthMethod(1) ...
        AuthMethod 0x02 = MAM1, 0x06 = MAM2.
        Returns RFFrame response or error frame.
        """
        from .frame import RFFrame

        data = frame.data
        if len(data) < 4:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        auth_method = data[3]
        if auth_method == 0x02:
            return self._handle_mam1(frame)
        elif auth_method == 0x06:
            return self._handle_mam2(frame)
        else:
            # Unsupported AuthMethod (0x80, 0x90, etc.)
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

    def _handle_mam1(self, frame):
        """Handle MAM1 (AuthMethod=0x02) — tag proves identity to reader.

        Format: flags(1) cmd(1) CSI(1) AuthMethod(1) KeyID(1) IChallenge_reversed(10)
        Total: 15 bytes minimum.

        Tag computes: AES-ECB-ENC(key, C_MAM1 || TChallenge[31:0] || IChallenge)
        Response: flags(04) header(A7) TChallenge_high_reversed(6) encrypted_reversed(16)
        """
        from .frame import RFFrame

        self._mam_state = None

        if not _HAS_AES:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        data = frame.data
        if len(data) < 15:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        key_id = data[4]
        if key_id > 3:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        if not self.is_key_active(key_id):
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        key = self._aes_keys[key_id]
        if key is None:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        # De-reverse IChallenge from wire format
        ichallenge_reversed = data[5:15]
        ichallenge = ichallenge_reversed[::-1]

        # Generate 10-byte TChallenge
        tchallenge = os.urandom(10)

        # Build plaintext: C_MAM1(2) + TChallenge[31:0](4) + IChallenge(10) = 16
        # TChallenge[31:0] = last 4 bytes of TChallenge (bytes 6-9)
        tc_low = tchallenge[6:10]
        plaintext = self._C_MAM1 + tc_low + ichallenge

        cipher = _AES.new(key, _AES.MODE_ECB)
        encrypted = cipher.encrypt(plaintext)

        # TChallenge_high = first 6 bytes (bytes 0-5), reversed for wire
        tc_high = tchallenge[0:6]
        tc_high_reversed = tc_high[::-1]

        # Encrypted block reversed for wire
        encrypted_reversed = encrypted[::-1]

        # Store state for MAM2 verification
        self._mam_state = {
            "key_id": key_id,
            "ichallenge": ichallenge,
            "tchallenge": tchallenge,
        }

        # Response: flags(04) header(A7) tc_high_reversed(6) encrypted_reversed(16)
        resp = bytes([0x04, 0xA7]) + tc_high_reversed + encrypted_reversed
        return RFFrame.from_bytes(resp)

    def _handle_mam2(self, frame):
        """Handle MAM2 (AuthMethod=0x06) — reader proves identity to tag.

        Format: flags(1) cmd(1) CSI(1) AuthMethod(1) IResponse_reversed(16)
        Total: 20 bytes minimum.

        Tag verifies: AES-ECB-ENC(key, iresponse) == C_MAM2_PURPOSE || IChallenge[31:0] || TChallenge
        On success: key_id added to _aes_authenticated, response = flags(00).
        """
        from .frame import RFFrame

        if not _HAS_AES:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        # Must have prior MAM1 state
        mam_state = getattr(self, "_mam_state", None)
        if mam_state is None:
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        data = frame.data
        if len(data) < 20:
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        key_id = mam_state["key_id"]
        ichallenge = mam_state["ichallenge"]
        tchallenge = mam_state["tchallenge"]

        key = self._aes_keys[key_id]
        if key is None:
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        # De-reverse IResponse from wire format
        iresponse_reversed = data[4:20]
        iresponse = iresponse_reversed[::-1]

        # Encrypt IResponse to get plaintext
        cipher = _AES.new(key, _AES.MODE_ECB)
        plaintext = cipher.encrypt(iresponse)

        # Parse: C_MAM2[11:0]||Purpose[3:0] (2 bytes) + IChallenge[31:0](4) + TChallenge[79:0](10)
        purpose = plaintext[0:2]
        ich_31_0 = plaintext[2:6]
        tc_check = plaintext[6:16]

        # Verify upper 12 bits of C_MAM2 field (0xDA8x)
        if purpose[0] != 0xDA or (purpose[1] & 0xF0) != 0x80:
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        purpose_nibble = purpose[1] & 0x0F

        # IChallenge[31:0] = last 4 bytes of original IChallenge
        if ich_31_0 != ichallenge[6:10]:
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        if tc_check != tchallenge:
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        # Check privilege requirements for non-standard purposes
        if purpose_nibble in (0x08, 0x09, 0x0A):
            if not (self._key_privileges[key_id] & PRIV_PRIVACY):
                self._mam_state = None
                return RFFrame.from_bytes(bytes([0x01, 0x0F]))
        elif purpose_nibble == 0x0B:
            if not (self._key_privileges[key_id] & PRIV_DESTROY):
                self._mam_state = None
                return RFFrame.from_bytes(bytes([0x01, 0x0F]))
        elif purpose_nibble != 0x00:
            # Unknown purpose nibble
            self._mam_state = None
            return RFFrame.from_bytes(bytes([0x01, 0x0F]))

        # Authentication successful
        self._aes_authenticated.add(key_id)

        # Dispatch purpose action
        if purpose_nibble == 0x08:
            # Temporary disable NFC Privacy Mode (until field reset)
            if hasattr(self, "_privacy_mode"):
                self._privacy_mode = False
        elif purpose_nibble == 0x09:
            # Enable NFC Privacy Mode
            if hasattr(self, "_privacy_mode"):
                self._privacy_mode = True
        elif purpose_nibble == 0x0A:
            # Permanent disable NFC Privacy Mode
            if hasattr(self, "_privacy_mode"):
                self._privacy_mode = False
        elif purpose_nibble == 0x0B:
            # Destroy
            if hasattr(self, "_destroyed"):
                self._destroyed = True

        self._mam_state = None

        return RFFrame.from_bytes(bytes([0x00]))