import unittest
from unittest.mock import MagicMock
from security.audit import SecureAudit
from security.rbac_engine import RBACEngine
import os

class TestSecureAudit(unittest.TestCase):
    def setUp(self):
        # Setup mock RBAC engine
        self.mock_rbac = MagicMock(spec=RBACEngine)
        
        # Generate test encryption key
        self.test_key = os.urandom(32)
        
        # Initialize SecureAudit with in-memory DB
        self.audit = SecureAudit(self.mock_rbac, ":memory:", self.test_key)

    def test_aes_encryption(self):
        """Test AES-256 encryption of sensitive data"""
        test_data = "test_cron_expression"
        encrypted = self.audit._encrypt_data(test_data)
        decrypted = self.audit._decrypt_data(encrypted)
        self.assertEqual(test_data, decrypted)

    def test_hmac_obfuscation(self):
        """Test HMAC-SHA256 obfuscation of task IDs"""
        task_id = "task_123"
        obfuscated = self.audit._obfuscate_id(task_id)
        self.assertEqual(len(obfuscated), 64)  # SHA256 hexdigest length
        self.assertNotEqual(task_id, obfuscated)

    def test_timestamp_integrity(self):
        """Verify timestamp integrity protection"""
        entry = {"operation": "test", "user": "admin"}
        hash_val = self.audit.log(entry)
        
        # Tamper with timestamp and verify detection
        with self.assertRaises(ValueError):
            self.audit._verify_integrity(hash_val, "tampered_timestamp")

    def test_log_retrieval(self):
        """Test encrypted log storage and retrieval"""
        entry = {
            "operation": "test",
            "user": "admin",
            "cron": "* * * * *",  # Sensitive data
            "task_id": "sensitive_task_123"
        }
        hash_val = self.audit.log(entry)

    def test_tls_handshake_logging(self):
        """Test TLS handshake parameter logging"""
        cert_info = {
            "subject": {"CN": "test.example.com"},
            "cert_chain": [
                {
                    "subject": "CN=test.example.com",
                    "issuer": "CN=Test CA",
                    "serial": "12345",
                    "valid_from": "2025-01-01",
                    "valid_to": "2026-01-01",
                    "key_algorithm": "RSA",
                    "key_size": 2048
                }
            ]
        }
        
        tls_params = {
            "protocol": "TLSv1.3",
            "cipher": "TLS_AES_256_GCM_SHA384",
            "key_exchange": "ECDHE",
            "authentication": "RSA",
            "encryption": "AES-256-GCM",
            "mac": "SHA384",
            "forward_secrecy": True,
            "session_resumed": False,
            "session_id": "abc123",
            "session_ticket": None,
            "ocsp_stapling": True,
            "sct_validation": True,
            "extensions": [
                {"type": "server_name", "data": "test.example.com"},
                {"type": "key_share", "data": "x25519"}
            ],
            "alpn_protocol": "h2"
        }
        
        # Log the TLS handshake
        self.audit.log_tls_handshake(cert_info, tls_params)
        
        # Verify the log entry was created
        logs = self.audit.get_logs(limit=1)
        self.assertEqual(len(logs), 1)
        self.assertEqual(logs[0]['event'], 'tls_handshake')
        self.assertEqual(logs[0]['client'], 'test.example.com')
        self.assertEqual(logs[0]['protocol'], 'TLSv1.3')
        self.assertTrue(logs[0]['cipher_suite']['forward_secrecy'])
        self.assertTrue(logs[0]['security_indicators']['ocsp_stapling'])
        
        logs = self.audit.get_logs()
        self.assertEqual(len(logs), 1)
        self.assertEqual(logs[0]["integrity_hash"], hash_val)
        
        # Verify sensitive data is encrypted
        self.assertTrue(logs[0]["encrypted_cron"].startswith("gAAAA"))
        self.assertEqual(len(logs[0]["obfuscated_task_id"]), 64)

if __name__ == "__main__":
    unittest.main()