"""
Tests for the security module.
"""
import os
import tempfile
import unittest
from cryptography.hazmat.primitives import serialization

from ..shared.security import SecurityManager

class TestSecurity(unittest.TestCase):
    """Test suite for the security module."""
    
    def setUp(self):
        """Set up test environment."""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.certs_dir = self.temp_dir.name
        self.security_manager = SecurityManager(self.certs_dir)
    
    def tearDown(self):
        """Clean up test environment."""
        self.temp_dir.cleanup()
    
    def test_rsa_key_generation(self):
        """Test RSA key generation."""
        private_key, public_key = self.security_manager.generate_rsa_keypair(key_size=2048)
        
        # Check that keys are of the correct type
        self.assertEqual(private_key.key_size, 2048)
        self.assertEqual(public_key.key_size, 2048)
    
    def test_key_save_load(self):
        """Test saving and loading keys."""
        private_key, public_key = self.security_manager.generate_rsa_keypair(key_size=2048)
        
        # Save keys
        self.security_manager.save_private_key(private_key, 'test_private.pem')
        self.security_manager.save_public_key(public_key, 'test_public.pem')
        
        # Check that files exist
        self.assertTrue(os.path.exists(os.path.join(self.certs_dir, 'test_private.pem')))
        self.assertTrue(os.path.exists(os.path.join(self.certs_dir, 'test_public.pem')))
        
        # Load keys
        loaded_private = self.security_manager.load_private_key('test_private.pem')
        loaded_public = self.security_manager.load_public_key('test_public.pem')
        
        # Check that loaded keys match original
        private_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        loaded_private_bytes = loaded_private.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        
        public_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        loaded_public_bytes = loaded_public.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        self.assertEqual(private_bytes, loaded_private_bytes)
        self.assertEqual(public_bytes, loaded_public_bytes)
    
    def test_rsa_encryption(self):
        """Test RSA encryption and decryption."""
        private_key, public_key = self.security_manager.generate_rsa_keypair(key_size=2048)
        
        # Test data
        test_data = b"This is a test message"
        
        # Encrypt
        encrypted = self.security_manager.encrypt_with_rsa(public_key, test_data)
        
        # Decrypt
        decrypted = self.security_manager.decrypt_with_rsa(private_key, encrypted)
        
        # Check that decrypted data matches original
        self.assertEqual(test_data, decrypted)
    
    def test_aes_encryption(self):
        """Test AES encryption and decryption."""
        # Generate key
        key = self.security_manager.generate_aes_key(key_size=256)
        
        # Test data
        test_data = b"This is a test message for AES encryption"
        
        # Encrypt
        encrypted = self.security_manager.encrypt_with_aes(key, test_data)
        
        # Decrypt
        decrypted = self.security_manager.decrypt_with_aes(key, encrypted)
        
        # Check that decrypted data matches original
        self.assertEqual(test_data, decrypted)
    
    def test_certificate_generation(self):
        """Test certificate generation and verification."""
        private_key, _ = self.security_manager.generate_rsa_keypair(key_size=2048)
        
        # Generate certificate
        domain_name = "test.example.domain"
        cert_path = self.security_manager.generate_certificate(private_key, domain_name)
        
        # Check that certificate file exists
        self.assertTrue(os.path.exists(cert_path))
        
        # Verify certificate
        valid, _ = self.security_manager.verify_certificate(cert_path, domain_name)
        self.assertTrue(valid)
        
        # Verify with wrong domain
        valid, _ = self.security_manager.verify_certificate(cert_path, "wrong.domain")
        self.assertFalse(valid)
    
    def test_pfs_session(self):
        """Test Perfect Forward Secrecy session."""
        # Initialize PFS session for client
        client_session = self.security_manager.pfs_init()
        
        # Initialize PFS session for server
        server_session = self.security_manager.pfs_init()
        
        # Get public keys
        client_public_key = client_session['public_key'].public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        server_public_key = server_session['public_key'].public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        # Exchange keys
        encrypted_secret = self.security_manager.pfs_exchange(client_session, server_public_key)
        
        # Receive on server side
        self.security_manager.pfs_receive(server_session, encrypted_secret)
        
        # Test message encryption/decryption
        test_message = b"This is a secure PFS message"
        
        # Client encrypts
        encrypted_message = self.security_manager.pfs_encrypt(client_session, test_message)
        
        # Server decrypts
        decrypted_message = self.security_manager.pfs_decrypt(server_session, encrypted_message)
        
        # Check that decrypted message matches original
        self.assertEqual(test_message, decrypted_message)