# Secure handshake with certificate verification (example)
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 crypto_improved.py
class ReplayProtectedAEAD(AEAD):
    def __init__(self, key: bytes, nonce_expiry_seconds: int = 3600):
        super().__init__(key)
        self.seen_nonces: Set[str] = set()
        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()
        with self.lock:
            # Clean up stale nonces first
            self._cleanup_stale_nonces()
            
            if nonce_str in self.seen_nonces:
                return False
            self.seen_nonces.add(nonce_str)
            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 RuntimeError(f"Nonce collision detected - possible replay attack")
        
        nonce, ct = super().encrypt(nonce, 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 RuntimeError(f"Nonce {nonce.hex()} is stale - possible replay attack")
        
        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)

# Safe binary serialization
class WeightSlice:
    def to_bytes(self) -> bytes:
        """Serialize weights to binary format (no pickle)."""
        packed = []
        for w, b in self.weights:
            packed.append(w.tobytes())
            packed.append(b.tobytes())
        return b'\x00'.join(packed)

# 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")

# 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, ""

# 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_improved 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()