133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from threading import Thread
|
|
from typing import Any, Self, Sequence, Optional
|
|
from http.cookies import SimpleCookie
|
|
from Abstracts.WebSocketServersAbstract import WebSocketServersAbstract
|
|
from Abstracts.ModelAbstract import ModelAbstract
|
|
from websockets.sync.server import serve as server_serve
|
|
from websockets import Server as WebSocketServer, ClientConnection as WebSocketClient
|
|
from Interfaces.Application.AnPInterface import AnPInterface
|
|
from Utils.Checks import Check
|
|
|
|
class WebSocketServerDriver(WebSocketServersAbstract, ModelAbstract):
|
|
|
|
def __init__(self:Self, anp:AnPInterface, key:str, inputs:Optional[dict[str, Any|None]|Sequence[Any|None]] = None) -> None:
|
|
super().__init__(anp, key, inputs)
|
|
|
|
self.__server:WebSocketServer
|
|
self.__clients:dict[str, WebSocketClient] = {}
|
|
self.__thread:Thread = None
|
|
|
|
anp.settings.get(("web_socket_server_autostart", "autostart"), inputs, True) and self.start()
|
|
|
|
def __run(self:Self) -> None:
|
|
self.__server = server_serve(self.__handler, self.host, self.port)
|
|
self.__server.serve_forever()
|
|
|
|
def start(self:Self) -> None:
|
|
self.__thread = Thread(target = self.__run)
|
|
self.__thread.start()
|
|
|
|
def close(self:Self) -> None:
|
|
|
|
id:str
|
|
|
|
for id in tuple(self.__clients.keys()):
|
|
self.close_client(id)
|
|
|
|
self.__server.shutdown()
|
|
|
|
def close_client(self:Self, ids:Optional[str|list[str]] = None) -> None:
|
|
|
|
id:str
|
|
|
|
for id in (
|
|
ids if Check.is_array(ids) else
|
|
list(self.__clients.keys()) if ids is None else
|
|
[ids] if Check.is_string(ids) else
|
|
[]):
|
|
if id in self.__clients:
|
|
try:
|
|
self.__clients[id].close()
|
|
except Exception as exception:
|
|
self.anp.exception(exception, "web_socket_server_client_close_exception", {
|
|
"id": id,
|
|
"port": self.port,
|
|
"host": self.host
|
|
})
|
|
|
|
def __handler(self:Self, client:WebSocketClient) -> None:
|
|
|
|
id:str = self.anp.unique_keys.get()
|
|
cookies:SimpleCookie = SimpleCookie()
|
|
|
|
self.__clients[id] = client
|
|
self.on_new_client.execute(id)
|
|
|
|
cookies.load(client.request.headers.get("Cookie", ''))
|
|
self.set_session(id, cookies.get("session").value)
|
|
|
|
self.anp.print("info", "web_socket_server_client_connected", {
|
|
"id": id,
|
|
"port": self.port,
|
|
"host": self.host,
|
|
"client_host" : client.remote_address[0],
|
|
"client_port" : client.remote_address[1]
|
|
})
|
|
|
|
self.send("web_socket_client", "set_id", id, id)
|
|
|
|
try:
|
|
while self.anp.working():
|
|
message:str = client.recv()
|
|
if message is None:
|
|
break
|
|
self.on_message.execute(id, message)
|
|
except Exception as exception:
|
|
self.anp.working() and self.anp.exception(exception, "web_socket_server_client_exception", {
|
|
"id": id,
|
|
"port": self.port,
|
|
"host": self.host
|
|
})
|
|
self.on_error.execute(id, exception)
|
|
finally:
|
|
try:
|
|
self.__clients[id].close()
|
|
except Exception as _:
|
|
pass
|
|
self.remove_session(id)
|
|
del self.__clients[id]
|
|
self.anp.unique_keys.remove(id)
|
|
self.on_close.execute(id)
|
|
self.anp.print("info", "web_socket_server_client_disconnected", {
|
|
"id": id,
|
|
"port": self.port,
|
|
"host": self.host
|
|
})
|
|
|
|
def send(self:Self, controller:str, action:str, data:str, ids:Optional[str|Sequence[str]] = None, code:int = 200) -> bool:
|
|
|
|
success:bool = True
|
|
id:str
|
|
|
|
# print([ids, list(self.__clients.keys())])
|
|
|
|
for id in (
|
|
list(self.__clients.keys()) if ids is None else
|
|
ids if Check.is_array(ids) else
|
|
[ids]):
|
|
if id in self.__clients:
|
|
try:
|
|
self.__clients[id].send(self.format_data(controller, action, data, id, code))
|
|
# print(["WS", controller, action, data, ids, code, list(self.__clients.keys())])
|
|
except Exception as exception:
|
|
self.anp.exception(exception, "web_socket_server_client_send_exception", {
|
|
"id": id,
|
|
"port": self.port,
|
|
"host": self.host
|
|
})
|
|
success = False
|
|
|
|
return success |