qssh 0.0.2-alpha

//! Comprehensive unit tests for QSSH components

use qssh::{
    crypto::{PqKeyExchange, SymmetricCrypto, SessionKeyDerivation, PqAlgorithm},
    transport::{Message, ClientHelloMessage, ServerHelloMessage, KeyExchangeMessage, AuthMessage, AuthMethod},
    QsshConfig, PortForward,
};
use tokio::net::{TcpListener, TcpStream};
use std::collections::HashMap;

#[cfg(test)]
mod crypto_tests {
    use super::*;
    
    #[test]
    fn test_symmetric_encryption_decryption() {
        let shared_secret = vec![0x42; 32];
        let crypto = SymmetricCrypto::from_shared_secret(&shared_secret).unwrap();
        
        let plaintext = b"Hello, Quantum World!";
        let (ciphertext, nonce) = crypto.encrypt(plaintext).unwrap();
        
        assert_ne!(ciphertext, plaintext);
        
        let decrypted = crypto.decrypt(&ciphertext, &nonce).unwrap();
        assert_eq!(decrypted, plaintext);
    }
    
    #[test]
    fn test_session_key_derivation() {
        let shared_secret = vec![0x11; 32];
        let client_random = vec![0x22; 32];
        let server_random = vec![0x33; 32];
        
        let keys = SessionKeyDerivation::derive_keys(&shared_secret, &client_random, &server_random).unwrap();
        
        // Verify all keys are different
        assert_ne!(keys.client_write_key, keys.server_write_key);
        assert_ne!(keys.client_write_key, keys.client_write_iv);
        assert_ne!(keys.server_write_key, keys.server_write_iv);
        
        // Verify key lengths
        assert_eq!(keys.client_write_key.len(), 32);
        assert_eq!(keys.server_write_key.len(), 32);
        assert_eq!(keys.client_write_iv.len(), 12);
        assert_eq!(keys.server_write_iv.len(), 12);
    }
    
    #[test]
    fn test_key_exchange_signature_verification() {
        let kex1 = PqKeyExchange::new().unwrap();
        let kex2 = PqKeyExchange::new().unwrap();
        
        // Create key shares
        let (share1, sig1) = kex1.create_key_share().unwrap();
        let (share2, sig2) = kex2.create_key_share().unwrap();
        
        // Verify shares
        let verified_share1 = kex2.process_key_share(&kex1.public_bytes(), &share1, &sig1).unwrap();
        let verified_share2 = kex1.process_key_share(&kex2.public_bytes(), &share2, &sig2).unwrap();
        
        assert_eq!(verified_share1, share1);
        assert_eq!(verified_share2, share2);
    }
    
    #[test]
    fn test_shared_secret_computation() {
        let kex1 = PqKeyExchange::new().unwrap();
        let kex2 = PqKeyExchange::new().unwrap();
        
        let client_random = vec![0xAA; 32];
        let server_random = vec![0xBB; 32];
        
        let (share1, _) = kex1.create_key_share().unwrap();
        let (share2, _) = kex2.create_key_share().unwrap();
        
        // Compute shared secrets from both sides
        let secret1 = kex1.compute_shared_secret(&share1, &share2, &client_random, &server_random);
        let secret2 = kex2.compute_shared_secret(&share2, &share1, &server_random, &client_random);
        
        // Both sides should derive the same secret
        assert_eq!(secret1, secret2);
    }
}

#[cfg(test)]
mod message_serialization_tests {
    use super::*;
    
    #[test]
    fn test_client_hello_serialization() {
        let msg = ClientHelloMessage {
            version: 1,
            random: vec![0x42; 32],
            kex_algorithms: vec![PqAlgorithm::Falcon512],
            sig_algorithms: vec![PqAlgorithm::SphincsPlus],
            ciphers: vec!["aes256-gcm".to_string()],
            qkd_capable: false,
            extensions: vec![],
        };
        
        let serialized = bincode::serialize(&Message::ClientHello(msg.clone())).unwrap();
        let deserialized: Message = bincode::deserialize(&serialized).unwrap();
        
        if let Message::ClientHello(hello) = deserialized {
            assert_eq!(hello.version, msg.version);
            assert_eq!(hello.random, msg.random);
            assert_eq!(hello.kex_algorithms, msg.kex_algorithms);
        } else {
            panic!("Wrong message type");
        }
    }
    
    #[test]
    fn test_auth_message_serialization() {
        let msg = AuthMessage {
            username: "testuser".to_string(),
            auth_method: AuthMethod::PublicKey {
                algorithm: PqAlgorithm::Falcon512,
                public_key: vec![0x99; 897],
            },
            signature: vec![0x88; 690],
            session_id: vec![0x77; 32],
        };
        
        let serialized = bincode::serialize(&Message::Auth(msg.clone())).unwrap();
        let deserialized: Message = bincode::deserialize(&serialized).unwrap();
        
        if let Message::Auth(auth) = deserialized {
            assert_eq!(auth.username, msg.username);
            assert_eq!(auth.session_id, msg.session_id);
            assert_eq!(auth.signature, msg.signature);
        } else {
            panic!("Wrong message type");
        }
    }
}

#[cfg(test)]
mod config_tests {
    use super::*;
    
    #[test]
    fn test_qssh_config_creation() {
        let config = QsshConfig {
            server: "localhost:22222".to_string(),
            username: "testuser".to_string(),
            port_forwards: vec![
                PortForward {
                    local_port: 8080,
                    remote_host: "localhost".to_string(),
                    remote_port: 80,
                }
            ],
            use_qkd: false,
            pq_algorithm: PqAlgorithm::Falcon512,
            key_rotation_interval: 3600,
        };
        
        assert_eq!(config.server, "localhost:22222");
        assert_eq!(config.username, "testuser");
        assert_eq!(config.port_forwards.len(), 1);
        assert!(!config.use_qkd);
    }
}

#[cfg(test)]
mod auth_tests {
    use super::*;
    use qssh::auth::{AuthorizedKeysManager, AuthorizedKey};
    use std::path::PathBuf;
    use tokio::fs;
    use tempfile::tempdir;
    
    #[tokio::test]
    async fn test_authorized_keys_loading() {
        let dir = tempdir().unwrap();
        let user_dir = dir.path().join("testuser").join(".ssh");
        fs::create_dir_all(&user_dir).await.unwrap();
        
        let auth_keys_path = user_dir.join("authorized_keys");
        let content = "qssh-falcon512 AAABBBCCC test@localhost\n";
        fs::write(&auth_keys_path, content).await.unwrap();
        
        let manager = AuthorizedKeysManager::new(dir.path().to_path_buf());
        let keys = manager.load_user_keys("testuser").await.unwrap();
        
        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0].algorithm, PqAlgorithm::Falcon512);
    }
    
    #[tokio::test]
    async fn test_public_key_verification() {
        let kex = PqKeyExchange::new().unwrap();
        let public_key = kex.public_bytes();
        
        let message = b"Test message for signing";
        let signature = kex.sign_falcon(message).unwrap();
        
        let valid = kex.verify_falcon(message, &signature, &public_key).unwrap();
        assert!(valid);
        
        // Test with wrong message
        let wrong_message = b"Different message";
        let invalid = kex.verify_falcon(wrong_message, &signature, &public_key).unwrap();
        assert!(!invalid);
    }
}

#[cfg(test)]
mod transport_tests {
    use super::*;
    use qssh::transport::{Transport, ChannelMessage, DisconnectMessage};
    
    #[tokio::test]
    async fn test_transport_message_framing() {
        // Create a pair of connected streams
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        
        let server_task = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            stream
        });
        
        let client_stream = TcpStream::connect(addr).await.unwrap();
        let server_stream = server_task.await.unwrap();
        
        // Create transports with encryption
        let shared_secret = vec![0x42; 32];
        let crypto1 = SymmetricCrypto::from_shared_secret(&shared_secret).unwrap();
        let crypto2 = SymmetricCrypto::from_shared_secret(&shared_secret).unwrap();
        
        let transport1 = Transport::new(client_stream, crypto1);
        let transport2 = Transport::new(server_stream, crypto2);
        
        // Send a message
        let msg = Message::Ping(12345);
        transport1.send_message(&msg).await.unwrap();
        
        // Receive and verify
        let received: Message = transport2.receive_message().await.unwrap();
        if let Message::Ping(nonce) = received {
            assert_eq!(nonce, 12345);
        } else {
            panic!("Wrong message type");
        }
    }
}

#[cfg(test)]
mod error_handling_tests {
    use super::*;
    use qssh::QsshError;
    
    #[test]
    fn test_crypto_error_handling() {
        // Test with invalid shared secret size
        let short_secret = vec![0x42; 16]; // Too short
        let result = SymmetricCrypto::from_shared_secret(&short_secret);
        assert!(result.is_err());
        
        if let Err(QsshError::Crypto(msg)) = result {
            assert!(msg.contains("too short"));
        } else {
            panic!("Wrong error type");
        }
    }
    
    #[test]
    fn test_invalid_signature_handling() {
        let kex = PqKeyExchange::new().unwrap();
        let public_key = kex.public_bytes();
        
        let message = b"Test message";
        let invalid_signature = vec![0xFF; 690]; // Invalid signature
        
        let result = kex.verify_falcon(message, &invalid_signature, &public_key);
        assert!(result.is_err() || !result.unwrap());
    }
}