# Secure handshake with certificate verification
from cryptography.hazmat.primitives.asymmetric import x25519
from cryptography.x509 import load_pem_x509_certificate

def verify_worker_identity(worker_cert_pem, trusted_ca_pem):
    """Verify worker identity against CA."""
    cert = load_pem_x509_certificate(trusted_ca_pem)
    # Verify signature chain
    cert.verify_signature()  # Check CA signature
    return True

# Replay protection in controller_secure.py
class ReplayProtectedAEAD(AEAD):
    def __init__(self, key: bytes, expected_sender_id: str, 
                 nonce_expiry_seconds: int = 3600):
        super().__init__(key)
        self.seen_nonces: Dict[str, float] = {}
        self.current_time = time.time()
    
    def is_nonce_fresh(self, nonce: bytes) -> bool:
        """Check if nonce hasn't been used recently."""
        nonce_str = nonce.hex()
        if nonce_str in self.seen_nonces:
            last_seen = self.seen_nonces[nonce_str]
            if self.current_time - last_seen < 3600:  # 1 hour window
                return False
        self.seen_nonces[nonce_str] = time.time()
        return True
    
    def encrypt(self, plaintext: bytes, aad: bytes = b''):
        nonce = os.urandom(12)
        
        # Check for replay before encryption
        if not self.is_nonce_fresh(nonce):
            raise ReplayError(f"Nonce {nonce.hex()} is stale")
        
        nonce, ct = super().encrypt(plaintext, aad)
        return nonce, ct
    
    def decrypt(self, nonce: bytes, ciphertext: bytes, aad: bytes = b''):
        # Check nonce freshness before decryption
        if not self.is_nonce_fresh(nonce):
            raise ReplayError(f"Nonce {nonce.hex()} is stale")
        
        return super().decrypt(nonce, ciphertext, aad)

# Encrypted weight storage
class SecureWeightStorage:
    def __init__(self, encryption_key_path: str):
        self.key = load_encryption_key(encryption_key_path)
    
    def store_weights(self, slice_id: str, weights: np.ndarray):
        """Store encrypted weights."""
        encrypted = ChaCha20Poly1305.encrypt(
            nonce=os.urandom(12),
            plaintext=pickle.dumps(weights),
            key=self.key
        )
        # Store with encryption metadata
        self._store_encrypted(slice_id, encrypted)
    
    def load_weights(self, slice_id: str):
        """Load and decrypt weights."""
        encrypted = self._load_encrypted(slice_id)
        decrypted = ChaCha20Poly1305.decrypt(
            nonce=encrypted[:12],
            ciphertext=encrypted[12:],
            key=self.key
        )
        return pickle.loads(decrypted)

# key_management.py
from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PrivateKey, Ed448PublicKey
from cryptography.hazmat.primitives.serialization import load_pem_private_key
import os

class SecureKeyManager:
    """Manages cryptographic keys with proper lifecycle."""
    
    def __init__(self, key_directory: str):
        self.key_directory = key_directory
        self.key_store = {}  # key_id -> (key_type, key_data)
    
    def load_private_key(self, key_path: str, password: bytes = None):
        """Load private key from PEM file."""
        with open(key_path, 'rb') as f:
            pem_data = f.read()
        
        if password:
            key = load_pem_private_key(pem_data, password=password)
        else:
            key = load_pem_private_key(pem_data, password=None)
        
        return key
    
    def generate_hybrid_keypair(self):
        """Generate hybrid classical + PQC keypair."""
        # Classical ECDH
        classical_priv = Ed448PrivateKey.generate()
        classical_pub = classical_priv.public_key()
        
        # PQC KEM (Kyber768)
        if OQS_AVAILABLE:
            kem = oqs.KeyEncapsulation("ML-KEM-768")
            pq_priv, pq_pub = kem.generate_keypair()
            
            return {
                'classical_private': classical_priv,
                'classical_public': classical_pub,
                'pq_private': pq_priv,
                'pq_public': pq_pub
            }
        
        raise RuntimeError("OQS not available for hybrid key generation")
    
    def rotate_key(self, old_key_id: str, new_key_path: str):
        """Rotate cryptographic key."""
        old_key = self.key_store.pop(old_key_id)
        
        # Generate new key
        new_keypair = self.generate_hybrid_keypair()
        
        # Migrate active sessions to new key
        self._migrate_sessions(old_key_id, new_keypair)
        
        # Store new key
        self.key_store[new_key_id] = new_keypair
    
    def _migrate_sessions(self, old_key_id: str, new_keypair):
        """Migrate active sessions from old to new key."""
        # Implement session migration logic
        pass

# input_validation.py
import re
from typing import Any, Dict

class SecureRequestValidator:
    """Validate and sanitize all incoming requests."""
    
    MAX_INPUT_SIZE = 10 * 1024 * 1024  # 10MB
    MAX_SESSION_ID_LENGTH = 36  # UUID length
    
    def validate_slice_id(self, slice_id: str) -> bool:
        """Validate slice ID format."""
        pattern = r'^slice_\d+_\d+$'
        return bool(re.match(pattern, slice_id))
    
    def validate_session_id(self, session_id: str) -> bool:
        """Validate session ID is a proper UUID."""
        import uuid
        try:
            uuid.UUID(session_id)
            return True
        except ValueError:
            return False
    
    def validate_input_size(self, data: bytes) -> tuple[bool, str]:
        """Check input size limits."""
        if len(data) > self.MAX_INPUT_SIZE:
            return False, f"Input exceeds {self.MAX_INPUT_SIZE} byte limit"
        return True, ""
    
    def sanitize_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Remove dangerous fields and validate types."""
        sanitized = {}
        
        for key, value in payload.items():
            # Reject control characters
            if isinstance(value, str) and any(c in value for c in ['<', '>', '"', "'", '&']):
                continue
            
            # Validate numeric fields
            if key in ['slice_id', 'manifest']:
                sanitized[key] = self._validate_literal(key, value)
            
        return sanitized
    
    def _validate_literal(self, field_name: str, value: Any) -> Any:
        """Validate a single literal field."""
        if isinstance(value, dict):
            # Validate nested dictionaries
            validated_dict = {}
            for k, v in value.items():
                if k.startswith('_'):  # Reject private fields
                    continue
                validated_dict[k] = v
            return validated_dict
        elif isinstance(value, (int, float)):
            if value < 0:
                raise ValueError(f"Negative value not allowed for {field_name}")
            return value
        return value

# incident_response/replay_attack_handler.py
import logging
from datetime import datetime

class ReplayAttackHandler:
    """Handle detected replay attack incidents."""
    
    def __init__(self, alert_endpoint: str):
        self.alert_endpoint = alert_endpoint
        self.incident_id = None
    
    def detect_replay(self, nonce: bytes, sender_id: str) -> bool:
        """Check for replay attack."""
        # Check if nonce was seen recently
        from prototype.crypto import ReplayProtectedAEAD
        
        aead = ReplayProtectedAEAD(
            key=self._get_key(sender_id),
            expected_sender_id=sender_id
        )
        
        try:
            aead.decrypt(nonce, ciphertext)
            logging.warning(f"Potential replay attack detected from {sender_id}")
            
            # Trigger incident response
            self.report_incident(
                severity="P2",
                type="REPLAY_ATTACK",
                source=sender_id,
                nonce=nonce.hex()
            )
            return True
        except Exception:
            return False
    
    def report_incident(self, severity: str, incident_type: str, **kwargs):
        """Report incident to monitoring system."""
        payload = {
            "severity": severity,
            "type": incident_type,
            "timestamp": datetime.utcnow().isoformat(),
            "details": kwargs
        }
        
        import requests
        response = requests.post(
            self.alert_endpoint,
            json=payload,
            headers={"Content-Type": "application/json"}
        )
        
        return response.json()