feat: 支持 AList

SilianZ · SilianZ · commit 8704c1481c3e · 2024-12-29T11:39:16.000+08:00
diff --git a/core/__init__.py b/core/__init__.py
@@ -35,7 +35,7 @@ async def syncFiles():
                 Config.get("advanced.retry"),
                 Config.get("advanced.delay"),
             )
-            asyncio.create_task(cluster.recycleFiles())
+            # asyncio.create_task(cluster.recycleFiles())
             if not cluster.enabled and cluster.socket:
                 await cluster.enable()
             if cluster.scheduler:
diff --git a/core/cluster.py b/core/cluster.py
@@ -8,7 +8,7 @@
 from core.logger import logger
 from core.scheduler import *
 from core.exceptions import ClusterIdNotSetError, ClusterSecretNotSetError
-from core.storages import getStorages, LocalStorage
+from core.storages import getStorages, LocalStorage, AListStorage
 from core.classes import FileInfo, FileList, AgentConfiguration
 from core.router import Router
 from core.orm import writeHits
@@ -69,7 +69,7 @@ async def fetchToken(self):
             ).hexdigest()
             response = await session.post(
                 "/openbmclapi-agent/token",
-                data={
+                json={
                     "clusterId": self.id,
                     "challenge": challenge,
                     "signature": signature,
@@ -226,6 +226,7 @@ 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")
 
@@ -271,7 +272,7 @@ async def downloadFile(
     async def report(self, error: ClientResponseError, session: aiohttp.ClientSession) -> None:
         history_urls = [urljoin(self.base_url), *error.history]
         try:
-            async with session.post("/openbmclapi/report", data={"url": history_urls, "error": error.message}) as response:
+            async with session.post("/openbmclapi/report", json={"url": history_urls, "error": error.message}) as response:
                 response.raise_for_status()
                 logger.tdebug("cluster.debug.report", url=history_urls)
         except Exception:
@@ -338,12 +339,12 @@ async def callback(data: List[Any]):
                     "flavor": {
                         "runtime": f"python/{sys.version.split()[0]} python-openbmclapi/{VERSION}",
                         "storage": "+".join(
-                            [
-                                "file"
-                                for storage in self.storages
-                                if isinstance(storage, LocalStorage)
-                            ]
-                        ),
+                        [
+                            "file" if isinstance(storage, LocalStorage) else
+                            "webdav" if isinstance(storage, AListStorage) else
+                            ""
+                            for storage in self.storages
+                        ])
                     },
                 },
                 callback=callback,
diff --git a/core/router.py b/core/router.py
@@ -2,6 +2,7 @@
 from core.config import Config
 from core.utils import checkSign
 from core.api import getStatus
+from core.storages import AListStorage
 from typing import Union
 from aiohttp import web
 import aiohttp
@@ -49,6 +50,12 @@ async def _(request: web.Request) -> web.StreamResponse:
                 ):
                     return web.Response(status=403 if size > 200 else 400)
 
+                response = None
+                for storage in self.storages:
+                    if isinstance(storage, AListStorage):
+                        await storage.measure(size, request, response)
+                        return response
+
                 buffer = b"\x00\x66\xcc\xff" * 256 * 1024
                 response = web.StreamResponse(
                     status=200,
diff --git a/core/storages/__init__.py b/core/storages/__init__.py
@@ -12,5 +12,12 @@ def getStorages() -> List[Storage]:
         if storage["type"] == "local":
             storages.append(LocalStorage(path=storage["path"]))
         if storage["type"] == "alist":
-            storages.append(AListStorage())
+            storages.append(
+                AListStorage(
+                    username=storage["username"],
+                    password=storage["password"],
+                    url=storage["url"],
+                    path=storage["path"],
+                )
+            )
     return storages
diff --git a/core/storages/alist.py b/core/storages/alist.py
@@ -1,6 +1,175 @@
-from dataclasses import dataclass
+from core.classes import Storage, FileInfo, FileList
+from core.scheduler import scheduler, IntervalTrigger
+from core.logger import logger
+from core.i18n import locale
+from typing import List, Set, Tuple, Dict, Any
+from tqdm import tqdm
+from aiohttp import web
+import aiohttp
+import secrets
+import io
+import asyncio
+import humanize
 
 
-class AListStorage(): 
-    def __init__(self) -> None:
+class AListStorage(Storage):
+    def __init__(self, username: str, password: str, url: str, path: str) -> None:
+        self.username = username
+        self.password = password
+        self.url = url
+        self.path = path
+        self.token = ""
+        self.scheduler = None
+        self.headers = {}
+
+    async def init(self) -> None:
+        async def fetchToken() -> None:
+            logger.tinfo("storage.info.alist.fetch_token")
+            async with aiohttp.ClientSession(self.url) as session:
+                try:
+                    async with session.post("/api/auth/login", json={"username": self.username, "password": self.password}) as response:
+                        response.raise_for_status()
+                        data = await response.json()
+                        if data["code"] != 200:
+                            raise aiohttp.ClientResponseError(status=data["code"], request_info=response.request_info, history=response.history)
+                        self.token = data["data"]["token"]
+                        self.headers = {"Authorization": self.token}
+                    logger.tsuccess("storage.success.alist.fetch_token")
+                except Exception as e:
+                    logger.terror("storage.error.alist.fetch_token", e=e)
+            if not self.scheduler:
+                self.scheduler = scheduler.add_job(fetchToken, IntervalTrigger(days=2))
+
+        await fetchToken()
+
+    async def check(self) -> None:
+        file_name = secrets.token_hex(8)
+        file_path = self.path + file_name
+        try:
+            async with aiohttp.ClientSession(self.url, headers={**self.headers, "File-Path": file_path, "Content-Type": "application/octet-stream"}) as session:
+                response = await session.put("/api/fs/put", data=b"")
+                response.raise_for_status()
+                data = await response.json()
+                if data["code"] != 200:
+                    raise aiohttp.ClientResponseError(status=data["code"], request_info=response.request_info, history=response.history)
+            async with aiohttp.ClientSession(self.url, headers=self.headers) as session:
+                response = await session.post("/api/fs/remove", json={"names": [file_name], "dir": self.path})
+                response.raise_for_status()
+                data = await response.json()
+                if data["code"] != 200:
+                    raise aiohttp.ClientResponseError(status=data["code"], request_info=response.request_info, history=response.history)
+            logger.tsuccess("storage.success.alist.check")
+        except Exception as e:
+            logger.terror("storage.error.alist.check", e=e)
+
+    async def getMissingFiles(self, files: FileList, pbar: tqdm) -> FileList:
+        existing_files: List[FileInfo] = []
+        async with aiohttp.ClientSession(self.url, headers=self.headers) as session:
+            async def getFileList(dir: str, pbar: tqdm) -> List[FileInfo]:
+                file_path = self.path + dir
+                response = await session.post("/api/fs/list", headers=self.headers, json={"path": file_path})
+                response.raise_for_status()
+                data = await response.json()
+                pbar.update(1)
+                if data["code"] != 200:
+                    return []
+                return [FileInfo(size=content["size"], hash=content["name"], path="", mtime=-1) for content in data["data"]["content"] if not content["is_dir"]]
+            
+            with tqdm(desc=locale.t("storage.tqdm.alist.get_filelist"), total=256) as pbar:
+                for i in range(256):
+                    dir = f"/{i:02x}"
+                    existing_files += await getFileList(dir, pbar)
+
+        existing_info: Set[Tuple[str, int]] = {(file.hash, file.size) for file in existing_files}
+        missing_files = [file for file in files.files if (file.hash, file.size) not in existing_info]
+
+        return FileList(files=missing_files)
+
+    async def measure(self, size: int, request: web.Request, response):
+        file_path = f"{self.path}/.{size}"
+        try:
+            async with aiohttp.ClientSession(self.url, headers=self.headers) as session:
+                response = await session.post("/api/fs/get", json={"path": file_path, "password": self.password})
+                response.raise_for_status()
+                data = await response.json()
+
+                if data["code"] == 200:
+                    response = web.HTTPFound(data["raw_url"])
+                    response.prepare(request)
+                    return
+
+                if data["code"] != 200:
+                    try:
+                        buffer = b"\x00\x66\xcc\xff" * 256 * 1024
+                        response = await session.put("/api/fs/put", data=buffer, headers={**self.headers, "File-Path": file_path, "Content-Type": "application/octet-stream"})
+                        response.raise_for_status()
+                        data = await response.json()
+                        if data["code"] != 200:
+                            raise aiohttp.ClientResponseError(status=500, request_info=response.request_info, history=response.history)
+                    except Exception as e:
+                        logger.terror("storage.error.alist.upload", e=e)
+                        raise
+
+                response = await session.post("/api/fs/get", json={"path": file_path, "password": self.password})
+                response.raise_for_status()
+                data = await response.json()
+
+                response = web.HTTPFound(data["raw_url"])
+                response.prepare(request)
+                return
+        except Exception as e:
+            logger.terror("storage.error.alist.measure", e=e)
+            return
+
+    async def express(self, hash: str, request: web.Request, response) -> Dict[str, Any]:
+        path = f"{self.path}/{hash[:2]}/{hash}"
+        async with aiohttp.ClientSession(self.url, headers=self.headers) as session:
+            res = await session.post("/api/fs/get", json={"path": path, "password": self.password})
+            data = await res.json()
+            if data["code"] != 200:
+                response = web.HTTPNotFound()
+                await response.prepare(request)
+                return {"bytes": 0, "hits": 0}
+            try:
+                response = web.HTTPFound(data["raw_url"])
+                response.headers["x-bmclapi-hash"] = hash
+                await response.prepare(request)
+                return {"bytes": data["size"], "hits": 1}
+            except Exception as e:
+                response = web.HTTPInternalServerError(reason=e)
+                await response.prepare(request)
+                logger.debug(e)
+                return {"bytes": 0, "hits": 0}
+
+    async def writeFile(self, file: FileInfo, content: io.BytesIO, delay: int, retry: int) -> bool:
+        file_path = f"{self.path}/{file.hash[:2]}/{file.hash}"
+        async def getSize() -> int:
+            async with aiohttp.ClientSession(self.url, headers=self.headers) as session:
+                response = await session.post("/api/fs/get", json={"path": file_path, "password": self.password})
+                response.raise_for_status()
+                data = await response.json()
+                if data["code"] != 200:
+                    raise aiohttp.ClientResponseError(status=data["code"], request_info=response.request_info, history=response.history)
+                return data["data"]["size"]
+
+        for _ in range(retry):
+            try:
+                async with aiohttp.ClientSession(self.url, headers={**self.headers, "File-Path": file_path, "Content-Type": "application/octet-stream"}) as session:
+                    response = await session.put("/api/fs/put", data=content.getvalue())
+                    response.raise_for_status()
+                    data = await response.json()
+                    if data["code"] != 200:
+                        raise aiohttp.ClientResponseError(status=data["code"], request_info=response.request_info, history=response.history)
+                    size = await getSize()
+                    if size == file.size:
+                        return True
+                    else:
+                        logger.terror("storage.error.alist.write_file.size_mismatch", file=file.hash, file_size=humanize.naturalsize(file.size, binary=True), actual_file_size=humanize.naturalsize(size, binary=True))
+                        return False
+            except Exception as e:
+                logger.terror("storage.error.alist.write_file.retry", file=file.hash, e=e, retry=delay)
+            await asyncio.sleep(delay)
+        return False
+
+    async def recycleFiles(files):
         pass
diff --git a/core/storages/local.py b/core/storages/local.py
@@ -13,7 +13,6 @@
 import humanize
 
 
-
 class LocalStorage(Storage):
     def __init__(self, path: str) -> None:
         self.path = path
@@ -22,20 +21,23 @@ async def init(self) -> None:
         os.makedirs(self.path, exist_ok=True)
 
     async def check(self) -> None:
-        logger.tinfo("storage.info.check")
+        logger.tinfo("storage.info.local.check")
         try:
             with tempfile.NamedTemporaryFile(dir=self.path, delete=True) as temp_file:
                 temp_file.write(b"")
-            logger.tsuccess("storage.success.check")
+            logger.tsuccess("storage.success.local.check")
         except Exception as e:
-            raise Exception(locale.t("storage.error.check", e=e))
+            raise Exception(locale.t("storage.error.local.check", e=e))
 
     async def writeFile(
         self, file: FileInfo, content: io.BytesIO, delay: int, retry: int
     ) -> bool:
         file_path = os.path.join(self.path, file.hash[:2], file.hash)
         os.makedirs(os.path.dirname(file_path), exist_ok=True)
-
+        if Path(file_path).exists and Path(file_path).stat().st_size == len(
+            content.getvalue()
+        ):
+            return True
         for _ in range(retry):
             try:
                 async with aiofiles.open(file_path, "wb") as f:
@@ -46,23 +48,27 @@ async def writeFile(
                     return True
                 else:
                     logger.terror(
-                        "storage.error.write_file.size_mismatch", file=file.hash,
+                        "storage.error.local.write_file.size_mismatch",
+                        file=file.hash,
                         file_size=humanize.naturalsize(file.size, binary=True),
-                        actual_file_size=humanize.naturalsize(size, binary=True)
+                        actual_file_size=humanize.naturalsize(size, binary=True),
                     )
                     return False
             except Exception as e:
                 logger.terror(
-                    "storage.error.write_file.retry", file=file.hash, e=e, retry=delay
+                    "storage.error.local.write_file.retry",
+                    file=file.hash,
+                    e=e,
+                    retry=delay,
                 )
 
             await asyncio.sleep(delay)
 
-        logger.terror("storage.error.write_file.failed", file=file.hash)
+        logger.terror("storage.error.local.write_file.failed", file=file.hash)
         return False
 
     async def getMissingFiles(self, files: FileList, pbar: tqdm) -> FileList:
-        async def check_file(file: FileInfo, pbar: tqdm) -> bool:
+        async def checkFile(file: FileInfo, pbar: tqdm) -> bool:
             pbar.update(1)
             file_path = os.path.join(self.path, file.hash[:2], file.hash)
             try:
@@ -71,9 +77,7 @@ async def check_file(file: FileInfo, pbar: tqdm) -> bool:
             except FileNotFoundError:
                 return True
 
-        results = await asyncio.gather(
-            *[check_file(file, pbar) for file in files.files]
-        )
+        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
         ]
@@ -84,8 +88,8 @@ async def express(
     ) -> Dict[str, Any]:
         path = os.path.join(self.path, hash[:2], hash)
         if not os.path.exists(path):
-            response = web.Response()
-            response.set_status(404, "File not found")
+            response = web.HTTPNotFound()
+            await response.prepare(request)
             return {"bytes": 0, "hits": 0}
         try:
             file_size = os.path.getsize(path)
@@ -117,9 +121,9 @@ async def recycleFiles(self, files: FileList) -> None:
                     delete_files.append(file)
 
         if len(delete_files) == 0:
-            logger.tinfo("storage.success.no_need_to_recycle")
+            logger.tinfo("storage.success.local.no_need_to_recycle")
             return
-        
+
         with tqdm(
             desc=locale.t("storage.tqdm.desc.recycling"),
             total=len(delete_files),
@@ -133,8 +137,9 @@ async def recycleFiles(self, files: FileList) -> None:
                 try:
                     file.unlink()
                 except Exception as e:
-                    logger.terror("storage.error.recycle", e=e)
-
-            logger.tsuccess("storage.success.recycled", size=humanize.naturalsize(size, binary=True))
+                    logger.terror("storage.error.local.recycle", e=e)
 
-        
+            logger.tsuccess(
+                "storage.success.local.recycled",
+                size=humanize.naturalsize(size, binary=True),
+            )
diff --git a/i18n/zh_cn.json b/i18n/zh_cn.json

Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,7 @@ async def syncFiles():`
`35`	`35`	`Config.get("advanced.retry"),`
`36`	`36`	`Config.get("advanced.delay"),`
`37`	`37`	`)`
`38`		`- asyncio.create_task(cluster.recycleFiles())`
	`38`	`+ # asyncio.create_task(cluster.recycleFiles())`
`39`	`39`	`if not cluster.enabled and cluster.socket:`
`40`	`40`	`await cluster.enable()`
`41`	`41`	`if cluster.scheduler:`