From d30a02987d4a370441c4c4158bbe90cdb42f73fe Mon Sep 17 00:00:00 2001 From: thsrite Date: Thu, 28 Sep 2023 11:10:34 +0800 Subject: [PATCH] =?UTF-8?q?feat=20=E6=AD=A3=E5=9C=A8=E4=B8=8B=E8=BD=BD?= =?UTF-8?q?=E8=BF=9B=E5=BA=A6=E6=8E=A8=E9=80=81=E6=8F=92=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- alembic/versions/b2f011d3a8b7_1_0_8.py | 29 +++ app/chain/download.py | 3 +- app/db/models/downloadhistory.py | 2 + app/plugins/downloading/__init__.py | 269 +++++++++++++++++++++++++ app/plugins/nastoolsync/__init__.py | 4 +- 5 files changed, 305 insertions(+), 2 deletions(-) create mode 100644 alembic/versions/b2f011d3a8b7_1_0_8.py create mode 100644 app/plugins/downloading/__init__.py diff --git a/alembic/versions/b2f011d3a8b7_1_0_8.py b/alembic/versions/b2f011d3a8b7_1_0_8.py new file mode 100644 index 00000000..9f7373e4 --- /dev/null +++ b/alembic/versions/b2f011d3a8b7_1_0_8.py @@ -0,0 +1,29 @@ +"""1_0_8 + +Revision ID: b2f011d3a8b7 +Revises: 30329639c12b +Create Date: 2023-09-28 10:15:58.410003 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = 'b2f011d3a8b7' +down_revision = '30329639c12b' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + try: + with op.batch_alter_table("downloadhistory") as batch_op: + batch_op.add_column(sa.Column('userid', sa.String, nullable=True)) + except Exception as e: + pass + # ### end Alembic commands ### + +def downgrade() -> None: + pass \ No newline at end of file diff --git a/app/chain/download.py b/app/chain/download.py index fcdc5a5a..20f29582 100644 --- a/app/chain/download.py +++ b/app/chain/download.py @@ -266,7 +266,8 @@ class DownloadChain(ChainBase): download_hash=_hash, torrent_name=_torrent.title, torrent_description=_torrent.description, - torrent_site=_torrent.site_name + torrent_site=_torrent.site_name, + userid=userid ) # 登记下载文件 diff --git a/app/db/models/downloadhistory.py b/app/db/models/downloadhistory.py index f624d09d..2fc90755 100644 --- a/app/db/models/downloadhistory.py +++ b/app/db/models/downloadhistory.py @@ -35,6 +35,8 @@ class DownloadHistory(Base): torrent_description = Column(String) # 种子站点 torrent_site = Column(String) + # 下载用户 + userid = Column(String) # 附加信息 note = Column(String) diff --git a/app/plugins/downloading/__init__.py b/app/plugins/downloading/__init__.py new file mode 100644 index 00000000..5320715c --- /dev/null +++ b/app/plugins/downloading/__init__.py @@ -0,0 +1,269 @@ +from apscheduler.schedulers.background import BackgroundScheduler + +from app.chain.download import DownloadChain +from app.core.config import settings +from app.db.downloadhistory_oper import DownloadHistoryOper +from app.plugins import _PluginBase +from typing import Any, List, Dict, Tuple, Optional +from app.log import logger +from app.schemas import NotificationType +from app.schemas.types import TorrentStatus +from app.utils.string import StringUtils + + +class Downloading(_PluginBase): + # 插件名称 + plugin_name = "正在下载" + # 插件描述 + plugin_desc = "定时推送正在下载进度。" + # 插件图标 + plugin_icon = "download.png" + # 主题色 + plugin_color = "#f2a026" + # 插件版本 + plugin_version = "1.0" + # 插件作者 + plugin_author = "thsrite" + # 作者主页 + author_url = "https://github.com/thsrite" + # 插件配置项ID前缀 + plugin_config_prefix = "downloading_" + # 加载顺序 + plugin_order = 16 + # 可使用的用户级别 + auth_level = 2 + + # 私有属性 + _enabled = False + # 任务执行间隔 + _seconds = None + _type = None + _adminuser = None + _downloadhis = None + + # 定时器 + _scheduler: Optional[BackgroundScheduler] = None + + def init_plugin(self, config: dict = None): + # 停止现有任务 + self.stop_service() + + if config: + self._enabled = config.get("enabled") + self._seconds = config.get("seconds") or 300 + self._type = config.get("type") or 'admin' + self._adminuser = config.get("adminuser") + + # 加载模块 + if self._enabled: + self._downloadhis = DownloadHistoryOper(self.db) + # 定时服务 + self._scheduler = BackgroundScheduler(timezone=settings.TZ) + + if self._seconds: + try: + self._scheduler.add_job(func=self.__downloading, + trigger='interval', + seconds=int(self._seconds), + name="正在下载") + except Exception as err: + logger.error(f"定时任务配置错误：{err}") + + # 启动任务 + if self._scheduler.get_jobs(): + self._scheduler.print_jobs() + self._scheduler.start() + + def __downloading(self): + """ + 定时推送正在下载进度 + """ + # 正在下载种子 + torrents = DownloadChain(self.db).list_torrents(status=TorrentStatus.DOWNLOADING) + if not torrents: + logger.info("当前没有正在下载的任务！") + return + # 推送用户 + if self._type == "admin" or self._type == "both": + if not self._adminuser: + logger.error("未配置管理员用户") + return + + for userid in str(self._adminuser).split(","): + self.__send_msg(torrents=torrents, userid=userid) + + if self._type == "user" or self._type == "both": + user_torrents = {} + # 根据正在下载种子hash获取下载历史 + for torrent in torrents: + downloadhis = self._downloadhis.get_by_hash(download_hash=torrent.hash) + if not downloadhis: + logger.warn(f"种子 {torrent.hash} 未获取到MoviePilot下载历史，无法推送下载进度") + continue + if not downloadhis.userid: + logger.warn(f"种子 {torrent.hash} 未获取到下载用户记录，无法推送下载进度") + continue + user_torrent = user_torrents.get(downloadhis.userid) or [] + user_torrent.append(torrent) + + if not user_torrents or not user_torrents.keys(): + logger.warn("未获取到用户下载记录，无法推送下载进度") + return + + # 推送用户下载任务进度 + for userid in list(user_torrents.keys()): + # 如果用户是管理员，无需重复推送 + if self._adminuser and userid in str(self._adminuser).split(","): + logger.debug("管理员已推送") + continue + + user_torrent = user_torrents.get(userid) + if not user_torrent: + logger.warn(f"未获取到用户 {userid} 下载任务") + continue + self.__send_msg(torrents=user_torrent, + userid=userid) + + if self._type == "all": + self.__send_msg(torrents=torrents) + + def __send_msg(self, torrents, userid=None): + """ + 发送消息 + """ + title = f"共 {len(torrents)} 个任务正在下载：" + messages = [] + index = 1 + for torrent in torrents: + messages.append(f"{index}. {torrent.title} " + f"{StringUtils.str_filesize(torrent.size)} " + f"{round(torrent.progress, 1)}%") + index += 1 + self.post_message(mtype=NotificationType.Download, + title=title, + text="\n".join(messages), + userid=userid) + + def get_state(self) -> bool: + return self._enabled + + @staticmethod + def get_command() -> List[Dict[str, Any]]: + pass + + def get_api(self) -> List[Dict[str, Any]]: + pass + + def get_form(self) -> Tuple[List[dict], Dict[str, Any]]: + """ + 拼装插件配置页面，需要返回两块数据：1、页面配置；2、数据结构 + """ + return [ + { + 'component': 'VForm', + 'content': [ + { + 'component': 'VRow', + 'content': [ + { + 'component': 'VCol', + 'props': { + 'cols': 12 + }, + 'content': [ + { + 'component': 'VSwitch', + 'props': { + 'model': 'enabled', + 'label': '启用插件', + } + } + ] + } + ] + }, + { + 'component': 'VRow', + 'content': [ + { + 'component': 'VCol', + 'props': { + 'cols': 12, + 'md': 4 + }, + 'content': [ + { + 'component': 'VTextField', + 'props': { + 'model': 'seconds', + 'label': '执行间隔', + 'placeholder': '单位（秒）' + } + } + ] + }, + { + 'component': 'VCol', + 'props': { + 'cols': 12, + 'md': 4 + }, + 'content': [ + { + 'component': 'VTextField', + 'props': { + 'model': 'adminuser', + 'label': '管理员用户', + 'placeholder': '多个用户,分割' + } + } + ] + }, + { + 'component': 'VCol', + 'props': { + 'cols': 12, + 'md': 4 + }, + 'content': [ + { + 'component': 'VSelect', + 'props': { + 'model': 'type', + 'label': '推送类型', + 'items': [ + {'title': 'admin', 'value': 'admin'}, + {'title': '下载用户', 'value': 'user'}, + {'title': 'admin和下载用户', 'value': 'both'}, + {'title': '所有用户', 'value': 'all'} + ] + } + } + ] + } + ] + } + ] + } + ], { + "enabled": False, + "seconds": 300, + "adminuser": "", + "type": "admin" + } + + def get_page(self) -> List[dict]: + pass + + def stop_service(self): + """ + 退出插件 + """ + try: + if self._scheduler: + self._scheduler.remove_all_jobs() + if self._scheduler.running: + self._scheduler.shutdown() + self._scheduler = None + except Exception as e: + logger.error("退出插件失败：%s" % str(e)) diff --git a/app/plugins/nastoolsync/__init__.py b/app/plugins/nastoolsync/__init__.py index b9ae06d4..fee9b890 100644 --- a/app/plugins/nastoolsync/__init__.py +++ b/app/plugins/nastoolsync/__init__.py @@ -3,6 +3,7 @@ import os import sqlite3 from datetime import datetime +from app.core.config import settings from app.db.downloadhistory_oper import DownloadHistoryOper from app.db.plugindata_oper import PluginDataOper from app.db.transferhistory_oper import TransferHistoryOper @@ -234,7 +235,8 @@ class NAStoolSync(_PluginBase): download_hash=mdownload_hash, torrent_name=mtorrent, torrent_description=mdesc, - torrent_site=msite + torrent_site=msite, + userid=settings.SUPERUSER ) cnt += 1 if cnt % 100 == 0: