238 lines
8.1 KiB
Python
238 lines
8.1 KiB
Python
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, path):
|
|
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")
|
|
|
|
self.from_webstudio = asyncio.create_task(self.handle_from_webstudio(websocket))
|
|
self.to_webstudio = asyncio.create_task(self.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()
|
|
|
|
async def handle_from_webstudio(self, websocket):
|
|
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)):
|
|
self.sendCommand(channel, data)
|
|
else:
|
|
channel = int(data["channel"])
|
|
self.sendCommand(channel, data)
|
|
|
|
await asyncio.wait(
|
|
[conn.send(message) for conn in self.baps_clients]
|
|
)
|
|
|
|
except websockets.exceptions.ConnectionClosedError as e:
|
|
self.logger.log.error(
|
|
"Client Disconncted {}, {}".format(websocket, e))
|
|
|
|
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(self, 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(self):
|
|
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)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Don't do this")
|