import asyncio from asyncio.futures import Future from asyncio.tasks import Task, shield import multiprocessing import queue from typing import List import websockets import json from os import _exit from helpers.logging_manager import LoggingManager from websockets.server import Serve class WebsocketServer: threads = Future baps_clients = set() channel_to_q: List[multiprocessing.Queue] webstudio_to_q: List[multiprocessing.Queue] server_name: str logger: LoggingManager to_webstudio: Task from_webstudio: Task websocket_server: Serve def __init__(self, in_q, out_q, state): self.channel_to_q = in_q self.webstudio_to_q = out_q self.logger = LoggingManager("Websockets") self.server_name = state.state["server_name"] self.websocket_server = websockets.serve( self.websocket_handler, state.state["host"], state.state["ws_port"] ) asyncio.get_event_loop().run_until_complete(self.websocket_server) try: asyncio.get_event_loop().run_forever() except Exception : # Sever died somehow, just quit out. self.quit() def quit(self): del self.websocket_server del self.logger _exit(0) def __del__(self): print("Deleting websocket server") self.quit() async def websocket_handler(self, websocket, _): self.baps_clients.add(websocket) await websocket.send( json.dumps({"message": "Hello", "serverName": self.server_name}) ) self.logger.log.info("New Client: {}".format(websocket)) for channel in self.channel_to_q: channel.put("WEBSOCKET:STATUS") async def handle_from_webstudio(): try: async for message in websocket: data = json.loads(message) if "channel" not in data: # Didn't specify a channel, send to all. for channel in range(len(self.channel_to_q)): sendCommand(channel, data) else: channel = int(data["channel"]) sendCommand(channel, data) async def send(conn, message): conn.send(message) await asyncio.wait( [send(conn, message) for conn in self.baps_clients] ) except websockets.exceptions.ConnectionClosedError as e: self.logger.log.error( "Client Disconncted {}, {}".format(websocket, e)) # TODO: Proper Logging except Exception as e: self.logger.log.exception( "Exception handling messages from Websocket.\n{}".format(e) ) finally: self.logger.log.info("Removing client: {}".format(websocket)) self.baps_clients.remove(websocket) def sendCommand(channel, data): if channel not in range(len(self.channel_to_q)): self.logger.log.exception( "Received channel number larger than server supported channels." ) return if "command" in data.keys(): command = data["command"] # Handle the general case # Message format: # SOURCE:COMMAND:EXTRADATA message = "WEBSOCKET:" + command # If we just want PLAY, PAUSE etc, we're all done. # Else, let's pipe in some extra info. extra = "" try: if command == "SEEK": extra += str(data["time"]) elif command == "LOAD": extra += str(data["weight"]) elif command == "AUTOADVANCE": extra += str(data["enabled"]) elif command == "PLAYONLOAD": extra += str(data["enabled"]) elif command == "REPEAT": extra += str(data["mode"]).lower() elif command == "ADD": extra += json.dumps(data["newItem"]) elif command == "REMOVE": extra += str(data["weight"]) elif command == "GET_PLAN": extra += str(data["timeslotId"]) # SPECIAL CASE ALERT! We need to talk to two channels here. elif command == "MOVE": # TODO Should we trust the client with the item info? # Tell the old channel to remove "weight" extra += str(data["weight"]) # Now modify the item with the weight in the new channel new_channel = int(data["new_channel"]) item = data["item"] item["weight"] = int(data["new_weight"]) # Now send the special case. self.channel_to_q[new_channel].put( "ADD:" + json.dumps(item)) except ValueError as e: self.logger.log.exception( "Error decoding extra data {} for command {} ".format( e, command ) ) pass # Stick the message together and send! if extra != "": message += ":" + extra try: self.channel_to_q[channel].put(message) except Exception as e: self.logger.log.exception( "Failed to send message {} to channel {}: {}".format( message, channel, e ) ) else: self.logger.log.error( "Command missing from message. Data: {}".format(data) ) async def handle_to_webstudio(): while True: for channel in range(len(self.webstudio_to_q)): try: message = self.webstudio_to_q[channel].get_nowait() source = message.split(":")[0] # TODO ENUM if source not in ["WEBSOCKET", "ALL"]: print( "ERROR: Message received from invalid source to websocket_handler. Ignored.", source, message, ) continue command = message.split(":")[1] # print("Websocket Out:", command) if command == "STATUS": try: message = message.split("OKAY:")[1] message = json.loads(message) except Exception : continue # TODO more logging elif command == "POS": try: message = message.split(":", 2)[2] except Exception : continue elif command == "QUIT": self.quit() else: continue data = json.dumps( {"command": command, "data": message, "channel": channel} ) await asyncio.wait( [conn.send(data) for conn in self.baps_clients] ) except queue.Empty: continue except ValueError: # Typically a "Set of coroutines/Futures is empty." when sending to a dead client. continue except Exception as e: self.logger.log.exception( "Exception trying to send to websocket:", e ) await asyncio.sleep(0.02) self.from_webstudio = asyncio.create_task(handle_from_webstudio()) self.to_webstudio = asyncio.create_task(handle_to_webstudio()) try: self.threads = await shield( asyncio.gather(self.from_webstudio, self.to_webstudio) ) finally: self.from_webstudio.cancel() self.to_webstudio.cancel() if __name__ == "__main__": print("Don't do this")