Merge pull request #58 from cikezhu/main

package.json

@@ -338,7 +338,7 @@
     "DownloadSiteTag": {
         "name": "下载任务分类与标签",
         "description": "自动给下载任务分类与打站点标签、剧集名称标签",
-        "version": "1.2",
+        "version": "1.3",
         "icon": "Youtube-dl_B.png",
         "author": "叮叮当",
         "level": 1

plugins/downloadsitetag/__init__.py

@@ -1,3 +1,5 @@
+import datetime
+import pytz
 from typing import List, Tuple, Dict, Any
 
 from app.core.context import Context
@@ -8,7 +10,9 @@ from app.log import logger
 from app.plugins import _PluginBase
 from app.modules.qbittorrent import Qbittorrent
 from app.modules.transmission import Transmission
-
+from app.db.downloadhistory_oper import DownloadHistoryOper
+from app.modules.themoviedb.tmdbapi import TmdbHelper
+from apscheduler.schedulers.background import BackgroundScheduler
 
 class DownloadSiteTag(_PluginBase):
     # 插件名称
@@ -18,7 +22,7 @@ class DownloadSiteTag(_PluginBase):
     # 插件图标
     plugin_icon = "Youtube-dl_B.png"
     # 插件版本
-    plugin_version = "1.2"
+    plugin_version = "1.3"
     # 插件作者
     plugin_author = "叮叮当"
     # 作者主页
@@ -29,11 +33,17 @@ class DownloadSiteTag(_PluginBase):
     plugin_order = 2
     # 可使用的用户级别
     auth_level = 1
+    # 日志前缀
+    LOG_TAG = "[DownloadSiteTag] "
 
     # 私有属性
-    downloader_name = None
-    downloader_example = None
+    downloader_qb = None
+    downloader_tr = None
+    downloadhistory_oper = None
+    tmdb_helper = None
+    _scheduler = None
     _enabled = False
+    _onlyonce = False
     _enabled_media_tag = False
     _enabled_tag = True
     _enabled_category = False
@@ -42,15 +52,41 @@ class DownloadSiteTag(_PluginBase):
     _category_anime = None
 
     def init_plugin(self, config: dict = None):
+        self.downloader_qb = Qbittorrent()
+        self.downloader_tr = Transmission()
+        self.downloadhistory_oper = DownloadHistoryOper()
+        self.tmdb_helper = TmdbHelper()
         # 读取配置
         if config:
             self._enabled = config.get("enabled")
+            self._onlyonce = config.get("onlyonce")
             self._enabled_media_tag = config.get("enabled_media_tag")
             self._enabled_tag = config.get("enabled_tag")
             self._enabled_category = config.get("enabled_category")
             self._category_movie = config.get("category_movie") or "电影"
             self._category_tv = config.get("category_tv") or "电视"
             self._category_anime = config.get("category_anime") or "动漫"
+        
+        # 停止现有任务
+        self.stop_service()
+
+        if self._onlyonce:
+            # 创建定时任务控制器
+            self._scheduler = BackgroundScheduler(timezone=settings.TZ)
+            # 执行一次, 关闭onlyonce
+            self._onlyonce = False
+            config.update({"onlyonce": self._onlyonce})
+            self.update_config(config)
+            # 添加 补全下载历史的标签与分类 任务
+            self._scheduler.add_job(func= self._complemented_history, trigger='date',
+                                            run_date=datetime.datetime.now(
+                                                tz=pytz.timezone(settings.TZ)) + datetime.timedelta(seconds=3)
+                                            )
+
+        if self._scheduler and self._scheduler.get_jobs():
+            # 启动服务
+            self._scheduler.print_jobs()
+            self._scheduler.start()
 
     def get_state(self) -> bool:
         return self._enabled
@@ -62,6 +98,164 @@ class DownloadSiteTag(_PluginBase):
     def get_api(self) -> List[Dict[str, Any]]:
         pass
 
+    def _complemented_history(self):
+        """
+        补全下载历史的标签与分类
+        """
+        logger.info(f"{self.LOG_TAG}开始执行: 补全下载历史的标签与分类 ...")
+        for DOWNLOADER in ["qbittorrent", "transmission"]:
+            logger.info(f"{self.LOG_TAG}开始扫描下载器 {DOWNLOADER} ...")
+            # 获取下载器中的种子
+            downloader_obj = self._get_downloader(DOWNLOADER)
+            if not downloader_obj:
+                logger.error(f"{self.LOG_TAG} 获取下载器失败 {DOWNLOADER}")
+                continue
+            torrents, error = downloader_obj.get_torrents()
+            # 如果下载器获取种子发生错误 或 没有种子 则跳过
+            if error or not torrents:
+                continue
+            else:
+                logger.info(f"{self.LOG_TAG}下载器 {DOWNLOADER} 种子数：{len(torrents)}")
+            for torrent in torrents:
+                # 获取种子hash
+                _hash = self._get_hash(torrent=torrent, dl_type=DOWNLOADER)
+                if not _hash:
+                    continue
+                # 提取种子hash对应的下载历史
+                history = self.downloadhistory_oper.get_by_hash(_hash)
+                if not history:
+                    continue
+                # 获取种子当前标签
+                torrent_tags = self._get_label(torrent=torrent, dl_type=DOWNLOADER)
+                torrent_cat = self._get_category(torrent=torrent, dl_type=DOWNLOADER)
+                # 按设置生成需要写入的标签与分类
+                _tags = []
+                _cat = None
+                tmdbid = history.tmdbid
+                mtype = history.type
+                # 站点标签, 如果勾选开关的话
+                if self._enabled_tag and history.torrent_site:
+                    _tags.append(history.torrent_site)
+                # 媒体标题标签, 如果勾选开关的话
+                if self._enabled_media_tag and history.title:
+                    _tags.append(history.title)
+                # 分类, 如果勾选开关的话 <tr暂不支持>
+                if DOWNLOADER == "qbittorrent" and self._enabled_category and not torrent_cat:
+                    # 如果是电视剧 需要区分是否动漫
+                    genre_ids = None
+                    if mtype == MediaType.TV or mtype == MediaType.TV.value:
+                        # tmdb_id获取tmdb信息
+                        tmdb_info = self.tmdb_helper.get_info(mtype=mtype, tmdbid=tmdbid)
+                        if tmdb_info:
+                            genre_ids = tmdb_info.get("genre_ids")
+                    _cat = self._genre_ids_get_cat(mtype, genre_ids)
+                
+                # 去除种子已经存在的标签
+                if _tags and torrent_tags:
+                    _tags = list(set(_tags) - set(torrent_tags))
+                # 如果分类一样, 那么不需要修改
+                if _cat == torrent_cat:
+                    _cat = None
+                # 判断当前种子是否不需要修改
+                if not _cat and not _tags:
+                    continue
+                # 执行通用方法, 设置种子标签与分类
+                self._set_torrent_info(DOWNLOADER=DOWNLOADER, _hash=_hash, _tags=_tags, _cat=_cat, _original_tags=torrent_tags)
+
+        logger.info(f"{self.LOG_TAG}执行完成: 补全下载历史的标签与分类 ...")
+
+    def _genre_ids_get_cat(self, mtype, genre_ids = None):
+        """
+        根据genre_ids判断是否<动漫>分类
+        """
+        _cat = None
+        if mtype == MediaType.MOVIE or mtype == MediaType.MOVIE.value:
+            # 电影
+            _cat = self._category_movie
+        elif mtype:
+            ANIME_GENREIDS = settings.ANIME_GENREIDS
+            if genre_ids \
+                    and set(genre_ids).intersection(set(ANIME_GENREIDS)):
+                # 动漫
+                _cat = self._category_anime
+            else:
+                # 电视剧
+                _cat = self._category_tv
+        return _cat
+
+    def _get_downloader(self, dtype: str):
+        """
+        根据类型返回下载器实例
+        """
+        if dtype == "qbittorrent":
+            return self.downloader_qb
+        elif dtype == "transmission":
+            return self.downloader_tr
+        else:
+            return None
+
+    def _get_hash(self, torrent: Any, dl_type: str):
+        """
+        获取种子hash
+        """
+        try:
+            return torrent.get("hash") if dl_type == "qbittorrent" else torrent.hashString
+        except Exception as e:
+            print(str(e))
+            return ""
+
+    def _get_label(self, torrent: Any, dl_type: str):
+        """
+        获取种子标签
+        """
+        try:
+            return [str(tag).strip() for tag in torrent.get("tags").split(',')] \
+                if dl_type == "qbittorrent" else torrent.labels or []
+        except Exception as e:
+            print(str(e))
+            return []
+
+    def _get_category(self, torrent: Any, dl_type: str):
+        """
+        获取种子分类
+        """
+        try:
+            return torrent.get("category") if dl_type == "qbittorrent" else None
+        except Exception as e:
+            print(str(e))
+            return None
+
+    def _set_torrent_info(self, DOWNLOADER: str, _hash: str, _tags: list, _cat: str, _original_tags: list = None):
+        """
+        设置种子标签与分类
+        """
+        # 当前下载器
+        downloader_obj = self._get_downloader(DOWNLOADER)
+        # 判断是否可执行
+        if DOWNLOADER and downloader_obj and _hash:
+            # 下载器api不通用, 因此需分开处理
+            if DOWNLOADER == "qbittorrent":
+                # 设置标签
+                if _tags:
+                    downloader_obj.set_torrents_tag(ids=_hash, tags=_tags)
+                # 设置分类 <tr暂不支持>
+                if _cat:
+                    downloader_obj.qbc.torrents_set_category(category=_cat, torrent_hashes=_hash)
+            else:
+                # 设置标签
+                if _tags:
+                    # _original_tags = None表示未指定, 因此需要获取原始标签
+                    if _original_tags == None:
+                        torrent = downloader_obj.trc.get_torrent(torrent_id=_hash)
+                        if torrent:
+                            _original_tags = self._get_label(torrent=torrent, dl_type=DOWNLOADER)
+                    # 如果原始标签不是空的, 那么合并原始标签
+                    if _original_tags:
+                        _tags = list(set(_original_tags).union(set(_tags)))
+                    downloader_obj.set_torrent_tag(ids=_hash, tags=_tags)
+            logger.warn(
+                f"{self.LOG_TAG}当前下载器: {DOWNLOADER} {('  TAG: ' + ','.join(_tags)) if _tags else ''} {('  CAT: ' + _cat) if _cat else ''}")
+
     @eventmanager.register(EventType.DownloadAdded)
     def DownloadAdded(self, event: Event):
         """
@@ -72,58 +266,25 @@ class DownloadSiteTag(_PluginBase):
 
         if not event.event_data:
             return
+        
+        context: Context = event.event_data.get("context")
+        _hash = event.event_data.get("hash")
+        _torrent = context.torrent_info
+        _media = context.media_info
+        _tags = []
+        _cat = None
+        # 站点标签, 如果勾选开关的话
+        if self._enabled_tag:
+           _tags.append(_torrent.site_name)
+        # 媒体标题标签, 如果勾选开关的话
+        if self._enabled_media_tag:
+            _tags.append(_media.title)
+        # 分类, 如果勾选开关的话 <tr暂不支持>
+        if self._enabled_category:
+            _cat = self._genre_ids_get_cat(_media.type, _media.genre_ids)
+        # 执行通用方法, 设置种子标签与分类
+        self._set_torrent_info(DOWNLOADER=settings.DOWNLOADER, _hash=_hash, _tags=_tags, _cat=_cat)
 
-        DOWNLOADER = getattr(settings, "DOWNLOADER")
-        if self.downloader_name != DOWNLOADER:
-            if DOWNLOADER == "qbittorrent":
-                self.downloader_example = Qbittorrent()
-                self.downloader_name = DOWNLOADER
-            elif DOWNLOADER == "transmission":
-                self.downloader_example = Transmission()
-                self.downloader_name = DOWNLOADER
-            else:
-                self.downloader_example = None
-                self.downloader_name = None
-        if self.downloader_name and self.downloader_example:
-            context: Context = event.event_data.get("context")
-            _hash = event.event_data.get("hash")
-            _torrent = context.torrent_info
-            _media = context.media_info
-            _media_type = None
-
-            if self.downloader_name == "qbittorrent":
-                # 设置标签, 如果勾选开关的话
-                if self._enabled_tag:
-                    self.downloader_example.set_torrents_tag(ids=_hash, tags=[_torrent.site_name])
-                # 设置媒体标题标签, 如果勾选开关的话
-                if self._enabled_media_tag:
-                    self.downloader_example.set_torrents_tag(ids=_hash, tags=[_media.title])
-                # 设置分类, 如果勾选开关的话 <tr暂不支持>
-                if self._enabled_category:
-                    if _media.type == MediaType.MOVIE:
-                        # 电影
-                        _media_type = self._category_movie
-                    else:
-                        ANIME_GENREIDS = getattr(settings, "ANIME_GENREIDS")
-                        if _media.genre_ids \
-                                and set(_media.genre_ids).intersection(set(ANIME_GENREIDS)):
-                            # 动漫
-                            _media_type = self._category_anime
-                        else:
-                            # 电视剧
-                            _media_type = self._category_tv
-                    # 使用qbapi进行设置
-                    if _media_type:
-                        self.downloader_example.qbc.torrents_set_category(category=_media_type, torrent_hashes=_hash)
-            else:
-                # 设置标签, 如果勾选开关的话
-                if self._enabled_tag:
-                    self.downloader_example.set_torrent_tag(ids=_hash, tags=[_torrent.site_name])
-                # 设置媒体标题标签, 如果勾选开关的话
-                if self._enabled_media_tag:
-                    self.downloader_example.set_torrent_tag(ids=_hash, tags=[_media.title])
-            logger.warn(
-                f"[DownloadSiteTag] 当前下载器: {self.downloader_name}{('  TAG: ' + _torrent.site_name) if self._enabled_tag else ''}{('  MEDIA_TAG: ' + _media.title) if self._enabled_media_tag else ''}{('  CAT: ' + _media_type) if _media_type else ''}")
 
     def get_form(self) -> Tuple[List[dict], Dict[str, Any]]:
         """
@@ -265,6 +426,21 @@ class DownloadSiteTag(_PluginBase):
                             }
                         ]
                     },
+                    {
+                        'component': 'VCol',
+                        'props': {
+                            'cols': 12,
+                        },
+                        'content': [
+                            {
+                                'component': 'VSwitch',
+                                'props': {
+                                    'model': 'onlyonce',
+                                    'label': '立即执行一次: 补全下载历史的标签与分类',
+                                }
+                            }
+                        ]
+                    },
                     {
                         'component': 'VRow',
                         'content': [
@@ -300,7 +476,7 @@ class DownloadSiteTag(_PluginBase):
                                         'props': {
                                             'type': 'info',
                                             'variant': 'tonal',
-                                            'text': '注意：分类功能只支持qbittorrent下载器'
+                                            'text': '注意：分类只支持qb, 并提前在qb创建好分类'
                                         }
                                     }
                                 ]
@@ -311,6 +487,7 @@ class DownloadSiteTag(_PluginBase):
             }
         ], {
             "enabled": False,
+            "onlyonce": False,
             "enabled_tag": True,
             "enabled_media_tag": False,
             "enabled_category": False,
@@ -324,6 +501,15 @@ class DownloadSiteTag(_PluginBase):
 
     def stop_service(self):
         """
-        退出插件
+        停止服务
         """
-        pass
+        try:
+            if self._scheduler:
+                self._scheduler.remove_all_jobs()
+                if self._scheduler.running:
+                    self._event.set()
+                    self._scheduler.shutdown()
+                    self._event.clear()
+                self._scheduler = None
+        except Exception as e:
+            print(str(e))