From 7665d295ee6ef00d9103f189164f1b14b89ef54a Mon Sep 17 00:00:00 2001 From: chenyanggao <2339083510@qq.com> Date: Sun, 7 Jul 2024 00:43:55 +0800 Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E5=86=99=20AutoFilm/app/modules/alist?= =?UTF-8?q?2strm.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- app/modules/alist2strm.py | 373 ++++++++++++++++++-------------------- requirements.txt | 3 +- 3 files changed, 184 insertions(+), 195 deletions(-) diff --git a/.gitignore b/.gitignore index 08d9163..2519d88 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ strm/* media/* __pycache__/* +.DS_Store *bak *test* -*pyc* \ No newline at end of file +*pyc* diff --git a/app/modules/alist2strm.py b/app/modules/alist2strm.py index 790218c..c3a66ec 100644 --- a/app/modules/alist2strm.py +++ b/app/modules/alist2strm.py @@ -1,222 +1,209 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + import logging -import hmac -import hashlib -import base64 -import asyncio -from urllib.parse import unquote -from requests import Session -from aiohttp import ClientSession as AsyncSession -from pathlib import Path -from typing import Optional -from alist import AlistFileSystem, AlistPath +from asyncio import run, to_thread, Semaphore, TaskGroup +from contextlib import closing, aclosing +from os import fsdecode, makedirs, PathLike +from os.path import dirname, exists, join as joinpath, normpath, splitext +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +from aiofile import async_open +from alist import AlistClient, AlistPath -VIDEO_EXT = ("mp4", "mkv", "flv", "avi", "wmv", "ts", "rmvb", "webm") -SUBTITLE_EXT = ("ass", "srt", "ssa", "sub") -IMG_EXT = ("png", "jpg") -ALL_EXT = (*VIDEO_EXT, *SUBTITLE_EXT, *IMG_EXT, "nfo") + +VIDEO_EXTS: Final = frozenset((".mp4", ".mkv", ".flv", ".avi", ".wmv", ".ts", ".rmvb", ".webm")) +SUBTITLE_EXTS: Final = frozenset((".ass", ".srt", ".ssa", ".sub")) +IMAGE_EXTS: Final = frozenset((".png", ".jpg")) +NFO_EXTS: Final = frozenset((".nfo",)) class Alist2Strm: + def __init__( - self, - alist_server_url: str, - alist_server_username: str, - alist_server_password: str, - alist_server_base_dir: Optional[str], - token: Optional[str], - output_dir: Path, - subtitle: bool = False, - img: bool = False, - nfo: bool = False, - library_mode: bool = True, - async_mode: bool = False, - overwrite: bool = False, + self, + origin: str = "http://localhost:5244", + username: str = "", + password: str = "", + token: str = "", + base_dir: str = "/", + output_dir: bytes | str | PathLike = "", + library_mode: bool = False, + subtitle: bool = False, + image: bool = False, + nfo: bool = False, + overwrite: bool = False, + async_mode: bool = False, + max_worders: int = 5, ) -> None: + """实例化 Alist2Strm 对象 + + :param origin: Alist 服务器地址，默认为 "http://localhost:5244" + :param username: Alist 用户名，默认为空 + :param password: Alist 密码，默认为空 + :param token: Alist 签名 token，默认为空 + :param base_dir: 需要同步的 Alist 的目录，默认为 "/" + :param output_dir: strm 文件输出目录，默认为当前工作目录 + :param library_mode: 是否启用媒体库模式，所有文件下载到同一个目录，默认为 False + :param subtitle: 是否下载字幕文件，默认为 False + :param image: 是否下载图片文件，默认为 False + :param nfo: 是否下载 .nfo 文件，默认为 False + :param overwrite: 本地路径存在同名文件时是否重新生成/下载该文件，默认为 False + :param async_mode: 是否用异步模式，默认为 True + :param max_worders: 最大并发数 """ - 实例化 Alist2Strm 对象 - - :param alist_server_url: Alist 服务器地址 - :param alist_server_username: Alist 用户名 - :param alist_server_username: Alist 密码 - :param alist_server_base_dir: 底层目录，默认为 "/" - :param token: Alist 签名 Token - :param output_dir: Strm 文件输出目录 - :param subtitle: 是否下载字幕文件，默认 "False" - :param img: 是否下载图片文件，默认 "False" - :param nfo: 是否下载NFO文件，默认 "False" - :param library_mode: 是否启用媒体库模式，默认 "True" - :param async_mode: 是否启用异步下载文件，默认 "False" - :param overwrite: 本地路径存在同名文件时是否重新生成/下载该文件，默认 "False" - """ - - self.alist_server_url = alist_server_url.rstrip("/") - self.alist_server_username = alist_server_username - self.alist_server_password = alist_server_password - if alist_server_base_dir == None or alist_server_base_dir == "": - self.alist_server_base_dir = "/" - else: - self.alist_server_base_dir = "/" + alist_server_base_dir.strip("/") + "/" - self.token = token if token else None - self.output_dir = output_dir - self.subtitle = subtitle - self.img = img - self.nfo = nfo + client = self.client = AlistClient( + origin=origin, + username=username, + password=password, + ) + self.token = token + self.base_dir = client.fs.abspath(base_dir) + self.output_dir = fsdecode(output_dir) self.library_mode = library_mode - self.async_mode = async_mode + download_exts: set[str] = set() + if subtitle: + download_exts |= SUBTITLE_EXTS + if image: + download_exts |= IMAGE_EXTS + if nfo: + download_exts |= NFO_EXTS + self.download_exts = download_exts self.overwrite = overwrite - - logging.debug( - f"Alist2Strm配置".center(50, "=") + "\n" - f"Alist地址：{self.alist_server_url}\n" - f"Alist用户名：{self.alist_server_username}\n" - f"Alist密码：{self.alist_server_password}\n" - f"Alist基本路径：{self.alist_server_base_dir}\n" - f"Alist签名Token：{self.token}\n" - f"输出目录：{self.output_dir}\n" - f"是否下载字幕：{self.subtitle}\n" - f"是否下载图片：{self.img}\n" - f"是否下载NFO：{self.nfo}\n" - f"是否为库模式：{self.library_mode}\n" - f"是否为启用异步下载：{self.async_mode}" - ) + self.async_mode = async_mode + self.max_worders = max_worders + self._async_semaphore = Semaphore(max_worders) + + logging.debug("Alist2Strm配置".center(50, "=") + f"""\ +Alist 地址： {origin!r} +Alist 用户名：{username!r} +Alist 密码： {password!r} +Alist token： {token!r} +Alist 目录： {base_dir!r} +输出目录： {output_dir!r} +媒体库模式： {library_mode} +下载字幕： {subtitle} +下载图片： {image} +下载 NFO： {nfo} +异步模式： {async_mode} +覆盖： {overwrite} +最大并发数： {max_worders}""") def run(self) -> None: """ - 异步启动程序 + 启动程序 """ + if self.async_mode: + run(self._processer_async()) + else: + self._processer() - asyncio.run(self._processer()) - - async def _processer(self) -> None: + def _processer(self, /): + """程序处理主体（多线程） """ - 程序处理主体 + with ThreadPoolExecutor(max_workers=self.max_worders) as executor: + submit = executor.submit + for path in self.client.fs.iter( + self.base_dir, + max_depth=-1, + predicate=lambda path: path.is_file(), + ): + submit(self._file_processer, path) + + async def _processer_async(self, /): + """程序处理主体（异步） """ - try: - fs = AlistFileSystem.login( - self.alist_server_url, - self.alist_server_username, - self.alist_server_password, - ) - except Exception as e: - logging.critical( - "登录失败".center(50, "=") + "\n" - f"错误信息：{str(e)}\n" - f"请检查Alist地址：{self.alist_server_url}\n" - f"用户名：{self.alist_server_username}\n" - f"密码：{self.alist_server_password}是否正确" - ) - return - - try: - fs.chdir(self.alist_server_base_dir) - except Exception as e: - logging.critical( - "切换目录失败".center(50, "=") + "\n" - f"错误信息：{str(e)}\n" - f"请检查Alist服务器中是否存在该目录：{self.alist_server_base_dir}" - ) + async with TaskGroup() as tg: + create_task = tg.create_task + async for path in self.client.fs.iter( + self.base_dir, + max_depth=-1, + predicate=lambda path: path.is_file(), + async_=True, + ): + create_task(self._file_processer_async(path)) + + def _file_processer(self, /, path: AlistPath): + """保存文件至本地（多线程） + + :param path: AlistPath 对象 + """ + suffix = path.suffix.lower() + if not (suffix in VIDEO_EXTS or suffix in self.download_exts): return - if self.async_mode: - self.session = AsyncSession() + if self.library_mode: + local_path = joinpath(self.output_dir, path.name) else: - self.session = Session() - - async with asyncio.TaskGroup() as tg: - for alist_path_cls in fs.rglob("*.*"): - tg.create_task(self._file_processer(alist_path_cls)) + local_path = joinpath(self.output_dir, normpath(path.relative_to(self.base_dir))) - if self.session: - if self.async_mode: - await self.session.close() + try: + if exists(local_path) and not self.overwrite: + logging.debug(f"跳过文件：{local_path!r}") + return + + url = path.get_url(token=self.token) + if dir_ := dirname(local_path): + makedirs(dir_, exist_ok=True) + + if suffix in VIDEO_EXTS: + local_path = splitext(local_path)[0] + ".strm" + open(local_path, mode="w", encoding="utf-8").write(url) + logging.debug(f"创建文件：{local_path!r}") else: - self.session.close() - - async def _file_processer(self, alist_path_cls: AlistPath) -> None: + with ( + closing(self.client.request(url, "GET", parse=None)) as resp, + open(local_path, mode="wb") as file, + ): + write = file.write + for chunk in resp.iter_bytes(1 << 16): + write(chunk) + logging.debug(f"下载文件：{local_path!r}") + except: + logging.exception(f"下载失败: {local_path!r}") + raise + + async def _file_processer_async(self, /, path: AlistPath): + """保存文件至本地（异步） + + :param path: AlistPath 对象 """ - 保存文件至本地 - - :param alist_path_cls: AlistPath 对象 - """ - - if not alist_path_cls.name.lower().endswith(ALL_EXT): - return - - file_output_path: Path = ( - self.output_dir / alist_path_cls.name - if self.library_mode - else self.output_dir - / str(alist_path_cls).replace(self.alist_server_base_dir, "") - ) - - if self.overwrite == False and file_output_path.exists(): - logging.debug( - f"{file_output_path.name}已存在，跳过下载，文件本地目录：{file_output_path.parent}" - ) - return - - file_alist_abs_path: str = alist_path_cls.url[ - alist_path_cls.url.index("/d/") + 2 : - ] - - file_download_url: str = alist_path_cls.url + self._sign( - secret_key=self.token, data=file_alist_abs_path - ) - - logging.debug( - f"正在处理:{alist_path_cls.name}\n" - f"本地文件目录：{file_output_path}\n" - f"文件远程路径：{file_alist_abs_path}\n" - f"下载URL：{file_download_url}" - ) - - if alist_path_cls.name.lower().endswith(VIDEO_EXT): - file_output_path.parent.mkdir(parents=True, exist_ok=True) - file_output_path = file_output_path.with_suffix(".strm") - with file_output_path.open(mode="w", encoding="utf-8") as f: - f.write(file_download_url) - logging.debug( - f"{file_output_path.name}创建成功，文件本地目录：{file_output_path.parent}" - ) + suffix = path.suffix.lower() + if not (suffix in VIDEO_EXTS or suffix in self.download_exts): return - if alist_path_cls.name.lower().endswith(IMG_EXT) and not self.img: - return - elif alist_path_cls.name.lower().endswith(SUBTITLE_EXT) and not self.subtitle: - return - elif alist_path_cls.name.lower().endswith("nfo") and not self.nfo: - return + if self.library_mode: + local_path = joinpath(self.output_dir, path.name) else: - file_output_path.parent.mkdir(parents=True, exist_ok=True) - - if self.async_mode: - resp = await self.session.get(file_download_url) - else: - resp = self.session.get(file_download_url) - - with file_output_path.open(mode="wb") as f: - if self.async_mode: - f.write(await resp.read()) + local_path = joinpath(self.output_dir, normpath(path.relative_to(self.base_dir))) + + async with self._async_semaphore: + try: + if exists(local_path) and not self.overwrite: + logging.debug(f"跳过文件：{local_path!r}") + return + + url = path.get_url(token=self.token) + if dir_ := dirname(local_path): + await to_thread(makedirs, dir_, exist_ok=True) + + if suffix in VIDEO_EXTS: + local_path = splitext(local_path)[0] + ".strm" + async with async_open(local_path, mode="w", encoding="utf-8") as file: + await file.write(url) + logging.debug(f"创建文件：{local_path!r}") else: - f.write(resp.content) - logging.debug( - f"{file_output_path.name}下载成功，文件本地目录：{file_output_path.parent}" - ) - - def _sign(self, secret_key: Optional[str], data: str) -> str: - """ - Alist 签名 Token 处理 - - :param secret_key: Alist 签名 Token - :param data: 待签名数据 （Alist 文件绝对路径） - """ - - if secret_key == "" or secret_key == None: - return "" + async with ( + aclosing(await self.client.request(url, "GET", parse=None, async_=True)) as resp, + async_open(local_path, mode="wb") as file, + ): + write = file.write + async for chunk in resp.aiter_bytes(1 << 16): + await write(chunk) + logging.debug(f"下载文件：{local_path!r}") + except: + logging.exception(f"下载失败: {local_path!r}") + raise - data = unquote(data) - h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256) - expire_time_stamp = str(0) - h.update((data + ":" + expire_time_stamp).encode()) - return f"?sign={base64.urlsafe_b64encode(h.digest()).decode()}:0" diff --git a/requirements.txt b/requirements.txt index 58b36b5..0f5af57 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ +aiofile == 3.8.8 PyYAML == 6.0.1 -python-alist == 0.0.12.2 \ No newline at end of file +python-alist == 0.0.13.0.2 \ No newline at end of file