# P2P Network Communication - Security Module Tests
"""
安全模块测试
测试传输加密、本地数据加密和密钥管理功能

需求: 10.1, 10.2, 10.3
"""

import os
import tempfile
import pytest
from pathlib import Path

from shared.security import (
    AESCipher,
    EncryptedData,
    MessageEncryptor,
    FileEncryptor,
    KeyManager,
    LocalDataEncryptor,
    TLSManager,
    EncryptionError,
    DecryptionError,
    create_message_encryptor,
    create_file_encryptor,
    create_local_data_encryptor,
    encrypt_message,
    decrypt_message,
)


class TestAESCipher:
    """AES-256-GCM 加密器测试"""
    
    def test_generate_key(self):
        """测试密钥生成"""
        key = AESCipher.generate_key()
        assert len(key) == 32  # 256 bits
        
        # 确保每次生成的密钥不同
        key2 = AESCipher.generate_key()
        assert key != key2
    
    def test_generate_iv(self):
        """测试IV生成"""
        iv = AESCipher.generate_iv()
        assert len(iv) == 12  # 96 bits
    
    def test_encrypt_decrypt(self):
        """测试加密解密往返"""
        cipher = AESCipher()
        plaintext = b"Hello, World! This is a test message."
        
        encrypted = cipher.encrypt(plaintext)
        decrypted = cipher.decrypt(encrypted)
        
        assert decrypted == plaintext
    
    def test_encrypt_decrypt_empty(self):
        """测试空数据加密解密"""
        cipher = AESCipher()
        plaintext = b""
        
        encrypted = cipher.encrypt(plaintext)
        decrypted = cipher.decrypt(encrypted)
        
        assert decrypted == plaintext
    
    def test_encrypt_decrypt_large_data(self):
        """测试大数据加密解密"""
        cipher = AESCipher()
        plaintext = os.urandom(1024 * 1024)  # 1MB
        
        encrypted = cipher.encrypt(plaintext)
        decrypted = cipher.decrypt(encrypted)
        
        assert decrypted == plaintext
    
    def test_encrypt_decrypt_unicode(self):
        """测试Unicode数据加密解密"""
        cipher = AESCipher()
        plaintext = "你好，世界！这是一条测试消息。🎉".encode('utf-8')
        
        encrypted = cipher.encrypt(plaintext)
        decrypted = cipher.decrypt(encrypted)
        
        assert decrypted == plaintext
    
    def test_different_keys_produce_different_ciphertext(self):
        """测试不同密钥产生不同密文"""
        cipher1 = AESCipher()
        cipher2 = AESCipher()
        plaintext = b"Same message"
        
        encrypted1 = cipher1.encrypt(plaintext)
        encrypted2 = cipher2.encrypt(plaintext)
        
        assert encrypted1.ciphertext != encrypted2.ciphertext
    
    def test_wrong_key_fails_decryption(self):
        """测试错误密钥解密失败"""
        cipher1 = AESCipher()
        cipher2 = AESCipher()
        plaintext = b"Secret message"
        
        encrypted = cipher1.encrypt(plaintext)
        
        with pytest.raises(DecryptionError):
            cipher2.decrypt(encrypted)
    
    def test_derive_key_from_password(self):
        """测试从密码派生密钥"""
        password = "my_secure_password"
        
        key1, salt1 = AESCipher.derive_key_from_password(password)
        assert len(key1) == 32
        assert len(salt1) == 16
        
        # 相同密码和盐应产生相同密钥
        key2, _ = AESCipher.derive_key_from_password(password, salt1)
        assert key1 == key2
        
        # 不同盐应产生不同密钥
        key3, salt3 = AESCipher.derive_key_from_password(password)
        assert key1 != key3
    
    def test_encrypt_with_password(self):
        """测试使用密码加密"""
        cipher = AESCipher()
        password = "test_password"
        plaintext = b"Secret data"
        
        encrypted = cipher.encrypt_with_password(plaintext, password)
        assert encrypted.salt  # 应包含盐值
        
        decrypted = cipher.decrypt_with_password(encrypted, password)
        assert decrypted == plaintext
    
    def test_wrong_password_fails(self):
        """测试错误密码解密失败"""
        cipher = AESCipher()
        plaintext = b"Secret data"
        
        encrypted = cipher.encrypt_with_password(plaintext, "correct_password")
        
        with pytest.raises(DecryptionError):
            cipher.decrypt_with_password(encrypted, "wrong_password")


class TestEncryptedData:
    """加密数据结构测试"""
    
    def test_to_bytes_from_bytes(self):
        """测试序列化和反序列化"""
        original = EncryptedData(
            ciphertext=b"encrypted_content",
            iv=b"123456789012",
            tag=b"1234567890123456",
            salt=b"salt_value_here!"
        )
        
        serialized = original.to_bytes()
        restored = EncryptedData.from_bytes(serialized)
        
        assert restored.ciphertext == original.ciphertext
        assert restored.iv == original.iv
        assert restored.tag == original.tag
        assert restored.salt == original.salt
    
    def test_to_base64_from_base64(self):
        """测试Base64编码和解码"""
        original = EncryptedData(
            ciphertext=b"encrypted_content",
            iv=b"123456789012",
            tag=b"1234567890123456"
        )
        
        base64_str = original.to_base64()
        restored = EncryptedData.from_base64(base64_str)
        
        assert restored.ciphertext == original.ciphertext
        assert restored.iv == original.iv
        assert restored.tag == original.tag


class TestMessageEncryptor:
    """消息加密器测试"""
    
    def test_encrypt_decrypt_message(self):
        """测试消息加密解密"""
        encryptor = MessageEncryptor()
        message_data = b'{"type": "text", "content": "Hello!"}'
        
        encrypted = encryptor.encrypt_message(message_data)
        decrypted = encryptor.decrypt_message(encrypted)
        
        assert decrypted == message_data
    
    def test_set_key(self):
        """测试设置密钥"""
        key = AESCipher.generate_key()
        encryptor = MessageEncryptor()
        encryptor.set_key(key)
        
        assert encryptor.key == key
    
    def test_shared_key_encryption(self):
        """测试共享密钥加密"""
        key = AESCipher.generate_key()
        
        encryptor1 = MessageEncryptor(AESCipher(key))
        encryptor2 = MessageEncryptor(AESCipher(key))
        
        message = b"Shared secret message"
        
        encrypted = encryptor1.encrypt_message(message)
        decrypted = encryptor2.decrypt_message(encrypted)
        
        assert decrypted == message


class TestFileEncryptor:
    """文件加密器测试"""
    
    def test_encrypt_decrypt_file(self):
        """测试文件加密解密"""
        encryptor = FileEncryptor()
        
        with tempfile.TemporaryDirectory() as tmpdir:
            # 创建测试文件
            input_file = Path(tmpdir) / "test.txt"
            encrypted_file = Path(tmpdir) / "test.enc"
            output_file = Path(tmpdir) / "test_decrypted.txt"
            
            original_content = b"This is test file content.\n" * 100
            input_file.write_bytes(original_content)
            
            # 加密
            encryptor.encrypt_file(str(input_file), str(encrypted_file))
            assert encrypted_file.exists()
            
            # 解密
            encryptor.decrypt_file(str(encrypted_file), str(output_file))
            assert output_file.exists()
            
            # 验证内容
            decrypted_content = output_file.read_bytes()
            assert decrypted_content == original_content
    
    def test_encrypt_decrypt_chunk(self):
        """测试文件块加密解密"""
        encryptor = FileEncryptor()
        chunk_data = os.urandom(64 * 1024)  # 64KB
        
        encrypted = encryptor.encrypt_chunk(chunk_data)
        decrypted = encryptor.decrypt_chunk(encrypted)
        
        assert decrypted == chunk_data


class TestKeyManager:
    """密钥管理器测试"""
    
    def test_generate_key_pair(self):
        """测试生成密钥对"""
        with tempfile.TemporaryDirectory() as tmpdir:
            manager = KeyManager(tmpdir)
            private_key, public_key = manager.generate_key_pair()
            
            assert b"PRIVATE KEY" in private_key
            assert b"PUBLIC KEY" in public_key
    
    def test_save_load_key(self):
        """测试保存和加载密钥"""
        with tempfile.TemporaryDirectory() as tmpdir:
            manager = KeyManager(tmpdir)
            key_data = AESCipher.generate_key()
            
            # 保存
            assert manager.save_key("test_key", key_data)
            
            # 加载
            loaded = manager.load_key("test_key")
            assert loaded == key_data
    
    def test_save_load_key_with_password(self):
        """测试使用密码保存和加载密钥"""
        with tempfile.TemporaryDirectory() as tmpdir:
            manager = KeyManager(tmpdir)
            key_data = AESCipher.generate_key()
            password = "secure_password"
            
            # 保存
            assert manager.save_key("protected_key", key_data, password)
            
            # 加载
            loaded = manager.load_key("protected_key", password)
            assert loaded == key_data
    
    def test_delete_key(self):
        """测试删除密钥"""
        with tempfile.TemporaryDirectory() as tmpdir:
            manager = KeyManager(tmpdir)
            key_data = AESCipher.generate_key()
            
            manager.save_key("to_delete", key_data)
            assert manager.delete_key("to_delete")
            
            # 验证已删除
            assert manager.load_key("to_delete") is None
    
    def test_session_key_management(self):
        """测试会话密钥管理"""
        with tempfile.TemporaryDirectory() as tmpdir:
            manager = KeyManager(tmpdir)
            
            # 获取或创建会话密钥
            key1 = manager.get_or_create_session_key("session1")
            assert len(key1) == 32
            
            # 再次获取应返回相同密钥
            key2 = manager.get_or_create_session_key("session1")
            assert key1 == key2
            
            # 不同会话应有不同密钥
            key3 = manager.get_or_create_session_key("session2")
            assert key1 != key3
            
            # 清除会话密钥
            manager.clear_session_key("session1")
            key4 = manager.get_or_create_session_key("session1")
            assert key1 != key4  # 应该是新密钥


class TestLocalDataEncryptor:
    """本地数据加密器测试"""
    
    def test_encrypt_decrypt_string(self):
        """测试字符串加密解密"""
        encryptor = LocalDataEncryptor(password="test_password")
        original = "Hello, World!"
        
        encrypted = encryptor.encrypt_data(original)
        decrypted = encryptor.decrypt_data(encrypted)
        
        assert decrypted.decode('utf-8') == original
    
    def test_encrypt_decrypt_dict(self):
        """测试字典加密解密"""
        encryptor = LocalDataEncryptor(password="test_password")
        original = {"name": "Test", "value": 123, "nested": {"key": "value"}}
        
        encrypted = encryptor.encrypt_data(original)
        decrypted = encryptor.decrypt_data(encrypted, as_json=True)
        
        assert decrypted == original
    
    def test_encrypt_decrypt_chat_history(self):
        """测试聊天记录加密解密"""
        encryptor = LocalDataEncryptor(password="chat_password")
        messages = [
            {"sender": "user1", "content": "Hello", "timestamp": 1234567890},
            {"sender": "user2", "content": "Hi there!", "timestamp": 1234567891},
        ]
        
        encrypted = encryptor.encrypt_chat_history(messages)
        decrypted = encryptor.decrypt_chat_history(encrypted)
        
        assert decrypted == messages
    
    def test_save_load_encrypted_file(self):
        """测试保存和加载加密文件"""
        with tempfile.TemporaryDirectory() as tmpdir:
            encryptor = LocalDataEncryptor(password="file_password")
            file_path = Path(tmpdir) / "encrypted_data.enc"
            
            original_data = {"key": "value", "list": [1, 2, 3]}
            
            # 保存
            assert encryptor.save_encrypted_file(original_data, str(file_path))
            assert file_path.exists()
            
            # 加载
            loaded = encryptor.load_encrypted_file(str(file_path), as_json=True)
            assert loaded == original_data


class TestConvenienceFunctions:
    """便捷函数测试"""
    
    def test_create_message_encryptor(self):
        """测试创建消息加密器"""
        encryptor = create_message_encryptor()
        assert encryptor is not None
        assert len(encryptor.key) == 32
    
    def test_create_file_encryptor(self):
        """测试创建文件加密器"""
        encryptor = create_file_encryptor()
        assert encryptor is not None
        assert len(encryptor.key) == 32
    
    def test_create_local_data_encryptor(self):
        """测试创建本地数据加密器"""
        encryptor = create_local_data_encryptor("password")
        assert encryptor is not None
    
    def test_encrypt_decrypt_message_functions(self):
        """测试快速加密解密函数"""
        key = AESCipher.generate_key()
        message = b"Quick encryption test"
        
        encrypted = encrypt_message(message, key)
        decrypted = decrypt_message(encrypted, key)
        
        assert decrypted == message


class TestTLSManager:
    """TLS管理器测试"""
    
    def test_create_client_ssl_context_no_verify(self):
        """测试创建客户端SSL上下文（不验证）"""
        manager = TLSManager()
        context = manager.create_client_ssl_context(verify=False)
        
        assert context is not None
        assert context.verify_mode.name == "CERT_NONE"
    
    def test_create_client_ssl_context_with_verify(self):
        """测试创建客户端SSL上下文（验证）"""
        manager = TLSManager()
        context = manager.create_client_ssl_context(verify=True)
        
        assert context is not None
        assert context.verify_mode.name == "CERT_REQUIRED"