mc-crafty-controller/app/classes/shared/websocket_manager.py

import json
import logging
from peewee import DoesNotExist

from app.classes.shared.singleton import Singleton
from app.classes.shared.console import Console
from app.classes.models.users import HelperUsers
from app.classes.models.server_permissions import (
    PermissionsServers,
)

logger = logging.getLogger(__name__)


class WebSocketManager(metaclass=Singleton):
    """Track active WebSocket clients and broadcast events to matching clients."""

    def __init__(self):
        """Initialize the shared client registry."""
        self.clients = set()

    def add_client(self, client):
        """Register a WebSocket client for future broadcasts.

        Args:
            client: WebSocket handler instance that implements the methods used
                by this manager, such as ``send_message`` and ``get_user_id``.
        """
        self.clients.add(client)

    def remove_client(self, client):
        """Remove a WebSocket client from the registry.

        Args:
            client: WebSocket handler instance to remove.
        """
        if client in self.clients:
            self.clients.remove(client)
        else:
            logger.exception("Error caught while removing unknown WebSocket client")

    def broadcast(self, event_type: str, data):
        """Send an event to every connected WebSocket client.

        Args:
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """
        logger.debug(
            f"Sending to {len(self.clients)} clients: "
            f"{json.dumps({'event': event_type, 'data': data})}"
        )
        for client in self.clients:
            try:
                client.send_message(event_type, data)
            except Exception as e:
                logger.exception(
                    f"Error caught while sending WebSocket message to "
                    f"{client.get_remote_ip()} {e}"
                )

    def broadcast_to_admins(self, event_type: str, data):
        """Send an event to connected clients whose users are super users.

        Args:
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """

        def filter_fn(client):
            if str(client.get_user_id()) in str(HelperUsers.get_super_user_list()):
                return True
            return False

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_to_non_admins(self, event_type: str, data):
        """Send an event to connected clients whose users are not super users.

        Args:
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """

        def filter_fn(client):
            if str(client.get_user_id()) not in str(HelperUsers.get_super_user_list()):
                return True
            return False

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_page(self, page: str, event_type: str, data):
        """Send an event to clients currently viewing a specific page.

        Args:
            page (str): Page path recorded on the WebSocket client.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """

        def filter_fn(client):
            return client.page == page

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_user(self, user_id: str, event_type: str, data):
        """Send an event to all connected clients for a specific user.

        Args:
            user_id (str): User id returned by the WebSocket client.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """

        def filter_fn(client):
            return client.get_user_id() == user_id

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_to_server_users(self, server_id: str, event_type: str, data):
        """Send an event to users with permission to access a server.

        Args:
            server_id (str): Server id used to look up permitted users.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """
        server_users = PermissionsServers.get_server_user_list(server_id)
        for user in server_users:
            self.broadcast_user(user, event_type, data)

    def broadcast_user_page(self, page: str, user_id: str, event_type: str, data):
        """Send an event to a user's clients on a specific page.

        Args:
            page (str): Page path recorded on the WebSocket client.
            user_id (str): User id returned by the WebSocket client.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """

        def filter_fn(client):
            if client.get_user_id() != user_id:
                return False
            if client.page != page:
                return False
            return True

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_page_params(
        self, page: str, params: dict, event_type: str, data, **kwargs
    ):
        """Send an event to clients on a page with matching query params.

        Args:
            page (str): Page path, for example ``"/panel/server_detail"``.
            params (dict): Query parameters that must match the client's page
                query parameters.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
            **kwargs: Optional filters. Supports ``required_permission`` to
                require a server permission before sending to a client.
        """

        def filter_fn(client):
            kwarg_perms = kwargs.get("required_permission")
            try:
                user_perms = PermissionsServers.get_user_id_permissions_list(
                    client.get_user_id(), params.get("id", "")
                )
            except DoesNotExist as why:
                logger.exception(
                    "User perms not found for websocket filter terminal buffer: %s",
                    why,
                )
                user_perms = []
            if kwarg_perms and (kwarg_perms not in user_perms):
                # Only send data to users with proper permission
                return False
            if client.page != page:
                return False
            for key, param in params.items():
                if param != client.page_query_params.get(key, None):
                    return False
            return True

        self.broadcast_with_fn(filter_fn, event_type, data)

    def broadcast_with_fn(self, filter_fn, event_type: str, data):
        """Send an event to clients accepted by a filter callback.

        Args:
            filter_fn: Callable that receives a client and returns ``True`` when
                the client should receive the event.
            event_type (str): Client-side event name to emit.
            data: JSON-serializable payload for the event.
        """
        # assign self.clients to a static variable here so hopefully
        # the set size won't change
        static_clients = self.clients
        clients = list(filter(filter_fn, static_clients.copy()))
        logger.debug(
            f"Sending to {len(clients)}  \
            out of {len(self.clients)} "
            f"clients: {json.dumps({'event': event_type, 'data': data})}"
        )

        for client in clients[:]:
            try:
                client.send_message(event_type, data)
            except Exception as e:
                logger.exception(
                    f"Error catched while sending WebSocket message to "
                    f"{client.get_remote_ip()} {e}"
                )

    def disconnect_all(self):
        """Close every connected WebSocket client."""
        Console.info("Disconnecting WebSocket clients")
        for client in self.clients:
            client.close()
        Console.info("Disconnected WebSocket clients")