diff --git a/royalnet/bots/telegram.py b/royalnet/bots/telegram.py index b37e3667..9d191876 100644 --- a/royalnet/bots/telegram.py +++ b/royalnet/bots/telegram.py @@ -20,7 +20,7 @@ class TelegramBot: self.should_run: bool = False self.offset: int = -100 self.missing_command: typing.Callable = missing_command - self.network: RoyalnetLink = RoyalnetLink(master_server_uri, "Telegram", null) + self.network: RoyalnetLink = RoyalnetLink(master_server_uri, "telegram", null) # Generate commands self.commands = {} for command in commands: diff --git a/royalnet/network/messages.py b/royalnet/network/messages.py index 8d1b5b48..59824eb2 100644 --- a/royalnet/network/messages.py +++ b/royalnet/network/messages.py @@ -12,5 +12,9 @@ class ErrorMessage(Message): self.reason = reason -class InvalidSecretErrorMessage(ErrorMessage): +class BadMessage(ErrorMessage): + pass + + +class InvalidSecretErrorMessage(BadMessage): pass diff --git a/royalnet/network/royalnetlink.py b/royalnet/network/royalnetlink.py index 08e38bab..d07baf4e 100644 --- a/royalnet/network/royalnetlink.py +++ b/royalnet/network/royalnetlink.py @@ -36,6 +36,7 @@ class PendingRequest: class RoyalnetLink: def __init__(self, master_uri: str, link_type: str, request_handler): + assert ":" not in link_type self.master_uri: str = master_uri self.link_type: str = link_type self.nid: str = str(uuid.uuid4()) @@ -70,7 +71,7 @@ class RoyalnetLink: @requires_connection async def identify(self, secret) -> None: - await self.websocket.send(f"Identify: {self.nid}:{secret}") + await self.websocket.send(f"Identify {self.nid}:{self.link_type}:{secret}") response_package = await self.receive() response = response_package.data if isinstance(response, ErrorMessage): diff --git a/royalnet/network/royalnetserver.py b/royalnet/network/royalnetserver.py new file mode 100644 index 00000000..ec9b4cb3 --- /dev/null +++ b/royalnet/network/royalnetserver.py @@ -0,0 +1,54 @@ +import typing +import websockets +import re +import datetime +from .messages import Message, ErrorMessage, BadMessage, InvalidSecretErrorMessage, IdentifySuccessfulMessage +from .packages import Package, TwoWayPackage + + +class ConnectedClient: + def __init__(self, socket: websockets.WebSocketServerProtocol): + self.socket: websockets.WebSocketServerProtocol = socket + self.nid: str = None + self.link_type: str = None + self.connection_datetime: datetime.datetime = datetime.datetime.now() + + @property + def is_identified(self) -> bool: + return bool(self.nid) + + +class RoyalnetServer: + def __init__(self, required_secret: str): + self.required_secret: str = required_secret + self.connected_clients: typing.List[ConnectedClient] = {} + self.server: websockets.server.WebSocketServer = websockets.server + + def find_client_by_nid(self, nid: str): + return [client for client in self.connected_clients if client.nid == nid][0] + + async def listener(self, websocket: websockets.server.WebSocketServerProtocol, request_uri: str): + connected_client = ConnectedClient(websocket) + # Wait for identification + identify_msg = websocket.recv() + if not isinstance(identify_msg, str): + websocket.send(BadMessage("Invalid identification message (not a str)")) + return + identification = re.match(r"Identify ([A-Za-z0-9\-]+):([a-z]+):([A-Za-z0-9\-])", identify_msg) + if identification is None: + websocket.send(BadMessage("Invalid identification message (regex failed)")) + return + secret = identification.group(3) + if secret != self.required_secret: + websocket.send(InvalidSecretErrorMessage("Invalid secret")) + return + # Identification successful + connected_client.nid = identification.group(1) + connected_client.link_type = identification.group(2) + self.connected_clients.append(connected_client) + # Main loop + while True: + pass + + +