NeuroSploit/backend/api/websocket.py

"""
NeuroSploit v3 - WebSocket Manager
"""
from typing import Dict, List, Optional
from fastapi import WebSocket
import json
import asyncio

try:
    from backend.core.notification_manager import notification_manager, NotificationEvent
    HAS_NOTIFICATIONS = True
except ImportError:
    HAS_NOTIFICATIONS = False


class ConnectionManager:
    """Manages WebSocket connections for real-time updates"""

    def __init__(self):
        # scan_id -> list of websocket connections
        self.active_connections: Dict[str, List[WebSocket]] = {}
        self._lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket, scan_id: str):
        """Accept a WebSocket connection and register it for a scan"""
        await websocket.accept()
        async with self._lock:
            if scan_id not in self.active_connections:
                self.active_connections[scan_id] = []
            self.active_connections[scan_id].append(websocket)
        print(f"WebSocket connected for scan: {scan_id}")

    def disconnect(self, websocket: WebSocket, scan_id: str):
        """Remove a WebSocket connection"""
        if scan_id in self.active_connections:
            if websocket in self.active_connections[scan_id]:
                self.active_connections[scan_id].remove(websocket)
            if not self.active_connections[scan_id]:
                del self.active_connections[scan_id]
        print(f"WebSocket disconnected for scan: {scan_id}")

    async def send_to_scan(self, scan_id: str, message: dict):
        """Send a message to all connections watching a specific scan"""
        if scan_id not in self.active_connections:
            return

        dead_connections = []
        for connection in self.active_connections[scan_id]:
            try:
                await connection.send_text(json.dumps(message))
            except Exception:
                dead_connections.append(connection)

        # Clean up dead connections
        for conn in dead_connections:
            self.disconnect(conn, scan_id)

    async def broadcast_scan_started(self, scan_id: str, target: str = ""):
        """Notify that a scan has started"""
        await self.send_to_scan(scan_id, {
            "type": "scan_started",
            "scan_id": scan_id
        })
        if HAS_NOTIFICATIONS:
            asyncio.create_task(notification_manager.notify(
                NotificationEvent.SCAN_STARTED, {"target": target, "scan_id": scan_id}
            ))

    async def broadcast_phase_change(self, scan_id: str, phase: str):
        """Notify phase change (recon, testing, reporting)"""
        await self.send_to_scan(scan_id, {
            "type": "phase_change",
            "scan_id": scan_id,
            "phase": phase
        })

    async def broadcast_progress(self, scan_id: str, progress: int, message: Optional[str] = None):
        """Send progress update"""
        await self.send_to_scan(scan_id, {
            "type": "progress_update",
            "scan_id": scan_id,
            "progress": progress,
            "message": message
        })

    async def broadcast_endpoint_found(self, scan_id: str, endpoint: dict):
        """Notify a new endpoint was discovered"""
        await self.send_to_scan(scan_id, {
            "type": "endpoint_found",
            "scan_id": scan_id,
            "endpoint": endpoint
        })

    async def broadcast_path_crawled(self, scan_id: str, path: str, status: int):
        """Notify a path was crawled"""
        await self.send_to_scan(scan_id, {
            "type": "path_crawled",
            "scan_id": scan_id,
            "path": path,
            "status": status
        })

    async def broadcast_url_discovered(self, scan_id: str, url: str):
        """Notify a URL was discovered"""
        await self.send_to_scan(scan_id, {
            "type": "url_discovered",
            "scan_id": scan_id,
            "url": url
        })

    async def broadcast_test_started(self, scan_id: str, vuln_type: str, endpoint: str):
        """Notify a vulnerability test has started"""
        await self.send_to_scan(scan_id, {
            "type": "test_started",
            "scan_id": scan_id,
            "vulnerability_type": vuln_type,
            "endpoint": endpoint
        })

    async def broadcast_test_completed(self, scan_id: str, vuln_type: str, endpoint: str, is_vulnerable: bool):
        """Notify a vulnerability test has completed"""
        await self.send_to_scan(scan_id, {
            "type": "test_completed",
            "scan_id": scan_id,
            "vulnerability_type": vuln_type,
            "endpoint": endpoint,
            "is_vulnerable": is_vulnerable
        })

    async def broadcast_vulnerability_found(self, scan_id: str, vulnerability: dict):
        """Notify a vulnerability was found"""
        await self.send_to_scan(scan_id, {
            "type": "vuln_found",
            "scan_id": scan_id,
            "vulnerability": vulnerability
        })
        if HAS_NOTIFICATIONS:
            asyncio.create_task(notification_manager.notify(
                NotificationEvent.VULN_FOUND, {
                    "title": vulnerability.get("title", "Vulnerability Found"),
                    "severity": vulnerability.get("severity", "medium"),
                    "vulnerability_type": vulnerability.get("vulnerability_type", "unknown"),
                    "endpoint": vulnerability.get("endpoint", ""),
                    "description": vulnerability.get("description", ""),
                }
            ))

    async def broadcast_log(self, scan_id: str, level: str, message: str):
        """Send a log message"""
        await self.send_to_scan(scan_id, {
            "type": "log_message",
            "scan_id": scan_id,
            "level": level,
            "message": message
        })

    async def broadcast_scan_completed(self, scan_id: str, summary: dict):
        """Notify that a scan has completed"""
        await self.send_to_scan(scan_id, {
            "type": "scan_completed",
            "scan_id": scan_id,
            "summary": summary
        })
        if HAS_NOTIFICATIONS:
            asyncio.create_task(notification_manager.notify(
                NotificationEvent.SCAN_COMPLETED, {
                    "total_vulnerabilities": summary.get("total_vulnerabilities", 0),
                    "critical": summary.get("critical", 0),
                    "high": summary.get("high", 0),
                    "medium": summary.get("medium", 0),
                }
            ))

    async def broadcast_scan_stopped(self, scan_id: str, summary: dict):
        """Notify that a scan was stopped by user"""
        await self.send_to_scan(scan_id, {
            "type": "scan_stopped",
            "scan_id": scan_id,
            "status": "stopped",
            "summary": summary
        })

    async def broadcast_scan_failed(self, scan_id: str, error: str, summary: dict = None):
        """Notify that a scan has failed"""
        await self.send_to_scan(scan_id, {
            "type": "scan_failed",
            "scan_id": scan_id,
            "status": "failed",
            "error": error,
            "summary": summary or {}
        })
        if HAS_NOTIFICATIONS:
            asyncio.create_task(notification_manager.notify(
                NotificationEvent.SCAN_FAILED, {"error": error}
            ))

    async def broadcast_stats_update(self, scan_id: str, stats: dict):
        """Broadcast updated scan statistics"""
        await self.send_to_scan(scan_id, {
            "type": "stats_update",
            "scan_id": scan_id,
            "stats": stats
        })

    async def broadcast_agent_task(self, scan_id: str, task: dict):
        """Broadcast agent task update (created, started, completed, failed)"""
        await self.send_to_scan(scan_id, {
            "type": "agent_task",
            "scan_id": scan_id,
            "task": task
        })

    async def broadcast_agent_task_started(self, scan_id: str, task: dict):
        """Broadcast when an agent task starts"""
        await self.send_to_scan(scan_id, {
            "type": "agent_task_started",
            "scan_id": scan_id,
            "task": task
        })

    async def broadcast_agent_task_completed(self, scan_id: str, task: dict):
        """Broadcast when an agent task completes"""
        await self.send_to_scan(scan_id, {
            "type": "agent_task_completed",
            "scan_id": scan_id,
            "task": task
        })

    async def broadcast_report_generated(self, scan_id: str, report: dict):
        """Broadcast when a report is generated"""
        await self.send_to_scan(scan_id, {
            "type": "report_generated",
            "scan_id": scan_id,
            "report": report
        })

    async def broadcast_error(self, scan_id: str, error: str):
        """Notify an error occurred"""
        await self.send_to_scan(scan_id, {
            "type": "error",
            "scan_id": scan_id,
            "error": error
        })


# Global instance
manager = ConnectionManager()