controller-wisetopdataserver/python/biopro/server/websocket.py

from typing import Set, Optional, List

# noinspection PyPackageRequirements
from SimpleWebSocketServer import SimpleWebSocketServer, WebSocket

from biopro.message.request import DataMessage, DataMessageHandler
from biopro.util.console import CYAN, RED, YELLOW, pc
from biopro.util.logger import LoggerFlag
from .controller import ControlServerAPI
from .socket import ServerThread

__all__ = ['WebsocketThread', 'WebsocketHandler']


class WebsocketThread(ServerThread):
    """Websocket server thread"""

    def __init__(self, server: ControlServerAPI, url: str, port: int):
        """

        :param server: data server to attach
        :param url: socket address
        :param port: socket port
        """
        super().__init__('WebSocketServer', RED)

        # server attribute
        self._websocket_url: str = url
        self._websocket_port: int = port
        self._websocket_server: SimpleWebSocketServer = None

        # socket attribute
        self.active_client: Set[WebsocketHandler] = set()

        # websocket command handler
        self._handler = server

    @property
    def websocket_url(self) -> str:
        return self._websocket_url

    @property
    def websocket_port(self) -> int:
        return self._websocket_port

    @property
    def websocket_address(self) -> str:
        return 'ws://%s:%d' % (self._websocket_url, self._websocket_port)

    def setup(self) -> None:
        super().setup()

        self.log_info('start listen %s:%d' % (self._websocket_url, self._websocket_port))
        self._websocket_server = server = SimpleWebSocketServer(self._websocket_url,
                                                                self._websocket_port,
                                                                WebsocketHandler)

        # add extra attribute to server
        server.biopro_websocket_thread = self

    @property
    def is_startup(self):
        return self._websocket_server is not None and super().is_startup

    def shutdown(self) -> None:
        self.log_info('close listen %s:%d' % (self._websocket_url, self._websocket_port))

        server: SimpleWebSocketServer = self._websocket_server
        if server is not None:
            server.close()

        self._websocket_server = None

    def routine(self, event = None) -> None:
        self._websocket_server.serveonce()

    def add_active_client(self, client: 'WebsocketHandler') -> None:
        self.active_client.add(client)

    def remove_active_client(self, client: 'WebsocketHandler') -> None:
        self.active_client.remove(client)

    # noinspection PyShadowingNames
    def broadcast_command(self, message: str, device: Optional[int] = None):
        if not self.is_closed:
            data = DataMessage.response(message)

            if device is None:
                for client in self.active_client:
                    client.sendMessage(data)
            else:
                for client in self.active_client:
                    if device in client.listen_device:
                        client.sendMessage(data)


class WebsocketHandler(WebSocket, LoggerFlag):

    def __init__(self, server, sock, address):
        """

        :param server: websocket server
        :param sock: socket
        :param address: socket address
        """
        WebSocket.__init__(self, server, sock, address)
        LoggerFlag.__init__(self, 'WebSocketHandler', YELLOW)

        self._websocket: Optional[WebsocketThread] = getattr(server, 'biopro_websocket_thread', None)
        self._handler: Optional[ControlServerAPI] = None

        self.data_message_handler = WebsocketDataMessageHandler(self)

        if self._websocket is not None:
            # noinspection PyProtectedMember
            self._handler = self._websocket._handler

    def handleConnected(self):
        """handle connection"""
        self.log_info('connection accept', self.address)

        if self._websocket is not None:
            self._websocket.add_active_client(self)

    def handleClose(self):
        """handle connection close"""
        self.log_info('connection close', self.address)

        self._handler.data_request_stream(False, self.data_message_handler)

        if self._websocket is not None:
            self._websocket.remove_active_client(self)

        self._handler = None

    @property
    def is_closed(self) -> bool:
        return self._handler is None

    # @property
    # def file_handle(self) -> Optional[int]:
    #     return self.data_message_handler.file_handle

    @property
    def listen_device(self) -> List[int]:
        return self.data_message_handler.listen_device

    def handleMessage(self):
        """handle websocket message"""

        command: str = self.data

        if len(command) > 0 and self._handler is not None:
            self.log_info('recv', pc(command, CYAN))

            try:
                self._handle_command(command)

            except RuntimeError as e:
                self.log_warn(e)
                self.sendMessage(DataMessage.response(DataMessage.warn(str(e))))
                # print_exception(self)

            except BaseException as e:
                self.log_warn(e)
                self.sendMessage(DataMessage.response(DataMessage.error(str(e))))
                # print_exception(self)

    def _handle_command(self, command: str):
        server: ControlServerAPI = self._handler

        if server is None:
            return

        if command == DataMessage.GREETING:
            message = DataMessage.greeting()

        elif command == DataMessage.SYNC:
            message = server.websocket_command(command)

        elif command == DataMessage.START_DATA_STREAM:
            server.data_request_stream(True, self.data_message_handler)
            return

        elif command == DataMessage.STOP_DATA_STREAM:
            server.data_request_stream(False, self.data_message_handler)
            return

        elif command.startswith(DataMessage.SYNC):
            message = server.websocket_command(command)

        elif command.startswith(DataMessage.BROADCAST):
            message = server.websocket_command(command)

        else:
            server.websocket_data_request(self.data_message_handler, command)
            return

        if message is None:
            return

        elif isinstance(message, str):
            self.sendMessage(DataMessage.response(message))

        elif isinstance(message, bytes):
            self.sendMessage(message)

        elif isinstance(message, list):
            for x in message:
                self.sendMessage(DataMessage.response(x))

        else:
            self.log_warn('unknown message response type')


class WebsocketDataMessageHandler(DataMessageHandler):
    __slots__ = '__websocket',

    def __init__(self, websocket: WebsocketHandler):
        super().__init__()
        self.__websocket = websocket

    def on_message(self, message: bytes):
        # self.__websocket.log_verbose('send', 'bytes', len(message))
        self.__websocket.sendMessage(message)

    def __str__(self):
        return 'DataMessageHandler[%s]' % ('websocket:' + str(self.__websocket.address))