websocket_monitoring.py

"""
Real-time WebSocket integration for live threat monitoring
"""

from fastapi import WebSocket, WebSocketDisconnect
import asyncio
import json
from datetime import datetime
from typing import Dict, List
import random
import logging

class ConnectionManager:
    """Manage WebSocket connections for real-time updates"""
    
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.logger = logging.getLogger(__name__)
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        self.logger.info(f"New WebSocket connection: {len(self.active_connections)} total")
    
    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        self.logger.info(f"WebSocket disconnected: {len(self.active_connections)} remaining")
    
    async def send_personal_message(self, message: str, websocket: WebSocket):
        try:
            await websocket.send_text(message)
        except Exception as e:
            self.logger.error(f"Failed to send personal message: {e}")
            self.disconnect(websocket)
    
    async def broadcast(self, message: str):
        """Broadcast message to all connected clients"""
        disconnected = []
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception as e:
                self.logger.error(f"Failed to broadcast to connection: {e}")
                disconnected.append(connection)
        
        # Clean up disconnected clients
        for connection in disconnected:
            self.disconnect(connection)

# Global connection manager
manager = ConnectionManager()

class ThreatFeedSimulator:
    """Simulate real-time threat intelligence feeds"""
    
    def __init__(self):
        self.threat_types = [
            "malware_detection",
            "network_intrusion", 
            "data_exfiltration",
            "brute_force_attack",
            "ddos_attempt",
            "suspicious_login",
            "privilege_escalation",
            "lateral_movement"
        ]
        
        self.threat_sources = [
            "firewall_logs",
            "ids_sensor",
            "endpoint_detection",
            "network_monitor", 
            "email_security",
            "web_filter",
            "dns_monitor",
            "user_behavior"
        ]
        
        self.severity_levels = ["LOW", "MEDIUM", "HIGH", "CRITICAL"]
        
    def generate_threat_event(self) -> Dict:
        """Generate a simulated threat event"""
        
        return {
            "event_id": f"evt_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{random.randint(1000, 9999)}",
            "timestamp": datetime.now().isoformat(),
            "threat_type": random.choice(self.threat_types),
            "source": random.choice(self.threat_sources),
            "severity": random.choice(self.severity_levels),
            "confidence": round(random.uniform(0.3, 0.95), 2),
            "source_ip": f"{random.randint(1, 255)}.{random.randint(1, 255)}.{random.randint(1, 255)}.{random.randint(1, 255)}",
            "target_ip": f"192.168.1.{random.randint(1, 254)}",
            "details": self._generate_threat_details(),
            "status": "active"
        }
    
    def _generate_threat_details(self) -> Dict:
        """Generate detailed threat information"""
        
        return {
            "attack_vector": random.choice([
                "network_based",
                "email_based", 
                "web_based",
                "endpoint_based",
                "social_engineering"
            ]),
            "mitre_technique": f"T{random.randint(1001, 1609)}",
            "indicators": [
                f"suspicious_process_{random.randint(1, 100)}.exe",
                f"malicious_domain_{random.randint(1, 50)}.com",
                f"unusual_network_traffic_port_{random.randint(1024, 65535)}"
            ],
            "recommendation": "Investigate immediately and implement containment measures"
        }

# Global threat feed simulator
threat_simulator = ThreatFeedSimulator()

async def threat_feed_worker():
    """Background worker that generates and broadcasts threat events"""
    
    while True:
        if manager.active_connections:
            # Generate threat event
            threat_event = threat_simulator.generate_threat_event()
            
            # Broadcast to all connected clients
            await manager.broadcast(json.dumps({
                "type": "threat_event",
                "data": threat_event
            }))
            
            # Log the event
            logging.getLogger(__name__).info(f"Broadcast threat event: {threat_event['event_id']}")
        
        # Wait before next event (simulate real-time frequency)
        await asyncio.sleep(random.uniform(2, 8))  # 2-8 seconds between events

class ThreatMonitor:
    """Advanced threat monitoring with analytics"""
    
    def __init__(self):
        self.active_threats: List[Dict] = []
        self.threat_history: List[Dict] = []
        self.alert_thresholds = {
            "CRITICAL": 1,  # Alert immediately
            "HIGH": 3,      # Alert after 3 events
            "MEDIUM": 10,   # Alert after 10 events  
            "LOW": 50       # Alert after 50 events
        }
        
    def process_threat_event(self, event: Dict) -> Dict:
        """Process and analyze threat event"""
        
        # Add to active threats
        self.active_threats.append(event)
        self.threat_history.append(event)
        
        # Analyze trends
        analysis = self._analyze_threat_trends()
        
        # Generate alerts if needed
        alerts = self._check_alert_conditions(event)
        
        return {
            "event": event,
            "analysis": analysis,
            "alerts": alerts,
            "statistics": self._generate_statistics()
        }
    
    def _analyze_threat_trends(self) -> Dict:
        """Analyze current threat trends"""
        
        if len(self.threat_history) < 2:
            return {"trend": "insufficient_data"}
        
        recent_events = self.threat_history[-10:]  # Last 10 events
        
        # Count by severity
        severity_counts = {}
        for event in recent_events:
            severity = event["severity"]
            severity_counts[severity] = severity_counts.get(severity, 0) + 1
        
        # Calculate trend
        critical_high_ratio = (severity_counts.get("CRITICAL", 0) + severity_counts.get("HIGH", 0)) / len(recent_events)
        
        if critical_high_ratio > 0.5:
            trend = "escalating"
        elif critical_high_ratio > 0.2:
            trend = "elevated" 
        else:
            trend = "normal"
        
        return {
            "trend": trend,
            "critical_high_ratio": round(critical_high_ratio, 2),
            "severity_distribution": severity_counts,
            "total_recent_events": len(recent_events)
        }
    
    def _check_alert_conditions(self, event: Dict) -> List[Dict]:
        """Check if alerts should be triggered"""
        
        alerts = []
        severity = event["severity"]
        
        # Count recent events of same severity
        recent_same_severity = [
            e for e in self.threat_history[-100:]  # Last 100 events
            if e["severity"] == severity
        ]
        
        threshold = self.alert_thresholds.get(severity, 10)
        
        if len(recent_same_severity) >= threshold:
            alerts.append({
                "type": f"{severity.lower()}_threshold_alert",
                "message": f"Threshold exceeded: {len(recent_same_severity)} {severity} events detected",
                "severity": severity,
                "recommended_action": self._get_recommended_action(severity)
            })
        
        return alerts
    
    def _get_recommended_action(self, severity: str) -> str:
        """Get recommended action based on severity"""
        
        actions = {
            "CRITICAL": "Initiate emergency response procedures immediately",
            "HIGH": "Escalate to security team and begin investigation",
            "MEDIUM": "Increase monitoring and prepare for potential escalation",
            "LOW": "Document and continue routine monitoring"
        }
        
        return actions.get(severity, "Review and assess threat significance")
    
    def _generate_statistics(self) -> Dict:
        """Generate current threat statistics"""
        
        total_active = len(self.active_threats)
        total_history = len(self.threat_history)
        
        if total_history == 0:
            return {"total_events": 0}
        
        # Calculate statistics
        severity_stats = {}
        for event in self.threat_history:
            severity = event["severity"]
            severity_stats[severity] = severity_stats.get(severity, 0) + 1
        
        return {
            "total_events": total_history,
            "active_threats": total_active,
            "severity_distribution": severity_stats,
            "average_confidence": round(
                sum(e["confidence"] for e in self.threat_history) / total_history, 2
            )
        }

# Global threat monitor
threat_monitor = ThreatMonitor()