chorus-services/modules/shhh/integrations/bzzz_interceptor.py

"""
BZZZ Message Interceptor for SHHH Secrets Sentinel
Intercepts and validates BZZZ P2P messages before network propagation.
"""

import asyncio
import json
import time
from typing import Dict, Any, Optional, Set, Callable
from dataclasses import dataclass
from datetime import datetime
import structlog

from ..core.hypercore_reader import BzzzMessage
from ..core.detector import SecretDetector, DetectionResult
from ..core.quarantine import QuarantineManager

logger = structlog.get_logger()


@dataclass
class BlockedMessage:
    """Represents a blocked BZZZ message"""
    message_id: str
    sender_agent: str
    block_reason: str
    secret_types: list
    timestamp: datetime
    quarantine_id: Optional[int] = None


class BzzzInterceptor:
    """
    Intercepts BZZZ P2P messages before transmission to prevent secret leakage.
    Integrates with the BZZZ network layer to scan messages in real-time.
    """
    
    def __init__(
        self, 
        detector: SecretDetector,
        quarantine_manager: QuarantineManager,
        bzzz_config: Dict[str, Any] = None
    ):
        self.detector = detector
        self.quarantine = quarantine_manager
        self.bzzz_config = bzzz_config or {}
        
        # Message blocking state
        self.blocked_messages: Dict[str, BlockedMessage] = {}
        self.message_hooks: Set[Callable] = set()
        self.is_active = False
        
        # Statistics
        self.stats = {
            'total_scanned': 0,
            'secrets_detected': 0,
            'messages_blocked': 0,
            'false_positives': 0,
            'last_reset': datetime.now()
        }
        
        logger.info("Initialized BzzzInterceptor")
    
    async def start(self):
        """Start the BZZZ message interception service"""
        self.is_active = True
        logger.info("BZZZ Interceptor started - all outgoing messages will be scanned")
    
    async def stop(self):
        """Stop the BZZZ message interception service"""
        self.is_active = False
        logger.info("BZZZ Interceptor stopped")
    
    def install_message_hook(self, hook_function: Callable):
        """Install a message hook for BZZZ network integration"""
        self.message_hooks.add(hook_function)
        logger.info(f"Installed BZZZ message hook: {hook_function.__name__}")
    
    def remove_message_hook(self, hook_function: Callable):
        """Remove a message hook"""
        self.message_hooks.discard(hook_function)
        logger.info(f"Removed BZZZ message hook: {hook_function.__name__}")
    
    async def intercept_outgoing_message(self, message: BzzzMessage) -> bool:
        """
        Intercept and scan an outgoing BZZZ message.
        Returns True if message should be allowed, False if blocked.
        """
        if not self.is_active:
            return True  # Pass through if interceptor is inactive
        
        start_time = time.time()
        self.stats['total_scanned'] += 1
        
        try:
            # Scan message for secrets
            detection_result = self.detector.scan_bzzz_message(message)
            
            if detection_result.has_secrets:
                await self._handle_secret_detection(message, detection_result)
                return False  # Block message
            
            # Message is clean, allow transmission
            processing_time = (time.time() - start_time) * 1000
            logger.debug(
                f"BZZZ message scanned clean",
                message_id=message.message_id,
                sender=message.sender_agent,
                processing_time_ms=processing_time
            )
            
            return True
            
        except Exception as e:
            logger.error(f"Error intercepting BZZZ message: {e}")
            # On error, default to blocking for security
            await self._block_message_on_error(message, str(e))
            return False
    
    async def _handle_secret_detection(self, message: BzzzMessage, detection_result: DetectionResult):
        """Handle detection of secrets in a BZZZ message"""
        self.stats['secrets_detected'] += 1
        self.stats['messages_blocked'] += 1
        
        # Extract secret types for blocking record
        secret_types = [match.secret_type for match in detection_result.matches]
        
        # Quarantine the detection result
        quarantine_entry = await self.quarantine.quarantine_detection(detection_result)
        
        # Create blocked message record
        blocked_msg = BlockedMessage(
            message_id=message.message_id,
            sender_agent=message.sender_agent,
            block_reason=f"Secrets detected: {', '.join(secret_types)}",
            secret_types=secret_types,
            timestamp=datetime.now(),
            quarantine_id=quarantine_entry.id
        )
        
        self.blocked_messages[message.message_id] = blocked_msg
        
        # Notify BZZZ network layer
        await self._notify_message_blocked(message, blocked_msg)
        
        logger.critical(
            f"BLOCKED BZZZ message containing secrets",
            message_id=message.message_id,
            sender=message.sender_agent,
            recipient=message.recipient_agent,
            secret_types=secret_types,
            severity=detection_result.max_severity,
            quarantine_id=quarantine_entry.id
        )
    
    async def _block_message_on_error(self, message: BzzzMessage, error_msg: str):
        """Block a message due to processing error"""
        self.stats['messages_blocked'] += 1
        
        blocked_msg = BlockedMessage(
            message_id=message.message_id,
            sender_agent=message.sender_agent,
            block_reason=f"Processing error: {error_msg}",
            secret_types=[],
            timestamp=datetime.now()
        )
        
        self.blocked_messages[message.message_id] = blocked_msg
        await self._notify_message_blocked(message, blocked_msg)
        
        logger.error(
            f"BLOCKED BZZZ message due to error",
            message_id=message.message_id,
            sender=message.sender_agent,
            error=error_msg
        )
    
    async def _notify_message_blocked(self, message: BzzzMessage, blocked_msg: BlockedMessage):
        """Notify BZZZ network and sender about blocked message"""
        notification = {
            'event': 'message_blocked',
            'message_id': message.message_id,
            'sender_agent': message.sender_agent,
            'recipient_agent': message.recipient_agent,
            'block_reason': blocked_msg.block_reason,
            'secret_types': blocked_msg.secret_types,
            'timestamp': blocked_msg.timestamp.isoformat(),
            'quarantine_id': blocked_msg.quarantine_id
        }
        
        # Notify all registered hooks
        for hook in self.message_hooks:
            try:
                await self._call_hook_safely(hook, 'message_blocked', notification)
            except Exception as e:
                logger.warning(f"Hook {hook.__name__} failed: {e}")
        
        # Send notification back to sender agent
        await self._notify_sender_agent(message.sender_agent, notification)
    
    async def _call_hook_safely(self, hook: Callable, event_type: str, data: Dict[str, Any]):
        """Safely call a hook function with error handling"""
        try:
            if asyncio.iscoroutinefunction(hook):
                await hook(event_type, data)
            else:
                hook(event_type, data)
        except Exception as e:
            logger.warning(f"Hook {hook.__name__} failed: {e}")
    
    async def _notify_sender_agent(self, sender_agent: str, notification: Dict[str, Any]):
        """Send notification to the sender agent about blocked message"""
        try:
            # This would integrate with the BZZZ network's agent communication system
            # For now, we'll log the notification
            logger.info(
                f"Notifying agent about blocked message",
                agent=sender_agent,
                message_id=notification['message_id'],
                reason=notification['block_reason']
            )
            
            # TODO: Implement actual agent notification via BZZZ network
            # This might involve:
            # - Sending a system message back to the agent
            # - Updating agent's message status
            # - Triggering agent's error handling workflow
            
        except Exception as e:
            logger.error(f"Failed to notify sender agent {sender_agent}: {e}")
    
    def is_message_blocked(self, message_id: str) -> Optional[BlockedMessage]:
        """Check if a message is blocked"""
        return self.blocked_messages.get(message_id)
    
    def unblock_message(self, message_id: str, reviewer: str, reason: str) -> bool:
        """Unblock a previously blocked message (for false positives)"""
        if message_id not in self.blocked_messages:
            return False
        
        blocked_msg = self.blocked_messages[message_id]
        
        # Mark as false positive in stats
        self.stats['false_positives'] += 1
        
        # Remove from blocked messages
        del self.blocked_messages[message_id]
        
        logger.info(
            f"Unblocked BZZZ message",
            message_id=message_id,
            reviewer=reviewer,
            reason=reason,
            original_block_reason=blocked_msg.block_reason
        )
        
        return True
    
    def get_blocked_messages(self, limit: int = 100) -> list[BlockedMessage]:
        """Get list of recently blocked messages"""
        blocked_list = list(self.blocked_messages.values())
        blocked_list.sort(key=lambda x: x.timestamp, reverse=True)
        return blocked_list[:limit]
    
    def get_stats(self) -> Dict[str, Any]:
        """Get interceptor statistics"""
        current_time = datetime.now()
        uptime_hours = (current_time - self.stats['last_reset']).total_seconds() / 3600
        
        stats = self.stats.copy()
        stats.update({
            'uptime_hours': round(uptime_hours, 2),
            'is_active': self.is_active,
            'blocked_messages_count': len(self.blocked_messages),
            'detection_rate': (
                self.stats['secrets_detected'] / max(1, self.stats['total_scanned'])
            ) * 100,
            'false_positive_rate': (
                self.stats['false_positives'] / max(1, self.stats['secrets_detected'])
            ) * 100 if self.stats['secrets_detected'] > 0 else 0
        })
        
        return stats
    
    def reset_stats(self):
        """Reset statistics counters"""
        self.stats = {
            'total_scanned': 0,
            'secrets_detected': 0,
            'messages_blocked': 0,
            'false_positives': 0,
            'last_reset': datetime.now()
        }
        
        logger.info("BZZZ Interceptor statistics reset")
    
    async def cleanup_old_blocked_messages(self, hours: int = 24):
        """Clean up old blocked message records"""
        cutoff_time = datetime.now() - timedelta(hours=hours)
        
        old_messages = [
            msg_id for msg_id, blocked_msg in self.blocked_messages.items()
            if blocked_msg.timestamp < cutoff_time
        ]
        
        for msg_id in old_messages:
            del self.blocked_messages[msg_id]
        
        if old_messages:
            logger.info(f"Cleaned up {len(old_messages)} old blocked message records")
        
        return len(old_messages)


class BzzzNetworkAdapter:
    """
    Adapter to integrate BzzzInterceptor with the actual BZZZ network layer.
    This would be customized based on the BZZZ implementation details.
    """
    
    def __init__(self, interceptor: BzzzInterceptor):
        self.interceptor = interceptor
        self.original_send_function = None
        
    def install_interceptor(self, bzzz_network_instance):
        """Install interceptor into BZZZ network layer"""
        # This would need to be customized based on actual BZZZ implementation
        # Example pattern:
        
        # Store original send function
        self.original_send_function = bzzz_network_instance.send_message
        
        # Replace with intercepting version
        bzzz_network_instance.send_message = self._intercepting_send_message
        
        logger.info("BzzzInterceptor installed into BZZZ network layer")
    
    async def _intercepting_send_message(self, message_data: Dict[str, Any]):
        """Intercepting version of BZZZ send_message function"""
        try:
            # Convert to BzzzMessage format
            bzzz_message = self._convert_to_bzzz_message(message_data)
            
            # Check with interceptor
            should_allow = await self.interceptor.intercept_outgoing_message(bzzz_message)
            
            if should_allow:
                # Call original send function
                return await self.original_send_function(message_data)
            else:
                # Message was blocked
                raise Exception(f"Message blocked by security interceptor: {bzzz_message.message_id}")
                
        except Exception as e:
            logger.error(f"Error in intercepting send: {e}")
            raise
    
    def _convert_to_bzzz_message(self, message_data: Dict[str, Any]) -> BzzzMessage:
        """Convert BZZZ network message format to BzzzMessage"""
        # This would need to be customized based on actual BZZZ message format
        return BzzzMessage(
            message_id=message_data.get('id', f"auto_{int(time.time())}"),
            sender_agent=message_data.get('sender', 'unknown'),
            recipient_agent=message_data.get('recipient'),
            message_type=message_data.get('type', 'unknown'),
            payload=json.dumps(message_data.get('payload', message_data)),
            timestamp=datetime.now(),
            network_metadata=message_data
        )