//! SMB3 encryption — AES-128-GCM / AES-128-CCM (MS-SMB2 §2.2.41, §3.1.4.3).
//!
//! Uses AEAD modes with the SMB2 TRANSFORM_HEADER as AAD
//! (Additional Authenticated Data).  Key derivation follows
//! SP 800-108 CTR-mode KDF (MS-SMB2 §3.1.4.2), re-using the
//! existing [`crate::proto::crypto::kdf::smb2_kdf`] primitive.
//!
//! Supported ciphers:
//!   * AES-128-GCM — 12-byte nonce, parallelisable, SMB 3.1.1+ (Windows 10+)
//!   * AES-128-CCM — 11-byte nonce, sequential, SMB 3.0 (Windows 8)

use aes_gcm::{
    aead::{Aead, KeyInit, Payload as GcmPayload},
    Aes128Gcm as Aes128GcmCipher, Nonce as GcmNonce,
};
use binrw::{binrw, BinWrite, BinRead, io::Cursor, Endian};
use ccm::{
    aead::{Aead as CcmAead, KeyInit as CcmKeyInit, Payload as CcmPayload},
    Ccm as Aes128CcmCipher, Nonce as CcmNonce,
};
use aes::Aes128;
use thiserror::Error;

type Aes128Ccm = Aes128CcmCipher<Aes128, typenum::U16, typenum::U11>;

// Re-export common AEAD traits for callers that need them.
pub use aes_gcm::aead::generic_array::typenum;

#[derive(Debug, Error)]
pub enum EncryptionError {
    #[error("Invalid transform header signature")]
    InvalidSignature,
    #[error("Unsupported cipher algorithm: {0}")]
    UnsupportedCipher(u16),
    #[error("Encryption failed: {0}")]
    EncryptionFailed(String),
    #[error("Decryption failed: {0}")]
    DecryptionFailed(String),
    #[error("Invalid key length")]
    InvalidKeyLength,
    #[error("Session key not set")]
    NoSessionKey,
}

/// SMB2 TRANSFORM_HEADER (MS-SMB2 §2.2.41) — 56 bytes.
///
/// For AES-128-GCM:
///   * Nonce = 12 bytes (first 12 of the 16-byte field; last 4 reserved).
///   * Signature = GCM authentication tag (16 bytes).
///
/// For AES-128-CCM:
///   * Nonce = 11 bytes (first 11 of the 16-byte field; last 5 reserved).
///   * Signature = CCM authentication tag (16 bytes).
///
/// In both cases AAD = entire header except the signature + encrypted data.
#[binrw]
#[brw(big, magic = 0x534D4272u32)]  // "SMBr" — SMB3 encrypted protocol id
pub struct TransformHeader {
    #[brw(little)]
    pub cipher_algorithm: u16,     // 0x0001 = AES-128-GCM, 0x0002 = AES-128-CCM
    #[brw(little)]
    pub cipher_key_length: u16,    // 16 bytes
    #[brw(little)]
    pub nonce: [u8; 16],           // 12 (GCM) or 11 (CCM) bytes used, rest reserved
    #[brw(little)]
    pub session_id: u64,
    #[brw(little)]
    pub original_message_size: u32,
    #[brw(little)]
    pub reserved1: u16,
    #[brw(little)]
    pub reserved2: u16,
    pub signature: [u8; 16],       // AEAD authentication tag
    // EncryptedData follows (variable length)
}

impl TransformHeader {
    pub const SIZE: usize = 56;

    pub fn write_to_bytes(&self) -> Result<Vec<u8>, EncryptionError> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&0x534D4272u32.to_be_bytes());
        bytes.extend_from_slice(&self.cipher_algorithm.to_le_bytes());
        bytes.extend_from_slice(&self.cipher_key_length.to_le_bytes());
        bytes.extend_from_slice(&self.nonce);
        bytes.extend_from_slice(&self.session_id.to_le_bytes());
        bytes.extend_from_slice(&self.original_message_size.to_le_bytes());
        bytes.extend_from_slice(&self.reserved1.to_le_bytes());
        bytes.extend_from_slice(&self.reserved2.to_le_bytes());
        bytes.extend_from_slice(&self.signature);
        Ok(bytes)
    }

    pub fn read_from_bytes(data: &[u8]) -> Result<Self, EncryptionError> {
        if data.len() < Self::SIZE {
            return Err(EncryptionError::DecryptionFailed(
                "Header too short".to_string(),
            ));
        }

        let magic = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
        if magic != 0x534D4272 {
            return Err(EncryptionError::InvalidSignature);
        }

        Ok(Self {
            cipher_algorithm: u16::from_le_bytes([data[4], data[5]]),
            cipher_key_length: u16::from_le_bytes([data[6], data[7]]),
            nonce: {
                let mut n = [0u8; 16];
                n.copy_from_slice(&data[8..24]);
                n
            },
            session_id: u64::from_le_bytes(data[24..32].try_into().unwrap()),
            original_message_size: u32::from_le_bytes(data[32..36].try_into().unwrap()),
            reserved1: u16::from_le_bytes([data[36], data[37]]),
            reserved2: u16::from_le_bytes([data[38], data[39]]),
            signature: {
                let mut s = [0u8; 16];
                s.copy_from_slice(&data[40..56]);
                s
            },
        })
    }

    /// Build AAD = header[0..52], i.e. everything before `signature`.
    fn build_aad(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(40);
        buf.extend_from_slice(&0x534D4272u32.to_be_bytes());
        buf.extend_from_slice(&self.cipher_algorithm.to_le_bytes());
        buf.extend_from_slice(&self.cipher_key_length.to_le_bytes());
        buf.extend_from_slice(&self.nonce);
        buf.extend_from_slice(&self.session_id.to_le_bytes());
        buf.extend_from_slice(&self.original_message_size.to_le_bytes());
        buf.extend_from_slice(&self.reserved1.to_le_bytes());
        buf.extend_from_slice(&self.reserved2.to_le_bytes());
        buf
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CipherAlgorithm {
    Aes128Gcm = 0x0001,
    Aes128Ccm = 0x0002,
}

impl CipherAlgorithm {
    pub fn from_u16(value: u16) -> Option<Self> {
        match value {
            0x0001 => Some(CipherAlgorithm::Aes128Gcm),
            0x0002 => Some(CipherAlgorithm::Aes128Ccm),
            _ => None,
        }
    }

    pub fn key_length(&self) -> u16 {
        16
    }

    /// Number of nonce bytes used by this cipher.
    pub fn nonce_length(&self) -> usize {
        match self {
            CipherAlgorithm::Aes128Gcm => 12,
            CipherAlgorithm::Aes128Ccm => 11,
        }
    }
}

/// Per-session SMB3 encryption helper.
///
/// Supports both AES-128-GCM (SMB 3.1.1+) and AES-128-CCM (SMB 3.0).
pub struct Smb3Encryption {
    encryption_key: [u8; 16],
    cipher: CipherAlgorithm,
}

impl Smb3Encryption {
    /// Create a new encryption context from the session key and cipher.
    ///
    /// Derives the AES-128 key via SP 800-108 KDF.
    pub fn new(session_key: &[u8], cipher_algorithm: CipherAlgorithm) -> Result<Self, EncryptionError> {
        if session_key.len() != 16 {
            return Err(EncryptionError::InvalidKeyLength);
        }

        let encryption_key = Self::derive_encryption_key_sp800108(session_key, b"SMB3ENC");

        Ok(Self {
            encryption_key,
            cipher: cipher_algorithm,
        })
    }

    /// Encrypt a plaintext SMB2 message.
    ///
    /// Returns a complete SMB3 TRANSFORM_HEADER + encrypted payload.
    pub fn encrypt_packet(&self, plaintext: &[u8], session_id: u64) -> Result<Vec<u8>, EncryptionError> {
        let nonce_len = self.cipher.nonce_length();

        // Generate random nonce, pad to 16 bytes in the header
        let mut nonce_full = [0u8; 16];
        getrandom::fill(&mut nonce_full[..nonce_len])
            .map_err(|e| EncryptionError::EncryptionFailed(format!("nonce: {}", e)))?;

        let header_no_tag = TransformHeader {
            cipher_algorithm: self.cipher as u16,
            cipher_key_length: 16,
            nonce: nonce_full,
            session_id,
            original_message_size: plaintext.len() as u32,
            reserved1: 0,
            reserved2: 0,
            signature: [0u8; 16],
        };

        let aad = header_no_tag.build_aad();

        // AEAD encrypt: returns ciphertext || tag (last 16 bytes)
        let ciphertext_with_tag = match self.cipher {
            CipherAlgorithm::Aes128Gcm => {
                let nonce12 = GcmNonce::from_slice(&nonce_full[..12]);
                let cipher = Aes128GcmCipher::new_from_slice(&self.encryption_key)
                    .map_err(|e| EncryptionError::EncryptionFailed(format!("GCM key: {}", e)))?;
                cipher
                    .encrypt(nonce12, GcmPayload { msg: plaintext, aad: &aad })
                    .map_err(|e| EncryptionError::EncryptionFailed(format!("GCM encrypt: {}", e)))?
            }
            CipherAlgorithm::Aes128Ccm => {
                let nonce11 = CcmNonce::from_slice(&nonce_full[..11]);
                let cipher = Aes128Ccm::new_from_slice(&self.encryption_key)
                    .map_err(|e| EncryptionError::EncryptionFailed(format!("CCM key: {}", e)))?;
                cipher
                    .encrypt(nonce11, CcmPayload { msg: plaintext, aad: &aad })
                    .map_err(|e| EncryptionError::EncryptionFailed(format!("CCM encrypt: {}", e)))?
            }
        };

        let tag_len = 16;
        let tag_pos = ciphertext_with_tag.len().saturating_sub(tag_len);
        let tag: [u8; 16] = ciphertext_with_tag[tag_pos..]
            .try_into()
            .map_err(|_| EncryptionError::EncryptionFailed("tag extraction".to_string()))?;
        let encrypted_data = &ciphertext_with_tag[..tag_pos];

        let header = TransformHeader {
            signature: tag,
            ..header_no_tag
        };

        let mut packet = header.write_to_bytes()?;
        packet.extend_from_slice(encrypted_data);
        Ok(packet)
    }

    /// Decrypt an SMB3 TRANSFORM_HEADER payload.
    ///
    /// The cipher algorithm is read from the header's `cipher_algorithm` field,
    /// so this is dispatch-safe — callers don't need to match the algorithm.
    pub fn decrypt_packet(&self, encrypted_packet: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        let header = TransformHeader::read_from_bytes(encrypted_packet)?;
        let encrypted_data = &encrypted_packet[TransformHeader::SIZE..];

        // Determine cipher from header (prefer the stored self.cipher but
        // also verify the header's opinion matches).
        let cipher = CipherAlgorithm::from_u16(header.cipher_algorithm)
            .unwrap_or(self.cipher);
        let _nonce_len = cipher.nonce_length();

        let aad = header.build_aad();

        // Build ciphertext_with_tag for AEAD verification
        let mut ct_with_tag = encrypted_data.to_vec();
        ct_with_tag.extend_from_slice(&header.signature);

        match cipher {
            CipherAlgorithm::Aes128Gcm => {
                let mut nonce_buf = [0u8; 12];
                nonce_buf.copy_from_slice(&header.nonce[..12]);
                let nonce12 = GcmNonce::from_slice(&nonce_buf);
                let cipher = Aes128GcmCipher::new_from_slice(&self.encryption_key)
                    .map_err(|e| EncryptionError::DecryptionFailed(format!("GCM key: {}", e)))?;
                cipher
                    .decrypt(nonce12, GcmPayload { msg: &ct_with_tag, aad: &aad })
                    .map_err(|_| EncryptionError::InvalidSignature)
            }
            CipherAlgorithm::Aes128Ccm => {
                let mut nonce_buf = [0u8; 11];
                nonce_buf.copy_from_slice(&header.nonce[..11]);
                let nonce11 = CcmNonce::from_slice(&nonce_buf);
                let cipher = Aes128Ccm::new_from_slice(&self.encryption_key)
                    .map_err(|e| EncryptionError::DecryptionFailed(format!("CCM key: {}", e)))?;
                cipher
                    .decrypt(nonce11, CcmPayload { msg: &ct_with_tag, aad: &aad })
                    .map_err(|_| EncryptionError::InvalidSignature)
            }
        }
    }

    /// Derive AES-128 encryption key via SP 800-108 KDF.
    ///
    /// Uses the existing [`crate::proto::crypto::kdf::smb2_kdf`] with
    /// Label = `label` (caller includes trailing NUL), Context = empty.
    ///
    /// MS-SMB2 §3.1.4.2: `encryption_key = KDF(session_key, label, "")`.
    pub fn derive_encryption_key_sp800108(session_key: &[u8], label: &[u8]) -> [u8; 16] {
        let mut label_with_nul = label.to_vec();
        label_with_nul.push(0x00);
        let context_with_nul = b"\x00";

        crate::proto::crypto::kdf::smb2_kdf(session_key, &label_with_nul, context_with_nul)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_encrypt_decrypt_roundtrip(cipher: CipherAlgorithm) {
        let session_key = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
        let enc = Smb3Encryption::new(&session_key, cipher).unwrap();

        let plaintext = b"Hello SMB3!";
        let session_id = 12345u64;

        let encrypted = enc.encrypt_packet(plaintext, session_id).unwrap();

        assert_eq!(encrypted.len(), TransformHeader::SIZE + plaintext.len());

        let magic = u32::from_be_bytes([encrypted[0], encrypted[1], encrypted[2], encrypted[3]]);
        assert_eq!(magic, 0x534D4272);

        // Verify cipher_algorithm field in header
        let header_cipher = u16::from_le_bytes([encrypted[4], encrypted[5]]);
        assert_eq!(header_cipher, cipher as u16);

        let decrypted = enc.decrypt_packet(&encrypted).unwrap();
        assert_eq!(plaintext.as_slice(), decrypted.as_slice());
    }

    #[test]
    fn test_gcm_roundtrip() {
        test_encrypt_decrypt_roundtrip(CipherAlgorithm::Aes128Gcm);
    }

    #[test]
    fn test_ccm_roundtrip() {
        test_encrypt_decrypt_roundtrip(CipherAlgorithm::Aes128Ccm);
    }

    #[test]
    fn test_gcm_and_ccm_interop() {
        // Verify packets encrypted with different ciphers produce different wire output
        let session_key = [1u8; 16];
        let plaintext = b"Cross-cipher test";

        let gcm_enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Gcm).unwrap();
        let ccm_enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Ccm).unwrap();

        let gcm_packet = gcm_enc.encrypt_packet(plaintext, 1).unwrap();
        let ccm_packet = ccm_enc.encrypt_packet(plaintext, 1).unwrap();

        // Different cipher algorithm IDs in the header
        assert_eq!(
            u16::from_le_bytes([gcm_packet[4], gcm_packet[5]]),
            CipherAlgorithm::Aes128Gcm as u16
        );
        assert_eq!(
            u16::from_le_bytes([ccm_packet[4], ccm_packet[5]]),
            CipherAlgorithm::Aes128Ccm as u16
        );

        // Ciphertext differs (different nonce length → different keystream offset)
        assert_ne!(gcm_packet, ccm_packet);

        // Each cipher can decrypt its own packet via the header-based dispatch
        assert!(gcm_enc.decrypt_packet(&gcm_packet).is_ok());
        assert!(ccm_enc.decrypt_packet(&ccm_packet).is_ok());
    }

    #[test]
    fn test_cipher_algorithm_conversion() {
        assert_eq!(CipherAlgorithm::from_u16(0x0001), Some(CipherAlgorithm::Aes128Gcm));
        assert_eq!(CipherAlgorithm::from_u16(0x0002), Some(CipherAlgorithm::Aes128Ccm));
        assert_eq!(CipherAlgorithm::from_u16(0x0003), None);
    }

    #[test]
    fn test_gcm_authentication_failure() {
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Gcm).unwrap();
        let encrypted = enc.encrypt_packet(b"Test data", 999).unwrap();

        let mut tampered = encrypted.clone();
        tampered[TransformHeader::SIZE] ^= 0xFF;

        let result = enc.decrypt_packet(&tampered);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Invalid transform header signature");
    }

    #[test]
    fn test_ccm_authentication_failure() {
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Ccm).unwrap();
        let encrypted = enc.encrypt_packet(b"Test data", 999).unwrap();

        let mut tampered = encrypted.clone();
        tampered[TransformHeader::SIZE] ^= 0xFF;

        let result = enc.decrypt_packet(&tampered);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Invalid transform header signature");
    }

    #[test]
    fn test_gcm_tag_tampering() {
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Gcm).unwrap();
        let encrypted = enc.encrypt_packet(b"Test data", 999).unwrap();

        let mut tampered = encrypted;
        tampered[48] ^= 0xFF;

        assert!(enc.decrypt_packet(&tampered).is_err());
    }

    #[test]
    fn test_ccm_tag_tampering() {
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Ccm).unwrap();
        let encrypted = enc.encrypt_packet(b"Test data", 999).unwrap();

        let mut tampered = encrypted;
        tampered[48] ^= 0xFF;

        assert!(enc.decrypt_packet(&tampered).is_err());
    }

    #[test]
    fn test_nonce_uniqueness() {
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Gcm).unwrap();

        let p1 = enc.encrypt_packet(b"Same data", 1).unwrap();
        let p2 = enc.encrypt_packet(b"Same data", 2).unwrap();

        let nonce1: [u8; 16] = p1[8..24].try_into().unwrap();
        let nonce2: [u8; 16] = p2[8..24].try_into().unwrap();
        assert_ne!(nonce1, nonce2);
    }

    #[test]
    fn test_ccm_nonce_length() {
        // CCM uses 11-byte nonce (verify the header stores it correctly)
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Ccm).unwrap();
        let encrypted = enc.encrypt_packet(b"nonce test", 1).unwrap();

        // The header nonce field is always 16 bytes, but CCM only uses 11
        let nonce: [u8; 16] = encrypted[8..24].try_into().unwrap();
        // Bytes 11-15 should be zero (padding/reserved)
        assert_eq!(&nonce[11..], &[0, 0, 0, 0, 0]);
    }

    #[test]
    fn test_gcm_nonce_length() {
        // GCM uses 12-byte nonce
        let session_key = [1u8; 16];
        let enc = Smb3Encryption::new(&session_key, CipherAlgorithm::Aes128Gcm).unwrap();
        let encrypted = enc.encrypt_packet(b"nonce test", 1).unwrap();

        let nonce: [u8; 16] = encrypted[8..24].try_into().unwrap();
        // Bytes 12-15 should be zero
        assert_eq!(&nonce[12..], &[0, 0, 0, 0]);
    }

    #[test]
    fn test_sp800108_kdf_known_answer() {
        let session_key = [0u8; 16];
        let key = Smb3Encryption::derive_encryption_key_sp800108(&session_key, b"SMB3ENC");

        let label = b"SMB3ENC\x00";
        let context = b"\x00";
        let expected = crate::proto::crypto::kdf::smb2_kdf(&session_key, label, context);
        assert_eq!(key, expected);
        assert_ne!(key, [0u8; 16]);
    }

    #[test]
    fn test_different_sessions_different_keys() {
        let key1 = Smb3Encryption::derive_encryption_key_sp800108(&[1u8; 16], b"SMB3ENC");
        let key2 = Smb3Encryption::derive_encryption_key_sp800108(&[2u8; 16], b"SMB3ENC");
        assert_ne!(key1, key2);
    }
}