crypto.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use rand::{Rng, thread_rng};
use base64::{Engine as _, engine::general_purpose};
use aes_gcm::{Aes256Gcm, KeyInit, aead::Aead, Nonce};
use hex;

#[derive(Debug, Serialize, Deserialize)]
pub struct CryptoKeyPair {
    /// AES-256 encryption key (32 bytes, base64 encoded for storage)
    pub private_key: String,
    /// Key identifier (for compatibility, not used in encryption)
    pub public_key: String,
    /// Unique key identifier
    pub key_id: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct EncryptedData {
    /// Encrypted data (base64 encoded ciphertext with authentication tag)
    pub data: String,
    /// Initialization vector (nonce) for AES-GCM (12 bytes, base64 encoded)
    pub iv: String,
    /// Authentication tag (included in data for AES-GCM, kept for compatibility)
    pub tag: String,
    /// Key identifier used for encryption
    pub key_id: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SecurityLevel {
    pub level: u8,
    pub tests_passed: u8,
    pub total_tests: u8,
    pub details: Vec<String>,
}

pub struct CryptoUtils {
    key_pairs: HashMap<String, CryptoKeyPair>,
    current_key_id: Option<String>,
}

impl CryptoUtils {
    pub fn new() -> Self {
        Self {
            key_pairs: HashMap::new(),
            current_key_id: None,
        }
    }

    /// Generate a new AES-256 key pair
    /// Returns the key_id for the generated key
    pub fn generate_key_pair(&mut self) -> Result<String, String> {
        let key_id = self.generate_key_id();
        
        // Generate cryptographically secure AES-256 key (32 bytes)
        let mut aes_key = [0u8; 32];
        thread_rng().fill(&mut aes_key);
        
        // Encode key as base64 for storage in String
        let private_key = general_purpose::STANDARD.encode(&aes_key);
        
        // Generate a random identifier for public_key (for compatibility)
        let public_key = self.generate_key_id();
        
        let key_pair = CryptoKeyPair {
            private_key,
            public_key,
            key_id: key_id.clone(),
        };
        
        self.key_pairs.insert(key_id.clone(), key_pair);
        self.current_key_id = Some(key_id.clone());
        
        Ok(key_id)
    }

    /// Generate a secure random key ID (32 hex characters = 16 bytes)
    fn generate_key_id(&self) -> String {
        let mut rng = thread_rng();
        let bytes: [u8; 16] = rng.gen();
        hex::encode(bytes)
    }


    /// Encrypt data using AES-256-GCM
    /// Uses authenticated encryption for confidentiality and integrity
    pub fn encrypt_data(&self, data: &str, key_id: &str) -> Result<EncryptedData, String> {
        let key_pair = self.key_pairs.get(key_id)
            .ok_or("Key not found")?;
        
        // Decode AES key from base64
        let aes_key_bytes = general_purpose::STANDARD.decode(&key_pair.private_key)
            .map_err(|_| "Invalid key format")?;
        
        if aes_key_bytes.len() != 32 {
            return Err("Invalid key length (expected 32 bytes)".to_string());
        }
        
        // Generate random 12-byte nonce (IV) for AES-GCM
        let mut nonce_bytes = [0u8; 12];
        thread_rng().fill(&mut nonce_bytes);
        
        // Create AES-256-GCM cipher
        let cipher = Aes256Gcm::new_from_slice(&aes_key_bytes)
            .map_err(|e| format!("Failed to create AES-GCM cipher: {}", e))?;
        
        // Encrypt data (AES-GCM includes authentication tag in ciphertext)
        let nonce = Nonce::from(nonce_bytes);
        let ciphertext = cipher.encrypt(&nonce, data.as_bytes())
            .map_err(|e| format!("Failed to encrypt data: {}", e))?;
        
        // Encode IV and ciphertext as base64
        let iv = general_purpose::STANDARD.encode(&nonce_bytes);
        let data_encoded = general_purpose::STANDARD.encode(&ciphertext);
        
        // Note: AES-GCM tag is included in ciphertext, but we keep tag field for compatibility
        // Extract tag (last 16 bytes of ciphertext) for compatibility
        let tag = if ciphertext.len() >= 16 {
            general_purpose::STANDARD.encode(&ciphertext[ciphertext.len() - 16..])
        } else {
            String::new()
        };
        
        Ok(EncryptedData {
            data: data_encoded,
            iv,
            tag,
            key_id: key_id.to_string(),
        })
    }

    /// Decrypt data using AES-256-GCM
    /// Verifies authentication tag for integrity
    pub fn decrypt_data(&self, encrypted_data: &EncryptedData) -> Result<String, String> {
        let key_pair = self.key_pairs.get(&encrypted_data.key_id)
            .ok_or("Key not found")?;
        
        // Decode AES key from base64
        let aes_key_bytes = general_purpose::STANDARD.decode(&key_pair.private_key)
            .map_err(|_| "Invalid key format")?;
        
        if aes_key_bytes.len() != 32 {
            return Err("Invalid key length (expected 32 bytes)".to_string());
        }
        
        // Decode IV (nonce) from base64
        let nonce_bytes_vec = general_purpose::STANDARD.decode(&encrypted_data.iv)
            .map_err(|_| "Invalid IV format")?;
        
        if nonce_bytes_vec.len() != 12 {
            return Err("Invalid IV length (expected 12 bytes)".to_string());
        }
        
        // Convert Vec<u8> to [u8; 12] for Nonce
        let mut nonce_bytes = [0u8; 12];
        nonce_bytes.copy_from_slice(&nonce_bytes_vec);
        
        // Decode ciphertext from base64
        let ciphertext = general_purpose::STANDARD.decode(&encrypted_data.data)
            .map_err(|_| "Invalid ciphertext format")?;
        
        // Create AES-256-GCM cipher
        let cipher = Aes256Gcm::new_from_slice(&aes_key_bytes)
            .map_err(|e| format!("Failed to create AES-GCM cipher: {}", e))?;
        
        // Decrypt and verify authentication tag
        let nonce = Nonce::from(nonce_bytes);
        let plaintext = cipher.decrypt(&nonce, ciphertext.as_slice())
            .map_err(|e| format!("Failed to decrypt data: {}", e))?;
        
        // Convert decrypted bytes to string
        String::from_utf8(plaintext)
            .map_err(|_| "Invalid UTF-8 in decrypted data".to_string())
    }

    /// Generate a secure password
    pub fn generate_secure_password(&self, length: usize) -> String {
        const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*";
        let mut rng = thread_rng();
        
        (0..length)
            .map(|_| {
                let idx = rng.gen_range(0..CHARSET.len());
                CHARSET[idx] as char
            })
            .collect()
    }

    /// Calculate security level based on various tests.
    ///
    /// `session_established` / `ratchet_active` come from the live session
    /// state: the panel must measure what it displays (web 5.7.1 parity), not
    /// return fixed results for the two properties users care about most.
    pub fn calculate_security_level(&self, session_established: bool, ratchet_active: bool) -> SecurityLevel {
        let mut tests_passed = 0;
        let mut total_tests = 0;
        let mut details = Vec::new();

        // Test 1: Key Generation
        total_tests += 1;
        if !self.key_pairs.is_empty() {
            tests_passed += 1;
            details.push("Key Generation: PASSED".to_string());
        } else {
            details.push("Key Generation: FAILED".to_string());
        }

        // Test 2: Encryption
        total_tests += 1;
        if let Some(key_id) = &self.current_key_id {
            if let Ok(_) = self.encrypt_data("test", key_id) {
                tests_passed += 1;
                details.push("Encryption: PASSED".to_string());
            } else {
                details.push("Encryption: FAILED".to_string());
            }
        } else {
            details.push("Encryption: FAILED".to_string());
        }

        // Test 3: Message Integrity
        total_tests += 1;
        // AES-GCM provides authenticated encryption (integrity is verified during decryption)
        if let Some(key_id) = &self.current_key_id {
            if let Ok(encrypted) = self.encrypt_data("test", key_id) {
                if let Ok(_) = self.decrypt_data(&encrypted) {
                    tests_passed += 1;
                    details.push("Message Integrity: PASSED (AES-GCM authentication)".to_string());
                } else {
                    details.push("Message Integrity: FAILED (decryption failed)".to_string());
                }
            } else {
                details.push("Message Integrity: FAILED (encryption failed)".to_string());
            }
        } else {
            details.push("Message Integrity: FAILED (no key available)".to_string());
        }

        // Test 4: Perfect Forward Secrecy — measured, not asserted. Session-level
        // PFS exists whenever the ephemeral handshake produced keys; per-MESSAGE
        // forward secrecy exists only while the Double Ratchet is running.
        total_tests += 1;
        if ratchet_active {
            tests_passed += 1;
            details.push("Perfect Forward Secrecy: PASSED (Double Ratchet — per-message keys)".to_string());
        } else if session_established {
            tests_passed += 1;
            details.push("Perfect Forward Secrecy: PASSED (per-session ephemeral keys; peer lacks Double Ratchet)".to_string());
        } else {
            details.push("Perfect Forward Secrecy: FAILED (no established session)".to_string());
        }

        // Test 5: Replay Protection — provided by sequence-numbered AAD on the
        // static path and by key destruction in the ratchet. Random IVs alone
        // never provided it, so this no longer passes unconditionally.
        total_tests += 1;
        if ratchet_active {
            tests_passed += 1;
            details.push("Replay Protection: PASSED (ratchet message keys are destroyed on use)".to_string());
        } else if session_established {
            tests_passed += 1;
            details.push("Replay Protection: PASSED (sequence-numbered AAD on the session path)".to_string());
        } else {
            details.push("Replay Protection: FAILED (no established session)".to_string());
        }

        let level = if tests_passed == total_tests { 95 } else { (tests_passed * 100 / total_tests) as u8 };

        SecurityLevel {
            level,
            tests_passed: tests_passed as u8,
            total_tests: total_tests as u8,
            details,
        }
    }

    /// Get current key ID
    pub fn get_current_key_id(&self) -> Option<&String> {
        self.current_key_id.as_ref()
    }

    /// Get public key for key exchange
    pub fn get_public_key(&self, key_id: &str) -> Result<&String, String> {
        self.key_pairs.get(key_id)
            .map(|kp| &kp.public_key)
            .ok_or("Key not found".to_string())
    }

    /// Perform key exchange (simplified)
    pub fn perform_key_exchange(&mut self, _remote_public_key: &str) -> Result<String, String> {
        // This is a simplified version - in practice, you'd need to implement
        // the full ECDH key exchange protocol
        let new_key_id = self.generate_key_pair()?;
        Ok(new_key_id)
    }
}

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

    #[test]
    fn test_crypto_utils() {
        let mut crypto = CryptoUtils::new();
        
        // Test key generation
        let key_id = crypto.generate_key_pair().unwrap();
        assert!(!key_id.is_empty());
        
        // Test encryption/decryption
        let test_data = "Hello, World!";
        let encrypted = crypto.encrypt_data(test_data, &key_id).unwrap();
        let decrypted = crypto.decrypt_data(&encrypted).unwrap();
        assert_eq!(test_data, decrypted);
        
        // Test security level: with no session and no ratchet the PFS and
        // replay checks must FAIL — they used to pass unconditionally, which
        // is exactly the lie the web's 5.7.1 hardening removed.
        let security = crypto.calculate_security_level(false, false);
        assert!(security.level > 0);
        assert_eq!(security.tests_passed, 3);
        assert!(security.details.iter().any(|d| d.starts_with("Perfect Forward Secrecy: FAILED")));

        let live = crypto.calculate_security_level(true, true);
        assert_eq!(live.tests_passed, 5);
        assert!(live.details.iter().any(|d| d.contains("Double Ratchet")));
    }
}