refactor: 类型批注

Author: SilianZ

core/__init__.py

@@ -9,7 +9,7 @@
 cluster = Cluster()
 
 
-async def main():
+async def main() -> None:
     try:
         await cluster.token.fetchToken()
         await cluster.getConfiguration()
@@ -23,7 +23,7 @@ async def main():
         except Exception as e:
             logger.terror("orm.error.failed", e=e)
 
-        async def syncFiles():
+        async def syncFiles() -> None:
             if cluster.scheduler:
                 cluster.scheduler.pause()
             if cluster.enabled and cluster.socket:

core/api.py

@@ -37,6 +37,3 @@ async def getStatus(cluster) -> web.Response:
         "version": VERSION,
     }
     return web.json_response(data=response)
-
-async def postConfig() -> web.Response:
-    pass

core/classes.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
-from typing import List, Dict, Any
+from typing import List
 from abc import ABC, abstractmethod
-import io
 from aiohttp import web
-import tqdm
+from tqdm import tqdm
+from typing import Union
+import io
 
 
 @dataclass
@@ -31,25 +32,25 @@ async def init(self) -> None:
         pass
 
     @abstractmethod
-    async def check(self) -> bool:
+    async def check(self) -> None:
         pass
 
     @abstractmethod
     async def writeFile(
-        self, path: str, content: io.BytesIO, delay: int, retry: int
-    ) -> int:
+        self, file: FileInfo, content: io.BytesIO, delay: int, retry: int
+    ) -> bool:
         pass
 
     @abstractmethod
-    async def getMissingFiles(files: FileList, pbar: tqdm) -> FileList:
+    async def getMissingFiles(self, files: FileList, pbar: tqdm) -> FileList:
         pass
 
     @abstractmethod
-    async def express(
-        hash: str, counters: dict
-    ) -> web.Response:
+    async def express(self,
+        hash: str, counter: dict
+    ) -> Union[web.Response, web.FileResponse]:
         pass
 
     @abstractmethod
-    async def recycleFiles(files: FileList) -> None:
+    async def recycleFiles(self, files: FileList) -> None:
         pass

core/cluster.py

@@ -13,7 +13,7 @@
 from core.router import Router
 from core.orm import writeHits
 from core.i18n import locale
-from typing import List, Any
+from typing import List, Any, Union
 from aiohttp import web, ClientResponseError
 from urllib.parse import urljoin
 from tqdm import tqdm
@@ -53,7 +53,7 @@ def __init__(self) -> None:
         if not self.id or not self.secret:
             raise ClusterIdNotSetError if not self.id else ClusterSecretNotSetError
 
-    async def fetchToken(self):
+    async def fetchToken(self) -> None:
         logger.tinfo("token.info.fetching")
         async with aiohttp.ClientSession(
             self.base_url, headers={"User-Agent": self.user_agent}
@@ -100,15 +100,15 @@ def __init__(self) -> None:
         self.filelist = FileList(files=[])
         self.storages = getStorages()
         self.configuration = None
-        self.semaphore = None
-        self.socket = None
-        self.router = None
+        self.semaphore = asyncio.Semaphore()
+        self.socket = socketio.AsyncClient(handle_sigint=False)
+        self.router: Router | None = None
         self.runner = None
         self.failed_filelist = FileList(files=[])
         self.enabled = False
         self.site = None
         self.want_enable = False
-        self.scheduler = None
+        self.scheduler = scheduler
         self.start_time = int(time.time() * 1000)
 
     async def fetchFileList(self) -> None:
@@ -227,7 +227,6 @@ async def recycleFiles(self) -> None:
     async def downloadFile(
         self, file: FileInfo, session: aiohttp.ClientSession, pbar: tqdm
     ) -> None:
-        # logger.debug(file)
         async with self.semaphore:
             delay, retry = Config.get("advanced.delay"), Config.get("advanced.retry")
 
@@ -255,7 +254,7 @@ async def downloadFile(
                         e=e.message,
                         retry=delay,
                     )
-                    self.report(e, session)
+                    await self.report(e, session)
 
                 except Exception as e:
                     logger.terror(
@@ -273,7 +272,7 @@ async def downloadFile(
     async def report(
         self, error: ClientResponseError, session: aiohttp.ClientSession
     ) -> None:
-        history_urls = [urljoin(self.base_url), *error.history]
+        history_urls = [*error.history]
         try:
             async with session.post(
                 "/openbmclapi/report",
@@ -402,6 +401,10 @@ async def keepAlive(self) -> bool:
             logger.terror("cluster.error.keep_alive.socket_not_setup")
             return False
 
+        if self.router is None:
+            logger.terror("cluster.error.keep_alive.router_not_setup")
+            return False
+
         future = asyncio.Future()
 
         async def callback(data: List[Any]):
@@ -482,28 +485,22 @@ async def callback(data: List[Any]):
     async def connect(self) -> None:
         if self.socket and self.socket.connected:
             return
-
-        self.socket = socketio.AsyncClient(handle_sigint=False)
-
-        @self.socket.on("connect")
-        async def _() -> None:
+        async def onConnect() -> None:
             logger.tsuccess("client.success.connected")
             await self.disable()
             if self.want_enable:
                 await self.enable()
             if self.scheduler:
                 self.scheduler.resume()
-
-        @self.socket.on("disconnect")
-        async def _() -> None:
+        async def onDisconnect() -> None:
             logger.twarning("client.warn.disconnected")
             if self.scheduler:
                 self.scheduler.pause()
-
-        @self.socket.on("message")
-        async def _(message: str) -> None:
+        async def onMessage(message: str) -> None:
             logger.tinfo("client.info.message", message=message)
-
+        self.socket.on("connect", onConnect)
+        self.socket.on("disconnect", onDisconnect)
+        self.socket.on("message", onMessage)
         await self.socket.connect(
             self.base_url,
             transports=["websocket"],

core/logger.py

@@ -1,4 +1,4 @@
-from loguru import logger as Logger
+from loguru import logger as Logger, Record
 from pathlib import Path
 import sys
 from core.config import Config
@@ -8,14 +8,14 @@
 debug_mode = Config.get("advanced.debug")
 
 
-def filter(record):
-    if "apscheduler" in record["name"]:
-        record["depth"] = 2
-    return 1
+def filter(record: Record) -> bool:
+    if record["name"] and "apscheduler" in record["name"]:
+        record["extra"] = {"depth": 2}
+    return True
 
 
 class LoggingLogger:
-    def __init__(self):
+    def __init__(self) -> None:
         self.log = Logger.opt(depth=1)
         self.log.remove()
         self.log.add(

core/router.py

@@ -1,10 +1,14 @@
 from core.orm import writeAgent
-from core.utils import checkSign
 from core.api import getStatus
 from core.storages import AListStorage
 from core.logger import logger
 from aiohttp import web
+from typing import Union
+from multidict import MultiMapping
 import aiohttp
+import base64
+import hashlib
+import time
 import random
 
 
@@ -19,6 +23,18 @@ def __init__(self, app: web.Application, cluster) -> None:
         self.ws_clients = []
         self.connection = 0
 
+    def checkSign(self, hash: str, secret: str, query: MultiMapping) -> 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()).digest()
+            )
+            .decode()
+            .rstrip("=")
+        )
+        return sign == s and time.time() < int(e, 36)
+
     def init(self) -> None:
         @self.route.get("/download/{hash}")
         async def _(
@@ -27,7 +43,7 @@ async def _(
             self.connection += 1
             writeAgent(request.headers["User-Agent"], 1)
             file_hash = request.match_info.get("hash", "").lower()
-            if not checkSign(file_hash, self.secret, request.query):
+            if not self.checkSign(file_hash, self.secret, request.query):
                 return web.Response(text="Invalid signature.", status=403)
 
             response = await random.choice(self.storages).express(
@@ -39,15 +55,15 @@ async def _(
             return response
 
         @self.route.get("/measure/{size}")
-        async def _(request: web.Request) -> web.Response:
+        async def _(request: web.Request) -> Union[web.Response, web.StreamResponse]:
             try:
                 size = int(request.match_info.get("size", "0"))
 
                 if (
-                    not checkSign(f"/measure/{size}", self.secret, request.query)
+                    not self.checkSign(f"/measure/{size}", self.secret, request.query)
                     or size > 200
                 ):
-                    return web.Response(status=403 if size > 200 else 400)
+                    return web.HTTPForbidden() if size > 200 else web.HTTPBadRequest()
 
                 response = None
                 for storage in self.storages:

core/storages/alist.py

@@ -189,7 +189,7 @@ async def express(self, hash: str, counter: dict) -> web.Response:
                 counter["hits"] += 1
                 return response
             except Exception as e:
-                response = web.HTTPError(text=e)
+                response = web.HTTPError(text=str(e))
                 logger.debug(e)
                 return response
 
@@ -254,5 +254,5 @@ async def getSize() -> int:
         logger.terror("storage.error.alist.write_file.failed", file=file.hash)
         return False
 
-    async def recycleFiles(files):
+    async def recycleFiles(self, files) -> None:
         pass

core/storages/local.py

@@ -2,7 +2,7 @@
 from core.logger import logger
 from core.i18n import locale
 from aiohttp import web
-from typing import Any, Dict
+from typing import Union
 from tqdm import tqdm
 from pathlib import Path
 import os
@@ -83,7 +83,7 @@ async def checkFile(file: FileInfo, pbar: tqdm) -> bool:
         ]
         return FileList(files=missing_files)
 
-    async def express(self, hash: str, counter: dict) -> web.Response:
+    async def express(self, hash: str, counter: dict) -> Union[web.Response, web.FileResponse]:
         path = os.path.join(self.path, hash[:2], hash)
         if not os.path.exists(path):
             response = web.HTTPNotFound()
@@ -96,7 +96,7 @@ async def express(self, hash: str, counter: dict) -> web.Response:
             counter["hits"] += 1
             return response
         except Exception as e:
-            response = web.HTTPError(text=e)
+            response = web.HTTPError(text=str(e))
             logger.debug(e)
             return response
 

core/storages/s3.py

@@ -2,7 +2,7 @@
 from core.logger import logger
 from botocore.config import Config
 from botocore.exceptions import ClientError
-from typing import Literal, Union
+from typing import Literal, Union, Self
 from tqdm import tqdm
 import boto3
 import humanize
@@ -86,18 +86,18 @@ async def check(self) -> None:
             logger.terror("storage.error.s3.check", e=e)
 
     async def getMissingFiles(self, files: FileList, pbar: tqdm) -> FileList:
-        async def checkFile(file: FileInfo, pbar: tqdm) -> bool:
-            pbar.update(1)
-            file_path = f"{file.hash[:2]}/{file.hash}"
-            try:
-                response = self.client.head_object(Bucket=self.bucket, Key=file_path)
-                s3_file_size = response["ContentLength"]
-                return s3_file_size != file.size
-            except Exception:
-                pass
+        s3_files = {}
+        try:
+            paginator = self.client.get_paginator("list_objects_v2")
+            async for page in paginator.paginate(Bucket=self.bucket):
+                for obj in page.get("Contents", []):
+                    s3_files[obj["Key"]] = obj["Size"]
+        except ClientError as e:
+            logger.terror("storage.error.s3.get_s3_files", e=e)
 
-        results = await asyncio.gather(*[checkFile(file, pbar) for file in files.files])
-        missing_files = [
-            file for file, is_missing in zip(files.files, results) if is_missing
-        ]
-        return FileList(files=missing_files)
+        missing_files = []
+        for file in files.files:
+            file_key = f"{file.hash[:2]}/{file.hash}"
+            if file_key not in s3_files or s3_files[file_key] != file.size:
+                missing_files.append(file)
+            pbar.update(1)