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))