Add existing project files

node.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+
+import asyncio
+
+from argparse import ArgumentParser
+from datetime import datetime
+from hashlib import md5
+from pathlib import Path
+
+from websockets.asyncio.client import connect
+from websockets.asyncio.server import broadcast
+from websockets.asyncio.server import serve
+from websockets.asyncio.server import ServerConnection
+from websockets.exceptions import ConnectionClosed, ConnectionClosedError
+
+from yaml import safe_load as yml_load
+
+
+def log(msg):
+    print(f"{datetime.now().isoformat()} {msg}")
+
+
+class NodeAddr:
+
+    def __init__(self, host: str, port: int):
+        self.host = str(host)
+        self.port = int(port)
+
+    def ws_uri(self) -> str:
+        return f"ws://{self.host}:{self.port}"
+
+    def identifier(self) -> str:
+        return f"{self.host}_{self.port}"
+
+
+class Node:
+
+    max_known_hashes_size = 1024 * 1024
+
+    def __init__(self, listen_address: NodeAddr, receivers: list[NodeAddr]):
+        self.listen_address = listen_address
+        self.receivers = list(receivers)
+        for r in receivers:
+            if not isinstance(r, NodeAddr):
+                raise TypeError(f"{r} must be of type NodeAddr")
+        # internal
+        self._connections: dict = {} # identifer: connection object
+        self._known_hashes: list[str] = [] # list of message hashes
+
+
+    def _relay(self, msg: str):
+        msg_hash = md5(msg.encode()).digest()
+        if not msg_hash in self._known_hashes:
+            self._known_hashes.append(msg_hash)
+            broadcast(self._connections.values(), msg)
+        while len(self._known_hashes) > self.max_known_hashes_size:
+            self._known_hashes.pop(0)
+
+
+    async def _receiver_connection(self, r: NodeAddr):
+        identifier = r.identifier()
+        async for ws in connect(r.ws_uri()):
+            log(identifier)
+            self._connections[identifier] = ws
+            try:
+                async for msg in ws:
+                    log(msg)
+                    self._relay(msg)
+            except (ConnectionClosed, ConnectionClosedError):
+                del self._connections[identifier]
+                log(f"Connection from {identifier} closed")
+                continue
+
+
+    async def _server_connection(self, ws: ServerConnection):
+        host, port = ws.remote_address
+        addr = NodeAddr(host, port)
+        identifier = addr.identifier()
+        self._connections[identifier] = ws
+        log(identifier)
+        async for msg in ws:
+            log(msg)
+            self._relay(msg)
+        del self._connections[identifier]
+
+
+    async def run(self):
+        # connect to receivers
+        receiver_tasks = []
+        for r in self.receivers:
+            t = asyncio.create_task(self._receiver_connection(r))
+            receiver_tasks.append(t)
+        # server loop
+        async with serve(self._server_connection, self.listen_address.host, self.listen_address.port) as server:
+            await server.serve_forever()
+        # wait for receivers
+        await asyncio.gather(receiver_tasks)
+
+
+
+def node_from_yml(yml_data: str) -> Node:
+    d = yml_load(yml_data)
+    listen_address = NodeAddr(d["listen"]["host"], d["listen"]["port"])
+    receivers = []
+    for e in d["receivers"]:
+        receivers.append(NodeAddr(e["host"], e["port"]))
+    return Node(listen_address, receivers)
+
+
+if __name__ == "__main__":
+
+    argp = ArgumentParser()
+    argp.add_argument("config", type=Path, help="Path to the config yaml file")
+    args = argp.parse_args()
+
+    node = node_from_yml(args.config.read_text())
+
+    asyncio.run(node.run())