"""Tests for AES TAM1/MAM authentication in NxpAesAuth mixin."""
import asyncio
import os

import pytest
from Crypto.Cipher import AES

from pm3py.sim.frame import RFFrame


def run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)


class TestTam1:
    def test_challenge_readbuffer(self):
        """Full TAM1 flow: CHALLENGE computes TResponse, READBUFFER returns it."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        key = b"\xAA" * 16
        tag = IcodeDnaTag(aes_keys=[key, None, None, None])
        tag._key_headers[0] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        ichallenge = bytes(range(10))

        # CHALLENGE: flags(02) cmd(39) CSI(00) AuthMethod(00) KeyID(00) IChallenge(10)
        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x00]) + ichallenge
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None  # CHALLENGE has no RF response

        # READBUFFER: flags(02) cmd(3A)
        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert len(resp.data) >= 17  # flags + 16 bytes TResponse

        tresponse_wire = resp.data[1:17]
        tresponse = tresponse_wire[::-1]  # reverse for decryption

        cipher = AES.new(key, AES.MODE_ECB)
        decrypted = cipher.decrypt(tresponse)

        # Verify C_TAM1 constant
        assert decrypted[0:2] == bytes([0x96, 0xC5])
        # Verify echoed IChallenge (bytes 6-15, reversed back)
        echoed = decrypted[6:16]
        assert echoed == ichallenge

    def test_challenge_key1(self):
        """TAM1 with key slot 1."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        key = b"\xCC" * 16
        tag = IcodeDnaTag(aes_keys=[None, key, None, None])
        tag._key_headers[1] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        ichallenge = b"\xFF" * 10
        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x01]) + ichallenge
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None

        tresponse = resp.data[1:17][::-1]
        cipher = AES.new(key, AES.MODE_ECB)
        decrypted = cipher.decrypt(tresponse)
        assert decrypted[0:2] == bytes([0x96, 0xC5])
        assert decrypted[6:16] == ichallenge

    def test_challenge_inactive_key_silent(self):
        """CHALLENGE with inactive key: tag stays silent, READBUFFER returns nothing."""
        from pm3py.sim.icode_dna import IcodeDnaTag

        tag = IcodeDnaTag(aes_keys=[b"\xBB" * 16, None, None, None])
        # Key header NOT active (default 0x81)
        run(tag.power_on())

        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x00]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

        # READBUFFER should have nothing
        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

    def test_readbuffer_without_challenge(self):
        """READBUFFER without prior CHALLENGE returns nothing."""
        from pm3py.sim.icode_dna import IcodeDnaTag

        tag = IcodeDnaTag()
        run(tag.power_on())

        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

    def test_readbuffer_one_shot(self):
        """TResponse is cleared after first READBUFFER (one-shot)."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        key = b"\xDD" * 16
        tag = IcodeDnaTag(aes_keys=[key, None, None, None])
        tag._key_headers[0] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        # CHALLENGE
        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x00]) + bytes(10)
        run(tag.handle_frame(RFFrame.from_bytes(cmd)))

        # First READBUFFER — should succeed
        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None

        # Second READBUFFER — should return nothing (consumed)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

    def test_challenge_no_key_in_slot(self):
        """CHALLENGE with None key slot stays silent."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        tag = IcodeDnaTag(aes_keys=[None, None, None, None])
        tag._key_headers[0] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x00]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

        cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

    def test_challenge_invalid_key_id(self):
        """CHALLENGE with key_id > 3 stays silent."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        tag = IcodeDnaTag(aes_keys=[b"\xAA" * 16, None, None, None])
        tag._key_headers[0] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        cmd = bytes([0x02, 0x39, 0x00, 0x00, 0x04]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

    def test_challenge_wrong_auth_method(self):
        """CHALLENGE with AuthMethod != 0x00 (not TAM1) stays silent."""
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        tag = IcodeDnaTag(aes_keys=[b"\xAA" * 16, None, None, None])
        tag._key_headers[0] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())

        # AuthMethod = 0x02 (MAM1, not TAM1)
        cmd = bytes([0x02, 0x39, 0x00, 0x02, 0x00]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is None

        rb_cmd = bytes([0x02, 0x3A])
        resp = run(tag.handle_frame(RFFrame.from_bytes(rb_cmd)))
        assert resp is None


class TestMam:
    """Tests for MAM1/MAM2 mutual authentication via AUTHENTICATE (0x35)."""

    def _make_tag(self, key=b"\xAA" * 16, key_slot=0):
        from pm3py.sim.icode_dna import IcodeDnaTag
        from pm3py.sim.auth_aes import KEY_HEADER_ACTIVE_LOCKED

        keys = [None] * 4
        keys[key_slot] = key
        tag = IcodeDnaTag(aes_keys=keys)
        tag._key_headers[key_slot] = KEY_HEADER_ACTIVE_LOCKED
        run(tag.power_on())
        return tag

    def _do_mam1(self, tag, key, ichallenge, key_id=0):
        """Send MAM1 and return (resp, tchallenge)."""
        cmd = bytes([0x02, 0x35, 0x00, 0x02, key_id]) + ichallenge[::-1]
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert len(resp.data) == 24
        assert resp.data[0] == 0x04
        assert resp.data[1] == 0xA7

        tc_high_reversed = resp.data[2:8]
        encrypted_reversed = resp.data[8:24]

        cipher = AES.new(key, AES.MODE_ECB)
        decrypted = cipher.decrypt(encrypted_reversed[::-1])

        assert decrypted[0:2] == bytes([0xDA, 0x83])
        tc_low = decrypted[2:6]
        echoed_ichallenge = decrypted[6:16]
        assert echoed_ichallenge == ichallenge

        tc_high = tc_high_reversed[::-1]
        tchallenge = tc_high + tc_low
        return resp, tchallenge

    def _compute_iresponse(self, key, ichallenge, tchallenge):
        """Compute IResponse for MAM2."""
        C_MAM2_PURPOSE = bytes([0xDA, 0x80])
        ich_31_0 = ichallenge[6:10]
        plaintext = C_MAM2_PURPOSE + ich_31_0 + tchallenge
        cipher = AES.new(key, AES.MODE_ECB)
        iresponse = cipher.decrypt(plaintext)
        return iresponse[::-1]

    def test_mam_complete_flow(self):
        """Complete MAM1 -> MAM2 flow results in authenticated key."""
        key = b"\xAA" * 16
        tag = self._make_tag(key)
        ichallenge = os.urandom(10)

        resp, tchallenge = self._do_mam1(tag, key, ichallenge)

        iresponse_wire = self._compute_iresponse(key, ichallenge, tchallenge)

        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + iresponse_wire
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert resp.data[0] == 0x00
        assert 0 in tag._aes_authenticated

    def test_mam1_response_structure(self):
        """MAM1 response has correct flags, header, and 24-byte length."""
        key = b"\xBB" * 16
        tag = self._make_tag(key)
        ichallenge = bytes(range(10))

        cmd = bytes([0x02, 0x35, 0x00, 0x02, 0x00]) + ichallenge[::-1]
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))

        assert resp is not None
        assert len(resp.data) == 24
        assert resp.data[0] == 0x04
        assert resp.data[1] == 0xA7

    def test_mam2_wrong_iresponse(self):
        """MAM2 with incorrect IResponse returns error."""
        key = b"\xCC" * 16
        tag = self._make_tag(key)
        ichallenge = os.urandom(10)

        self._do_mam1(tag, key, ichallenge)

        bad_iresponse = os.urandom(16)
        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + bad_iresponse
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert resp.data[0] == 0x01
        assert 0 not in tag._aes_authenticated

    def test_mam2_without_prior_mam1(self):
        """MAM2 without prior MAM1 returns error."""
        key = b"\xDD" * 16
        tag = self._make_tag(key)

        iresponse = os.urandom(16)
        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + iresponse
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert resp.data[0] == 0x01
        assert 0 not in tag._aes_authenticated

    def test_mam_key1(self):
        """MAM flow with key slot 1."""
        key = b"\xEE" * 16
        tag = self._make_tag(key, key_slot=1)
        ichallenge = os.urandom(10)

        resp, tchallenge = self._do_mam1(tag, key, ichallenge, key_id=1)

        iresponse_wire = self._compute_iresponse(key, ichallenge, tchallenge)
        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + iresponse_wire
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp.data[0] == 0x00
        assert 1 in tag._aes_authenticated

    def test_mam1_inactive_key_error(self):
        """MAM1 with inactive key returns error."""
        from pm3py.sim.icode_dna import IcodeDnaTag

        tag = IcodeDnaTag(aes_keys=[b"\xAA" * 16, None, None, None])
        run(tag.power_on())

        cmd = bytes([0x02, 0x35, 0x00, 0x02, 0x00]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp is not None
        assert resp.data[0] == 0x01

    def test_mam1_invalid_key_id(self):
        """MAM1 with key_id > 3 returns error."""
        tag = self._make_tag()

        cmd = bytes([0x02, 0x35, 0x00, 0x02, 0x05]) + bytes(10)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp.data[0] == 0x01

    def test_unsupported_auth_method(self):
        """AUTHENTICATE with AuthMethod 0x80 returns error (not supported)."""
        tag = self._make_tag()

        cmd = bytes([0x02, 0x35, 0x00, 0x80]) + bytes(16)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp.data[0] == 0x01

    def test_mam_state_cleared_after_mam2_failure(self):
        """After MAM2 failure, state is cleared -- second MAM2 also fails."""
        key = b"\xAA" * 16
        tag = self._make_tag(key)
        ichallenge = os.urandom(10)

        self._do_mam1(tag, key, ichallenge)

        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + os.urandom(16)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp.data[0] == 0x01

        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + os.urandom(16)
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))
        assert resp.data[0] == 0x01

    def test_mam1_ichallenge_byte_reversal(self):
        """Verify IChallenge is de-reversed correctly from wire format."""
        key = b"\xAA" * 16
        tag = self._make_tag(key)

        ichallenge = bytes([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A])

        cmd = bytes([0x02, 0x35, 0x00, 0x02, 0x00]) + ichallenge[::-1]
        resp = run(tag.handle_frame(RFFrame.from_bytes(cmd)))

        encrypted_reversed = resp.data[8:24]
        cipher = AES.new(key, AES.MODE_ECB)
        decrypted = cipher.decrypt(encrypted_reversed[::-1])
        assert decrypted[6:16] == ichallenge

    def _do_mam_full(self, tag, key, ichallenge, key_id=0, purpose_nibble=0x00):
        """Complete MAM1+MAM2 flow with a specific purpose nibble. Returns MAM2 response."""
        _, tchallenge = self._do_mam1(tag, key, ichallenge, key_id=key_id)

        purpose_byte = 0x80 | purpose_nibble
        ich_31_0 = ichallenge[6:10]
        plaintext = bytes([0xDA, purpose_byte]) + ich_31_0 + tchallenge
        cipher = AES.new(key, AES.MODE_ECB)
        iresponse = cipher.decrypt(plaintext)
        iresponse_reversed = iresponse[::-1]

        cmd = bytes([0x02, 0x35, 0x00, 0x06]) + iresponse_reversed
        return run(tag.handle_frame(RFFrame.from_bytes(cmd)))

    def test_mam2_purpose_enable_privacy(self):
        """MAM2 with purpose=0x09 sets privacy mode."""
        from pm3py.sim.auth_aes import PRIV_PRIVACY

        key = b"\xAA" * 16
        tag = self._make_tag(key)
        tag._key_privileges[0] = PRIV_PRIVACY
        ichallenge = os.urandom(10)

        resp = self._do_mam_full(tag, key, ichallenge, purpose_nibble=0x09)
        assert resp.data[0] == 0x00
        assert 0 in tag._aes_authenticated
        assert tag._privacy_mode is True

    def test_mam2_purpose_disable_privacy(self):
        """MAM2 with purpose=0x0A clears privacy mode."""
        from pm3py.sim.auth_aes import PRIV_PRIVACY

        key = b"\xAA" * 16
        tag = self._make_tag(key)
        tag._key_privileges[0] = PRIV_PRIVACY
        tag._privacy_mode = True
        ichallenge = os.urandom(10)

        resp = self._do_mam_full(tag, key, ichallenge, purpose_nibble=0x0A)
        assert resp.data[0] == 0x00
        assert 0 in tag._aes_authenticated
        assert tag._privacy_mode is False

    def test_mam2_purpose_destroy(self):
        """MAM2 with purpose=0x0B sets destroyed."""
        from pm3py.sim.auth_aes import PRIV_DESTROY

        key = b"\xAA" * 16
        tag = self._make_tag(key)
        tag._key_privileges[0] = PRIV_DESTROY
        ichallenge = os.urandom(10)

        resp = self._do_mam_full(tag, key, ichallenge, purpose_nibble=0x0B)
        assert resp.data[0] == 0x00
        assert 0 in tag._aes_authenticated
        assert tag._destroyed is True

    def test_mam2_purpose_requires_privilege(self):
        """MAM2 with privacy purpose but no PRIV_PRIVACY returns error."""
        key = b"\xAA" * 16
        tag = self._make_tag(key)
        # No privileges set (default 0x00)
        ichallenge = os.urandom(10)

        resp = self._do_mam_full(tag, key, ichallenge, purpose_nibble=0x09)
        assert resp.data[0] == 0x01
        assert 0 not in tag._aes_authenticated
        assert tag._privacy_mode is False