"""Tests for i2p_crypto.aes_cbc — AES-256-CBC for NTCP2 handshake obfuscation.

Covers:
  - Round-trip encrypt/decrypt for 16, 32, 48 bytes
  - 32-byte block specifically (the NTCP2 use case)
  - Known NIST AES-256-CBC test vectors (SP 800-38A F.2.5/F.2.6)
  - Ciphertext differs from plaintext
  - Different IVs produce different ciphertext
  - Wrong key/IV fails to decrypt correctly
  - IV chaining: encrypting 32 bytes, bytes 16-31 of ciphertext become new IV
  - Input validation (key size, IV size, block alignment)
"""

from __future__ import annotations

import hashlib
import os

import pytest

from i2p_crypto.aes_cbc import aes_cbc_decrypt, aes_cbc_encrypt


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _random_key() -> bytes:
    """Return a random 32-byte AES-256 key."""
    return os.urandom(32)


def _random_iv() -> bytes:
    """Return a random 16-byte IV."""
    return os.urandom(16)


# ---------------------------------------------------------------------------
# Round-trip tests
# ---------------------------------------------------------------------------

class TestRoundTrip:
    """encrypt then decrypt must return the original plaintext."""

    def test_16_bytes(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(16)
        ciphertext = aes_cbc_encrypt(key, iv, plaintext)
        assert len(ciphertext) == 16
        assert aes_cbc_decrypt(key, iv, ciphertext) == plaintext

    def test_32_bytes(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(32)
        ciphertext = aes_cbc_encrypt(key, iv, plaintext)
        assert len(ciphertext) == 32
        assert aes_cbc_decrypt(key, iv, ciphertext) == plaintext

    def test_48_bytes(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(48)
        ciphertext = aes_cbc_encrypt(key, iv, plaintext)
        assert len(ciphertext) == 48
        assert aes_cbc_decrypt(key, iv, ciphertext) == plaintext

    def test_deterministic(self):
        """Same key+IV+plaintext must produce identical ciphertext."""
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(32)
        c1 = aes_cbc_encrypt(key, iv, plaintext)
        c2 = aes_cbc_encrypt(key, iv, plaintext)
        assert c1 == c2


# ---------------------------------------------------------------------------
# NTCP2 32-byte use case
# ---------------------------------------------------------------------------

class TestNTCP2UseCase:
    """NTCP2 handshake message 1: 32-byte ephemeral key X encrypted."""

    def test_ntcp2_encrypt_decrypt_32_byte_key(self):
        """Simulate NTCP2 message 1 obfuscation of ephemeral key X.

        Key = SHA-256(responder RouterInfo hash)
        IV  = responder 'i' parameter (16 bytes)
        Plaintext = 32 bytes (ephemeral key X, exactly 2 AES blocks)
        """
        router_info_hash = os.urandom(32)
        aes_key = hashlib.sha256(router_info_hash).digest()
        iv = os.urandom(16)  # responder's "i" parameter
        ephemeral_key_x = os.urandom(32)

        ciphertext = aes_cbc_encrypt(aes_key, iv, ephemeral_key_x)
        assert len(ciphertext) == 32
        recovered = aes_cbc_decrypt(aes_key, iv, ciphertext)
        assert recovered == ephemeral_key_x

    def test_iv_chaining_for_message_2(self):
        """After encrypting 32 bytes, bytes 16-31 of ciphertext become
        the IV for message 2 encryption.

        This tests that the IV chaining property of CBC holds: the last
        ciphertext block is usable as an IV for subsequent encryption.
        """
        aes_key = _random_key()
        iv1 = _random_iv()
        plaintext_msg1 = os.urandom(32)

        ciphertext_msg1 = aes_cbc_encrypt(aes_key, iv1, plaintext_msg1)

        # In NTCP2, bytes 16-31 of the first ciphertext become IV for msg2
        iv2 = ciphertext_msg1[16:32]

        plaintext_msg2 = os.urandom(32)
        ciphertext_msg2 = aes_cbc_encrypt(aes_key, iv2, plaintext_msg2)
        recovered_msg2 = aes_cbc_decrypt(aes_key, iv2, ciphertext_msg2)
        assert recovered_msg2 == plaintext_msg2

        # The two ciphertexts should differ (different IVs, different plaintexts)
        assert ciphertext_msg1 != ciphertext_msg2


# ---------------------------------------------------------------------------
# Known test vectors — NIST SP 800-38A AES-256-CBC (F.2.5 / F.2.6)
# ---------------------------------------------------------------------------

class TestKnownVectors:
    """NIST SP 800-38A AES-256-CBC test vectors."""

    NIST_KEY = bytes.fromhex(
        "603deb1015ca71be2b73aef0857d7781"
        "1f352c073b6108d72d9810a30914dff4"
    )
    NIST_IV = bytes.fromhex("000102030405060708090a0b0c0d0e0f")

    NIST_PLAINTEXT = bytes.fromhex(
        "6bc1bee22e409f96e93d7e117393172a"
        "ae2d8a571e03ac9c9eb76fac45af8e51"
        "30c81c46a35ce411e5fbc1191a0a52ef"
        "f69f2445df4f9b17ad2b417be66c3710"
    )

    NIST_CIPHERTEXT = bytes.fromhex(
        "f58c4c04d6e5f1ba779eabfb5f7bfbd6"
        "9cfc4e967edb808d679f777bc6702c7d"
        "39f23369a9d9bacfa530e26304231461"
        "b2eb05e2c39be9fcda6c19078c6a9d1b"
    )

    def test_encrypt(self):
        ct = aes_cbc_encrypt(self.NIST_KEY, self.NIST_IV, self.NIST_PLAINTEXT)
        assert ct == self.NIST_CIPHERTEXT

    def test_decrypt(self):
        pt = aes_cbc_decrypt(self.NIST_KEY, self.NIST_IV, self.NIST_CIPHERTEXT)
        assert pt == self.NIST_PLAINTEXT

    def test_round_trip(self):
        ct = aes_cbc_encrypt(self.NIST_KEY, self.NIST_IV, self.NIST_PLAINTEXT)
        pt = aes_cbc_decrypt(self.NIST_KEY, self.NIST_IV, ct)
        assert pt == self.NIST_PLAINTEXT


# ---------------------------------------------------------------------------
# Ciphertext differs from plaintext
# ---------------------------------------------------------------------------

class TestCiphertextDiffers:
    """Ciphertext must not equal plaintext."""

    def test_ciphertext_not_plaintext_16(self):
        plaintext = os.urandom(16)
        ct = aes_cbc_encrypt(_random_key(), _random_iv(), plaintext)
        assert ct != plaintext

    def test_ciphertext_not_plaintext_32(self):
        plaintext = os.urandom(32)
        ct = aes_cbc_encrypt(_random_key(), _random_iv(), plaintext)
        assert ct != plaintext


# ---------------------------------------------------------------------------
# Different IVs produce different ciphertext
# ---------------------------------------------------------------------------

class TestDifferentIVs:
    """Same key + plaintext but different IVs must yield different ciphertext."""

    def test_different_ivs_different_ciphertext(self):
        key = _random_key()
        plaintext = os.urandom(32)
        iv1 = _random_iv()
        iv2 = _random_iv()
        # Ensure IVs are actually different (astronomically unlikely to collide)
        assert iv1 != iv2

        ct1 = aes_cbc_encrypt(key, iv1, plaintext)
        ct2 = aes_cbc_encrypt(key, iv2, plaintext)
        assert ct1 != ct2


# ---------------------------------------------------------------------------
# Wrong key / wrong IV fails to decrypt correctly
# ---------------------------------------------------------------------------

class TestWrongKeyIV:
    """Decrypting with wrong key or IV must not return original plaintext."""

    def test_wrong_key(self):
        key1, key2 = _random_key(), _random_key()
        iv = _random_iv()
        plaintext = os.urandom(32)
        ciphertext = aes_cbc_encrypt(key1, iv, plaintext)
        decrypted = aes_cbc_decrypt(key2, iv, ciphertext)
        assert decrypted != plaintext

    def test_wrong_iv(self):
        key = _random_key()
        iv1, iv2 = _random_iv(), _random_iv()
        plaintext = os.urandom(32)
        ciphertext = aes_cbc_encrypt(key, iv1, plaintext)
        decrypted = aes_cbc_decrypt(key, iv2, ciphertext)
        assert decrypted != plaintext


# ---------------------------------------------------------------------------
# Input validation
# ---------------------------------------------------------------------------

class TestValidation:
    """Key, IV, and plaintext sizes must be enforced."""

    def test_key_not_32_bytes(self):
        with pytest.raises(ValueError, match="32 bytes"):
            aes_cbc_encrypt(os.urandom(16), _random_iv(), os.urandom(16))

    def test_iv_not_16_bytes(self):
        with pytest.raises(ValueError, match="16 bytes"):
            aes_cbc_encrypt(_random_key(), os.urandom(12), os.urandom(16))

    def test_plaintext_not_aligned(self):
        with pytest.raises(ValueError, match="multiple of 16"):
            aes_cbc_encrypt(_random_key(), _random_iv(), os.urandom(15))

    def test_plaintext_empty(self):
        with pytest.raises(ValueError, match="empty"):
            aes_cbc_encrypt(_random_key(), _random_iv(), b"")

    def test_ciphertext_not_aligned(self):
        with pytest.raises(ValueError, match="multiple of 16"):
            aes_cbc_decrypt(_random_key(), _random_iv(), os.urandom(17))

    def test_ciphertext_empty(self):
        with pytest.raises(ValueError, match="empty"):
            aes_cbc_decrypt(_random_key(), _random_iv(), b"")