"""
Tests for lockstock_core.keyring — Envelope encryption key management.

Tests verify:
1. Account Key generation (32 random bytes)
2. KEK derivation determinism and differentiation
3. DEK wrap/unwrap round-trip and failure cases
4. Blob key derivation with pinned test vectors (CANARY — if these change,
   all previously encrypted data becomes unrecoverable)
5. Blob nonce derivation safety
6. Full lifecycle: account_key → KEK → DEK wrap → unwrap → blob_key
"""

import hashlib
import os
import pytest

from lockstock_core.keyring import (
    generate_account_key,
    account_key_verification_hash,
    derive_kek,
    generate_dek,
    wrap_dek,
    unwrap_dek,
    derive_blob_key,
    derive_blob_nonce,
)


# --- Fixed test inputs for pinned vectors ---

# These are NOT real keys — they are deterministic test values.
FIXED_ACCOUNT_KEY = bytes.fromhex(
    "a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2"
)
FIXED_AGENT_ID = "agent_test_keyring_001"
FIXED_DEK = bytes.fromhex(
    "1122334455667788990011223344556677889900112233445566778899001122"
)
FIXED_STAMP_HASH = "52d9e6bb163ecd57d2ec3feb440901dda37d8a564f9f3881988d869927c68e26"
FIXED_PAYLOAD_HASH = "fade8d4214d42417d4309c878cda679e7763336230f44f25a03322f1fa959c06"

# Different values for differentiation tests
ALT_AGENT_ID = "agent_test_keyring_002"
ALT_STAMP_HASH = "b3fc9283fd83539c35effb972eeb9d5990cc0f758480dd8dfa6e5e8385725ef2"
ALT_PAYLOAD_HASH = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
ALT_DEK = bytes.fromhex(
    "ffeeddccbbaa99887766554433221100ffeeddccbbaa99887766554433221100"
)


class TestAccountKey:
    """Test Account Key generation."""

    def test_generate_length(self):
        """Account key is 32 bytes."""
        key = generate_account_key()
        assert len(key) == 32

    def test_generate_random(self):
        """Two generated keys are different."""
        key1 = generate_account_key()
        key2 = generate_account_key()
        assert key1 != key2

    def test_verification_hash_deterministic(self):
        """Same key always produces same verification hash."""
        h1 = account_key_verification_hash(FIXED_ACCOUNT_KEY)
        h2 = account_key_verification_hash(FIXED_ACCOUNT_KEY)
        assert h1 == h2
        assert len(h1) == 64  # SHA-256 hex

    def test_verification_hash_different_keys(self):
        """Different keys produce different verification hashes."""
        h1 = account_key_verification_hash(FIXED_ACCOUNT_KEY)
        h2 = account_key_verification_hash(ALT_DEK)  # Using alt DEK as different key
        assert h1 != h2


class TestKEKDerivation:
    """Test Key Encryption Key derivation."""

    def test_deterministic(self):
        """Same (account_key, agent_id) always produces same KEK."""
        kek1 = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        kek2 = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        assert kek1 == kek2
        assert len(kek1) == 32  # AES-256

    def test_different_agent_different_kek(self):
        """Different agent_id produces different KEK (agent isolation)."""
        kek1 = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        kek2 = derive_kek(FIXED_ACCOUNT_KEY, ALT_AGENT_ID)
        assert kek1 != kek2

    def test_different_account_key_different_kek(self):
        """Different account_key produces different KEK."""
        kek1 = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        kek2 = derive_kek(ALT_DEK, FIXED_AGENT_ID)  # Different key, same agent
        assert kek1 != kek2

    def test_pinned_vector(self):
        """
        CANARY TEST: Pinned KEK derivation vector.

        If this test fails, the HKDF parameters have changed and all
        previously wrapped DEKs for this (account_key, agent_id) pair
        are unrecoverable. This is a breaking change.
        """
        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        assert kek.hex() == "ea99d5cd36472c5fea59889ad11c717bcabe52ab612b5ad2b72d5f1a60c23c55"


class TestDEKWrapping:
    """Test DEK wrap/unwrap round-trip."""

    def test_round_trip(self):
        """Wrap then unwrap returns original DEK."""
        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        wrapped = wrap_dek(kek, FIXED_DEK)
        recovered = unwrap_dek(kek, wrapped)
        assert recovered == FIXED_DEK

    def test_wrapped_size(self):
        """Wrapped DEK is nonce(12) + ciphertext(32) + tag(16) = 60 bytes."""
        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        wrapped = wrap_dek(kek, FIXED_DEK)
        assert len(wrapped) == 60

    def test_wrapped_different_each_time(self):
        """Random nonce means wrapping same DEK twice produces different output."""
        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        wrapped1 = wrap_dek(kek, FIXED_DEK)
        wrapped2 = wrap_dek(kek, FIXED_DEK)
        assert wrapped1 != wrapped2
        # But both unwrap to the same DEK
        assert unwrap_dek(kek, wrapped1) == FIXED_DEK
        assert unwrap_dek(kek, wrapped2) == FIXED_DEK

    def test_wrong_kek_fails(self):
        """Unwrapping with wrong KEK raises InvalidTag."""
        from cryptography.exceptions import InvalidTag

        kek_correct = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        kek_wrong = derive_kek(FIXED_ACCOUNT_KEY, ALT_AGENT_ID)
        wrapped = wrap_dek(kek_correct, FIXED_DEK)

        with pytest.raises(InvalidTag):
            unwrap_dek(kek_wrong, wrapped)

    def test_tampered_wrapped_fails(self):
        """Modifying wrapped DEK raises InvalidTag."""
        from cryptography.exceptions import InvalidTag

        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        wrapped = wrap_dek(kek, FIXED_DEK)

        tampered = bytearray(wrapped)
        tampered[20] ^= 0xFF  # Flip a byte in the ciphertext region
        with pytest.raises(InvalidTag):
            unwrap_dek(kek, bytes(tampered))

    def test_truncated_wrapped_fails(self):
        """Truncated wrapped DEK fails."""
        from cryptography.exceptions import InvalidTag

        kek = derive_kek(FIXED_ACCOUNT_KEY, FIXED_AGENT_ID)
        wrapped = wrap_dek(kek, FIXED_DEK)

        with pytest.raises((InvalidTag, Exception)):
            unwrap_dek(kek, wrapped[:-1])


class TestBlobKeyDerivation:
    """Test per-blob encryption key derivation."""

    def test_deterministic(self):
        """Same (DEK, stamp_hash, payload_hash) always produces same blob_key."""
        key1 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        key2 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        assert key1 == key2
        assert len(key1) == 32  # AES-256

    def test_different_dek(self):
        """Different DEK produces different blob_key (confidentiality gate)."""
        key1 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        key2 = derive_blob_key(ALT_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        assert key1 != key2

    def test_different_stamp_hash(self):
        """Different stamp_hash produces different blob_key (position binding)."""
        key1 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        key2 = derive_blob_key(FIXED_DEK, ALT_STAMP_HASH, FIXED_PAYLOAD_HASH)
        assert key1 != key2

    def test_different_payload_hash(self):
        """Different payload_hash produces different blob_key (content binding)."""
        key1 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        key2 = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, ALT_PAYLOAD_HASH)
        assert key1 != key2

    def test_pinned_vector(self):
        """
        CANARY TEST: Pinned blob key derivation vector.

        If this test fails, the HKDF parameters have changed and all
        previously encrypted blobs for this (DEK, stamp_hash, payload_hash)
        combination are unrecoverable. This is a breaking change.

        The hex value is computed once from the implementation and frozen.
        """
        key = derive_blob_key(FIXED_DEK, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        assert key.hex() == "a6a4ba42fd6d521ccbe62ec358b4a2fef7f6bb2d7e5e68db7d63f7ed0f93e842"


class TestBlobNonce:
    """Test blob nonce derivation."""

    def test_deterministic(self):
        """Same (DEK, stamp_hash) always produces same nonce."""
        n1 = derive_blob_nonce(FIXED_DEK, FIXED_STAMP_HASH)
        n2 = derive_blob_nonce(FIXED_DEK, FIXED_STAMP_HASH)
        assert n1 == n2
        assert len(n1) == 12  # 96-bit for AES-GCM

    def test_different_stamp_different_nonce(self):
        """Different stamp_hash produces different nonce."""
        n1 = derive_blob_nonce(FIXED_DEK, FIXED_STAMP_HASH)
        n2 = derive_blob_nonce(FIXED_DEK, ALT_STAMP_HASH)
        assert n1 != n2

    def test_different_dek_different_nonce(self):
        """Different DEK produces different nonce (key rotation safety)."""
        n1 = derive_blob_nonce(FIXED_DEK, FIXED_STAMP_HASH)
        n2 = derive_blob_nonce(ALT_DEK, FIXED_STAMP_HASH)
        assert n1 != n2


class TestFullLifecycle:
    """Test the complete envelope encryption lifecycle."""

    def test_account_to_blob_key(self):
        """Full chain: account_key → KEK → wrap DEK → unwrap → derive blob_key."""
        # 1. Operator generates account key
        account_key = generate_account_key()

        # 2. Agent derives KEK from account key + agent ID
        kek = derive_kek(account_key, FIXED_AGENT_ID)
        assert len(kek) == 32

        # 3. Agent generates DEK
        dek = generate_dek()
        assert len(dek) == 32

        # 4. Agent wraps DEK for server storage
        wrapped = wrap_dek(kek, dek)
        assert len(wrapped) == 60

        # 5. On reboot: derive same KEK, unwrap
        kek_again = derive_kek(account_key, FIXED_AGENT_ID)
        assert kek_again == kek
        recovered_dek = unwrap_dek(kek_again, wrapped)
        assert recovered_dek == dek

        # 6. Derive blob key from DEK + chain material
        blob_key = derive_blob_key(recovered_dek, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)
        assert len(blob_key) == 32

        # 7. Same derivation from original DEK produces same blob key
        assert blob_key == derive_blob_key(dek, FIXED_STAMP_HASH, FIXED_PAYLOAD_HASH)

    def test_cross_agent_isolation(self):
        """Two agents with same account key have different KEKs and can't unwrap each other's DEKs."""
        from cryptography.exceptions import InvalidTag

        account_key = FIXED_ACCOUNT_KEY

        # Agent 1
        kek1 = derive_kek(account_key, FIXED_AGENT_ID)
        dek1 = FIXED_DEK
        wrapped1 = wrap_dek(kek1, dek1)

        # Agent 2
        kek2 = derive_kek(account_key, ALT_AGENT_ID)
        dek2 = ALT_DEK
        wrapped2 = wrap_dek(kek2, dek2)

        # Each can unwrap their own
        assert unwrap_dek(kek1, wrapped1) == dek1
        assert unwrap_dek(kek2, wrapped2) == dek2

        # Cannot unwrap each other's
        with pytest.raises(InvalidTag):
            unwrap_dek(kek1, wrapped2)
        with pytest.raises(InvalidTag):
            unwrap_dek(kek2, wrapped1)

    def test_server_cannot_derive_blob_key(self):
        """Without the DEK, knowing stamp_hash and payload_hash is insufficient."""
        # The server sees these (from the chain):
        stamp_hash = FIXED_STAMP_HASH
        payload_hash = FIXED_PAYLOAD_HASH

        # The real blob key (requires DEK):
        real_key = derive_blob_key(FIXED_DEK, stamp_hash, payload_hash)

        # Any other DEK produces a different key:
        fake_key = derive_blob_key(ALT_DEK, stamp_hash, payload_hash)
        assert real_key != fake_key

        # Zero DEK produces a different key:
        zero_key = derive_blob_key(b'\x00' * 32, stamp_hash, payload_hash)
        assert real_key != zero_key
