feat: 支持 download 路由并适配一系列动作

- 支持 WebSocket 连接
- 支持证书获取
- 使用 aiohttp 来服务文件

Author: SilianZ

.gitignore

@@ -1,6 +1,6 @@
 .vscode/
-config/config.properties
 **/__pycache__
 logs
 config/config.yml
-cache
+cache
+cert

core/__init__.py

@@ -1,5 +1,6 @@
 from core.cluster import Cluster
 from core.config import Config
+import asyncio
 
 
 async def init():
@@ -13,3 +14,13 @@ async def init():
     delay = Config.get("advanced.delay")
     retry = Config.get("advanced.retry")
     await cluster.syncFiles(missing_filelist, retry, delay)
+    await cluster.connect()
+    protocol = "http" if Config.get("cluster.byoc") else "https"
+    if protocol == "https":
+        await cluster.socket.requestCertificate()
+    await cluster.setupExpress(protocol == "https", port=Config.get('cluster.port'))
+    try:
+        while True:
+            await asyncio.sleep(1000)
+    except KeyboardInterrupt:
+        raise KeyboardInterrupt

core/client.py

@@ -0,0 +1,73 @@
+from core.logger import logger
+from core.config import Config
+from typing import List, Any
+import socketio
+import asyncio
+import aiofiles
+import os
+
+
+class WebSocketClient:
+    def __init__(self, token: str) -> None:
+        self.socket = None
+        self.connected = None
+        self.base_url = Config.get("cluster.base_url")
+        self.cert_path = Config.get("advanced.paths.cert")
+        self.key_path = Config.get("advanced.paths.key")
+        self.token = token
+        os.makedirs(os.path.dirname(self.cert_path), exist_ok=True)
+        os.makedirs(os.path.dirname(self.key_path), exist_ok=True)
+
+    async def connect(self) -> None:
+        if self.socket and self.connected:
+            return
+
+        self.socket = socketio.AsyncClient()
+
+        @self.socket.on("connect")
+        async def _() -> None:
+            logger.tsuccess("client.success.connected")
+
+        @self.socket.on("message")
+        async def _(message: str) -> None:
+            logger.tinfo("client.info.message", message=message)
+
+        @self.socket.on("exception")
+        async def _(error: str) -> None:
+            logger.tinfo("client.error.exception", error=error)
+
+        @self.socket.on("reconnect")
+        async def _() -> None:
+            pass
+
+        @self.socket.on("reconnect_error")
+        async def _(error: str) -> None:
+            logger.terror("client.error.reconnect", error=error)
+
+        @self.socket.on("reconnect_failed")
+        async def _() -> None:
+            pass
+
+        await self.socket.connect(
+            self.base_url, transports=["websocket"], auth={"token": str(self.token)}
+        )
+
+    async def requestCertificate(self) -> None:
+        future = asyncio.Future()
+
+        async def callback(data: List[Any]):
+            error, cert = data
+            future.set_result((error, cert))
+
+        try:
+            await self.socket.emit("request-cert", callback=callback)
+            error, cert = await future
+            if error:
+                raise Exception(error)
+            async with aiofiles.open(self.cert_path, "w") as f:
+                await f.write(cert["cert"])
+            async with aiofiles.open(self.key_path, "w") as f:
+                await f.write(cert["key"])
+            logger.tsuccess("client.success.request_certificate")
+        except Exception as e:
+            logger.terror("client.error.request_certificate", e=e)

core/cluster.py

@@ -5,6 +5,7 @@
 from core.storages import getStorages
 from core.classes import FileInfo, FileList, AgentConfiguration
 from core.router import Router
+from core.client import WebSocketClient
 from core.i18n import locale
 from aiohttp import web
 from tqdm import tqdm
@@ -14,7 +15,7 @@
 import asyncio
 import hmac
 import hashlib
-import socketio
+import ssl
 import humanize
 import io
 
@@ -82,6 +83,9 @@ def __init__(self) -> None:
         self.storages = getStorages()
         self.configuration = None
         self.semaphore = None
+        self.socket = None
+        self.router = None
+        self.server = None
         self.failed_filelist = FileList(files=[])
 
     async def fetchFileList(self) -> None:
@@ -214,7 +218,9 @@ async def downloadFile(
                         content = await response.read()
                         results = await asyncio.gather(
                             *(
-                                storage.writeFile(file, io.BytesIO(content), delay, retry)
+                                storage.writeFile(
+                                    file, io.BytesIO(content), delay, retry
+                                )
                                 for storage in self.storages
                             )
                         )
@@ -226,18 +232,37 @@ async def downloadFile(
                         "cluster.error.download_file.retry",
                         file=file.hash,
                         e=e,
-                        retry=delay
+                        retry=delay,
                     )
                 await asyncio.sleep(delay)
             logger.terror("cluster.error.download_file.failed", file=file.hash)
             self.failed_filelist.files.append(file)
 
-    async def setupExpress(self, https: bool) -> None:
+    async def setupExpress(self, https: bool, port: int) -> None:
         logger.tinfo("cluster.info.router.creating")
-        app = web.Application
-        router = Router(https, app)
-        router.init()
+        app = web.Application()
+        try:
+            self.router = Router(app, self.storages)
+            self.router.init()
+            ssl_context = None
+            # if https:
+            #     ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
+            #     ssl_context.load_cert_chain(
+            #         certfile=Config.get("advanced.paths.cert"),
+            #         keyfile=Config.get("advanced.paths.key")
+            #     )
+            self.server = web.AppRunner(app)
+            logger.tsuccess("cluster.success.router.created")
+            await self.server.setup()
+            site = web.TCPSite(self.server, "0.0.0.0", port, ssl_context=ssl_context)
+            await site.start()
+            logger.tsuccess("cluster.success.site.start")
+        except Exception as e:
+            logger.terror("cluster.error.router.exception", e=e)
 
+    async def connect(self) -> None:
+        self.socket = WebSocketClient(self.token.token)
+        await self.socket.connect()
 
     async def init(self) -> None:
         await asyncio.gather(*(storage.init() for storage in self.storages))

core/config.py

@@ -12,7 +12,12 @@
     "cluster.base_url": "https://openbmclapi.bangbang93.com",
     "cluster.id": "",
     "cluster.secret": "",
+    "cluster.byoc": False,
+    "cluster.public_port": 8080,
+    "cluster.port": 8800,
     "storages": [{"type": "local", "path": "./cache"}],
+    "advanced.paths.cert": "./cert/cert.pem",
+    "advanced.paths.key": "./cert/key.pem",
 }
 
 

core/router.py

@@ -1,49 +1,59 @@
 from functools import wraps
 from core.config import Config
 from core.classes import Storage
-from typing import Any, List
+from core.logger import logger
+from typing import List
 from aiohttp import web
 import random
-import aiohttp
-import asyncio
 import base64
 import hashlib
 import time
 
+
 class Router:
-    def __init__(self, https: bool, app: web.Application) -> None:
-        self.https = https
+    def __init__(self, app: web.Application, storages: List[Storage]) -> None:
         self.app = app
-        self.secret = Config.get('cluster.secret')
+        self.secret = Config.get("cluster.secret")
+        self.storages = storages
 
     def route(self, path, method="GET"):
         def decorator(func):
             @wraps(func)
             async def wrapper(request):
                 return await func(request)
+
             self.app.router.add_route(method, path, wrapper)
             return wrapper
+
         return decorator
-    
+
     def init(self) -> None:
-        @self.route("/auth")
-        async def _():
-            pass
+        # @self.route("/auth")
+        # async def _():
+        #     pass
 
         @self.route("/download/{hash}")
-        async def _(self, request: web.Request, storages: List[Storage]) -> web.Response | web.StreamResponse:
+        async def _(
+            request: web.Request,
+        ) -> web.Response | web.FileResponse | web.StreamResponse:
             def check_sign(hash: str, secret: str, query: dict) -> bool:
-                if not (s := query.get('s')) or not (e := query.get('e')): return False
-                sign = base64.urlsafe_b64encode(hashlib.sha1(f"{secret}{hash}{e}".encode('utf-8')).digest()).decode('utf-8').rstrip('=')
+                # return True
+                if not (s := query.get("s")) or not (e := query.get("e")):
+                    return False
+                sign = (
+                    base64.urlsafe_b64encode(
+                        hashlib.sha1(f"{secret}{hash}{e}".encode("utf-8")).digest()
+                    )
+                    .decode("utf-8")
+                    .rstrip("=")
+                )
                 return sign == s and time.time() < int(e, 36)
-            
-            hash = request.match_info.get('hash').lower()
+
+            hash = request.match_info.get("hash").lower()
             valid = check_sign(hash, self.secret, request.query)
             if not valid:
                 return web.Response(text="invalid sign", status=403)
-            response = web.StreamResponse(status=200)
-            response.headers['x-bmclapi-hash'] = hash
-            data = await random.choice(storages).express(hash, request, response)
+            response = None
+            data = await random.choice(self.storages).express(hash, request, response)
+            logger.debug(data)
             return response
-
-            

core/storages/local.py

@@ -17,8 +17,7 @@ def __init__(self, path: str) -> None:
         self.checked = False
 
     async def init(self) -> None:
-        if not os.path.exists(self.path):
-            os.makedirs(self.path)
+        os.makedirs(self.path, exist_ok=True)
 
     async def check(self) -> None:
         logger.tinfo("storage.info.check")
@@ -77,24 +76,21 @@ async def check_file(file: FileInfo, pbar: tqdm) -> bool:
             file for file, is_missing in zip(files.files, results) if is_missing
         ]
         return FileList(files=missing_files)
-    
-    async def express(self, hash: str, request: web.Request, response: web.StreamResponse) -> Dict[str, Any]:
+
+    async def express(
+        self, hash: str, request: web.Request, response
+    ) -> Dict[str, Any]:
         path = os.path.join(self.path, hash[:2], hash)
         if not os.path.exists(path):
-            response.set_status(404, 'File not found')
-            return {'bytes': 0, 'hits': 0}
-        file_size = os.path.getsize(path)
-        response.content_length = file_size
-        response.content_type = 'application/octet-stream'
-        response.headers['Cache-Control'] = 'max-age=2592000'
-        response.prepare(request)
-
+            response = web.Response()
+            response.set_status(404, "File not found")
+            return {"bytes": 0, "hits": 0}
         try:
-            transport = request.transport
-            socket = transport.get_extra_info('socket')
-            with open(path, 'rb') as f:
-                await asyncio.get_event_loop().sendfile(socket, f, 0, file_size)
-            return {'bytes': file_size, 'hits': 1}
-
-        except Exception:
-            return {'bytes': 0, 'hits': 0}
+            file_size = (os.path.getsize(path),)
+            response = web.FileResponse(path, status=200)
+            response.headers["x-bmclapi-hash"] = hash
+            await response.prepare(request)
+            return {"bytes": file_size, "hits": 1}
+        except Exception as e:
+            logger.debug(e)
+            return {"bytes": 0, "hits": 0}

i18n/zh_cn.json

@@ -21,8 +21,17 @@
     "cluster.error.sync_files.retry": "无法下载所有文件，将在 ${retry}s 后重试。",
     "cluster.error.sync_files.failed": "无法下载所有文件，已达到最高重试次数。",
     "cluster.info.router.creating": "正在创建路由……",
+    "cluster.success.router.created": "成功创建路由！",
+    "cluster.error.router.exception": "在尝试创建路由时发生错误：${e}。",
     "cluster.tqdm.desc.get_missing": "获取缺失文件中",
     "cluster.tqdm.desc.sync_files": "同步文件中",
     "cluster.tqdm.unit.files": " 个文件",
+    "client.success.connected": "成功与主控建立连接！",
+    "client.info.message": "[主控] ${message}",
+    "client.error.exception": "[主控] ${error}",
+    "client.error.reconnect": "在尝试重连至主控时发生错误：${error}。",
+    "client.error.request_certificate": "在尝试获取证书时发生错误：${e}。",
+    "client.success.request_certificate": "成功获取证书！",
+    "cluster.success.site.start": "成功启动服务端！",
     "configuration.debug.get": "同步策略：${sync}。"
 }