perf: 优化文件后缀处理逻辑

source/application/download.py

@@ -1,5 +1,6 @@
 from asyncio import Semaphore
 from asyncio import gather
+from os import path
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -12,6 +13,7 @@ from source.module import Manager
 from source.module import logging
 from source.module import retry as re_download
 from source.module import sleep_time
+from source.module.static import FILE_HEADER_MAX_LENGTH, MAGIC_DICT
 
 if TYPE_CHECKING:
     from httpx import AsyncClient
@@ -162,7 +164,6 @@ class Download:
                 # )
                 return False
             temp = self.temp.joinpath(f"{name}.{suffix}")
-            real = path.joinpath(f"{name}.{suffix}")
             self.__update_headers_range(headers, temp, )
             try:
                 async with self.client.stream("GET", url, headers=headers, ) as response:
@@ -178,6 +179,7 @@ class Download:
                         async for chunk in response.aiter_bytes(self.chunk):
                             await f.write(chunk)
                             # self.__update_progress(bar, len(chunk))
+                real = await self.__fix_suffix(temp, path, name, suffix, log)
                 self.manager.move(temp, real)
                 # self.__create_progress(bar, None)
                 logging(log, self.message("文件 {0} 下载成功").format(real.name))
@@ -218,7 +220,7 @@ class Download:
             headers: dict[str, str],
             suffix: str,
             # sleep_args: tuple[int, int],
-    ) -> [int, str]:
+    ) -> tuple[int, str]:
         response = await self.client.head(
             url,
             headers=headers,
@@ -242,3 +244,26 @@ class Download:
     ) -> int:
         headers["Range"] = f"bytes={(p := self.__get_resume_byte_position(file))}-"
         return p
+
+    @staticmethod
+    async def __fix_suffix(
+        temp: Path,
+        path: Path,
+        name: str,
+        suffix: str,
+        log
+    ) -> Path:
+        try:
+            async with open(temp, "rb") as f:
+                file_start = await f.read(FILE_HEADER_MAX_LENGTH)
+            for offset, magic, suffix in MAGIC_DICT:
+                if file_start[offset:offset + len(magic)] == magic:
+                    return path.joinpath(f"{name}.{suffix}")
+            return path.joinpath(f"{name}.{suffix}")
+        except Exception as error:
+            logging(
+                log,
+                f"文件 {temp.name} 后缀修复失败，错误信息: {repr(error)}",
+                ERROR,
+            )
+            return path.joinpath(f"{name}.{suffix}")

source/module/static.py

@@ -66,3 +66,14 @@ PROGRESS = "b bright_magenta"
 ERROR = "b bright_red"
 WARNING = "b bright_yellow"
 INFO = "b bright_green"
+
+MAGIC_DICT = {
+    # 分别为偏移量(字节)、魔数、后缀名
+    # 参考：https://en.wikipedia.org/wiki/List_of_file_signatures
+    (0, b"\xFF\xD8\xFF", "jpg"),
+    (0, b"\x89\x50\x4E\x47\x0D\x0A\x1A\x0A", "png"),
+    (0, b"\x00\x00\x00", "avif"),
+    (4, b"\x66\x74\x79\x70\x68\x65\x69\x63", "heic"),
+    (8, b"\x57\x45\x42\x50", "webp"),
+}
+FILE_HEADER_MAX_LENGTH = max(offset + len(magic) for offset, magic, _ in MAGIC_DICT)