Improvement

moon/moon.py

@@ -0,0 +1,79 @@
+import json
+from loguru import logger
+import yaml
+from websocket_server import WebsocketServer
+import os
+
+# 设置当前工作目录
+os.chdir(os.path.dirname(os.path.abspath(__file__)))
+
+# 配置日志
+logger.add("logs.log", rotation="1 MB", level="INFO", format="{time} - {level} - {message}")
+
+moon_addr = None
+port = None
+
+with open('moon_config.yaml', 'r') as config_file:
+    config = yaml.safe_load(config_file)
+    moon_addr = config.get("addr", "127.0.0.1")
+    port = config.get("port", 8765)
+
+clients = dict()  # id -> name
+clients_rev = dict()  # name -> id
+clients_connection = dict()  # id -> client obj
+
+def new_client(client, server):
+    clients_connection[client["id"]] = client
+    logger.info(f"新客户端连接: {client['id']}")
+
+def client_left(client, server):
+    logger.info(f"客户端断开连接: {client['id']}")
+    # 连接关闭时从集合中移除
+    if client in clients.values():
+        clients.pop(clients_rev[client['id']])
+        clients_rev.pop(client['id'])
+
+def message_received(client, server, message):
+    logger.info(f"处理信息: {message}")
+    try:
+        msg = json.loads(message)
+        clients[msg["id"]] = client['id']
+        clients_rev[client['id']] = msg["id"]
+        
+        if msg["receiver"] == "auth":
+            return
+        
+        if msg["receiver"] == "*":
+            try:
+                for i in clients_rev.values():
+                    content = msg["content"]
+                    content["emitter"] = msg["id"]
+                    server.send_message(
+                        clients_connection[i], 
+                        json.dumps(content))
+            except KeyError:
+                logger.warning("无法发送至端点")
+        
+        try:
+            content = msg["content"]
+            content["emitter"] = msg["id"]
+            server.send_message(
+                clients_connection[clients[msg["receiver"]]], 
+                json.dumps(content))
+        except KeyError:
+            logger.warning("无法发送至端点")
+    except json.JSONDecodeError:
+        logger.error("无法解码 JSON 消息")
+
+def main():
+    global moon_addr
+    global port
+    server = WebsocketServer(host=moon_addr, port=port)
+    server.set_fn_new_client(new_client)
+    server.set_fn_client_left(client_left)
+    server.set_fn_message_received(message_received)
+    logger.info(f'启动"卫星"中继服务器: ws://{moon_addr}:{port}')
+    server.run_forever()
+
+if __name__ == "__main__":
+    main()