websocket_handler.py

"""
WebSocket handler for real-time updates to frontend.

Broadcasts events to all connected clients when:
- New MCP servers are registered
- New messages/events arrive
- Detection engine results are available
"""

import asyncio
import json
from typing import Set
from aiohttp import web, WSMsgType
from utils import safe_print


class WebSocketHandler:
    """
    Manages WebSocket connections and broadcasts real-time updates.
    """

    def __init__(self):
        self.connections: Set[web.WebSocketResponse] = set()
        self.running = False

    async def start(self):
        """Start the WebSocket handler."""
        self.running = True
        safe_print('[WebSocket] Handler started')

    async def stop(self):
        """Stop the WebSocket handler and close all connections."""
        self.running = False

        # Close all active connections
        if self.connections:
            safe_print(f'[WebSocket] Closing {len(self.connections)} connections...')
            close_tasks = []
            for ws in list(self.connections):
                close_tasks.append(ws.close())

            await asyncio.gather(*close_tasks, return_exceptions=True)
            self.connections.clear()

        safe_print('[WebSocket] Handler stopped')

    async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
        """
        Handle incoming WebSocket connection.

        Args:
            request: aiohttp request object

        Returns:
            WebSocketResponse
        """
        ws = web.WebSocketResponse(heartbeat=30)
        await ws.prepare(request)

        # Add to active connections
        self.connections.add(ws)
        client_id = id(ws)
        safe_print(f'[WebSocket] Client {client_id} connected. Total: {len(self.connections)}')

        # Send initial connection success message
        await self.send_to_client(ws, {
            'type': 'connection',
            'status': 'connected',
            'message': 'WebSocket connection established'
        })

        try:
            # Listen for messages from client
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    data = json.loads(msg.data)
                    msg_type = data.get('type')

                    if msg_type == 'ping':
                        await self.send_to_client(ws, {
                            'type': 'pong',
                            'timestamp': data.get('timestamp')
                        })
                    elif msg_type == 'blocking_decision':
                        # Handle user's blocking decision
                        await self._handle_blocking_decision(data)
                elif msg.type == WSMsgType.ERROR:
                    safe_print(f'[WebSocket] Client {client_id} error: {ws.exception()}')
                elif msg.type == WSMsgType.CLOSE:
                    safe_print(f'[WebSocket] Client {client_id} closed connection')
                    break
        except Exception as e:
            safe_print(f'[WebSocket] Client {client_id} exception: {e}')
        finally:
            # Remove from active connections
            self.connections.discard(ws)
            safe_print(f'[WebSocket] Client {client_id} disconnected. Total: {len(self.connections)}')

        return ws

    async def send_to_client(self, ws: web.WebSocketResponse, data: dict):
        """
        Send data to a specific client.

        Args:
            ws: WebSocket connection
            data: Dictionary to send as JSON
        """
        try:
            if not ws.closed:
                await ws.send_json(data)
        except Exception as e:
            safe_print(f'[WebSocket] Error sending to client: {e}')

    async def broadcast(self, event_type: str, data: dict):
        """
        Broadcast event to all connected clients.

        Args:
            event_type: Type of event ('server_update', 'message_update', etc.)
            data: Event data to broadcast
        """
        if not self.running or not self.connections:
            return

        message = {
            'type': event_type,
            'data': data
        }

        # Broadcast to all clients
        dead_connections = set()
        for ws in self.connections:
            try:
                if ws.closed:
                    dead_connections.add(ws)
                else:
                    await ws.send_json(message)
            except Exception as e:
                safe_print(f'[WebSocket] Error broadcasting to client: {e}')
                dead_connections.add(ws)

        # Clean up dead connections
        if dead_connections:
            self.connections -= dead_connections
            safe_print(f'[WebSocket] Removed {len(dead_connections)} dead connections')

    async def broadcast_server_update(self):
        """Notify clients that server list has changed."""
        await self.broadcast('server_update', {
            'message': 'Server list updated',
            'action': 'refresh_servers'
        })
        safe_print('[WebSocket] Broadcasted server_update')

    async def broadcast_message_update(self, server_id: int, server_name: str):
        """
        Notify clients that new messages are available for a server.

        Args:
            server_id: Server database ID
            server_name: Server name/tag
        """
        await self.broadcast('message_update', {
            'server_id': server_id,
            'server_name': server_name,
            'action': 'refresh_messages'
        })
        safe_print(f'[WebSocket] Broadcasted message_update for {server_name}')

    async def broadcast_detection_result(self, event_id: int, engine_name: str, severity: str):
        """
        Notify clients that a detection result is available.

        Args:
            event_id: Raw event ID
            engine_name: Name of detection engine
            severity: Severity level (none/low/medium/high)
        """
        await self.broadcast('detection_result', {
            'event_id': event_id,
            'engine_name': engine_name,
            'severity': severity,
            'action': 'refresh_detections'
        })
        safe_print(f'[WebSocket] Broadcasted detection_result: {engine_name} ({severity})')

    async def broadcast_reload_all(self):
        """Notify clients to reload all data (nuclear option)."""
        await self.broadcast('reload_all', {
            'message': 'Full reload requested',
            'action': 'reload_all'
        })
        safe_print('[WebSocket] Broadcasted reload_all')

    async def broadcast_tool_safety_update(self, mcp_tag: str, tool_name: str, safety: int):
        """
        Notify clients that a tool's safety status has been updated.

        Args:
            mcp_tag: MCP server tag
            tool_name: Tool name
            safety: New safety value (0-3)
        """
        await self.broadcast('tool_safety_update', {
            'mcp_tag': mcp_tag,
            'tool_name': tool_name,
            'safety': safety,
            'action': 'refresh_tools'
        })
        safe_print(f'[WebSocket] Broadcasted tool_safety_update: {mcp_tag}/{tool_name} -> {safety}')

    async def broadcast_custom_rule_update(self, engine_name: str):
        """
        Notify clients that custom rules have been updated for an engine.

        Args:
            engine_name: Engine name (e.g., 'pii_leak_engine')
        """
        await self.broadcast('custom_rule_update', {
            'engine_name': engine_name,
            'message': f'Custom rules updated for {engine_name}',
            'action': 'refresh_custom_rules'
        })
        safe_print(f'[WebSocket] Broadcasted custom_rule_update: {engine_name}')

    async def broadcast_blocking_request(self, request_id: str, event_data: dict,
                                         detection_results: list, engine_name: str,
                                         severity: str, server_name: str, tool_name: str):
        """
        Notify clients that a blocking decision is needed.

        Args:
            request_id: Unique ID for this blocking request
            event_data: Original event data
            detection_results: List of detection findings
            engine_name: Name of detection engine
            severity: Severity level
            server_name: Server name
            tool_name: Tool name being called
        """
        await self.broadcast('blocking_request', {
            'request_id': request_id,
            'event_data': event_data,
            'detection_results': detection_results,
            'engine_name': engine_name,
            'severity': severity,
            'server_name': server_name,
            'tool_name': tool_name,
            'action': 'show_blocking_modal'
        })
        safe_print(f'[WebSocket] Broadcasted blocking_request: {request_id} ({engine_name}: {severity})')

    async def _handle_blocking_decision(self, data: dict):
        """
        Handle user's blocking decision from frontend.

        Args:
            data: Decision data with request_id and decision (allow/block)
        """
        from state import state

        request_id = data.get('request_id')
        decision = data.get('decision')  # 'allow' or 'block'

        if not request_id or not decision:
            safe_print(f'[WebSocket] Invalid blocking decision: {data}')
            return

        blocking_request = await state.get_blocking_request(request_id)
        if not blocking_request:
            safe_print(f'[WebSocket] Blocking request not found: {request_id}')
            return

        # Resolve the future with user's decision
        if blocking_request.future and not blocking_request.future.done():
            blocking_request.future.set_result(decision == 'allow')
            safe_print(f'[WebSocket] Blocking decision received: {request_id} -> {decision}')

        # Clean up
        await state.remove_blocking_request(request_id)


# Global singleton instance
ws_handler = WebSocketHandler()