Merge pull request #289 from EstrellaXD/3.0-dev

3.0.2

src/module/api/download.py

@@ -32,8 +32,10 @@ async def download_collection(
         )
     if data:
         with SeasonCollector() as collector:
-            collector.collect_season(data, data.rss_link[0])
-        return {"status": "Success"}
+            if collector.collect_season(data, data.rss_link[0], proxy=True):
+                return {"status": "Success"}
+            else:
+                return {"status": "Failed to add torrent"}
     else:
         return {"status": "Failed to parse link"}
 

src/module/downloader/client/qb_downloader.py

@@ -80,14 +80,16 @@ class QbDownloader:
     def torrents_info(self, status_filter, category, tag=None):
         return self._client.torrents_info(status_filter=status_filter, category=category, tag=tag)
 
-    def torrents_add(self, urls, save_path, category):
-        return self._client.torrents_add(
+    def torrents_add(self, urls, save_path, category, torrent_files=None):
+        resp = self._client.torrents_add(
             is_paused=False,
             urls=urls,
+            torrent_files=torrent_files,
             save_path=save_path,
             category=category,
             use_auto_torrent_management=False
         )
+        return resp == "Ok."
 
     def torrents_delete(self, hash):
         return self._client.torrents_delete(delete_files=True, torrent_hashes=hash)

src/module/downloader/download_client.py

@@ -45,7 +45,7 @@ class DownloadClient(TorrentPath):
     def auth(self):
         self.authed = self.client.auth()
         if self.authed:
-            logger.info("[Downloader] Authed.")
+            logger.debug("[Downloader] Authed.")
         else:
             logger.error("[Downloader] Auth failed.")
 
@@ -112,9 +112,17 @@ class DownloadClient(TorrentPath):
         logger.info(f"[Downloader] Remove torrents.")
 
     def add_torrent(self, torrent: dict):
-        self.client.torrents_add(
-            urls=torrent["url"], save_path=torrent["save_path"], category="Bangumi"
-        )
+        if self.client.torrents_add(
+            urls=torrent.get("urls"),
+            torrent_files=torrent.get("torrent_files"),
+            save_path=torrent.get("save_path"),
+            category="Bangumi"
+        ):
+            logger.debug(f"[Downloader] Add torrent: {torrent.get('save_path')}")
+            return True
+        else:
+            logger.error(f"[Downloader] Add torrent failed: {torrent.get('save_path')}")
+            return False
 
     def move_torrent(self, hashes, location):
         self.client.move_torrent(hashes=hashes, new_location=location)

src/module/manager/collector.py

@@ -11,23 +11,30 @@ logger = logging.getLogger(__name__)
 
 
 class SeasonCollector(DownloadClient):
-    def add_season_torrents(self, data: BangumiData, torrents):
-        for torrent in torrents:
+    def add_season_torrents(self, data: BangumiData, torrents, torrent_files=None):
+        if torrent_files:
             download_info = {
-                "url": torrent.torrent_link,
+                "torrent_files": torrent_files,
                 "save_path": self._gen_save_path(data),
             }
-            self.add_torrent(download_info)
+            return self.add_torrent(download_info)
+        else:
+            download_info = {
+                "urls": [torrent.torrent_link for torrent in torrents],
+                "save_path": self._gen_save_path(data),
+            }
+            return self.add_torrent(download_info)
 
-    def collect_season(self, data: BangumiData, link: str = None):
+    def collect_season(self, data: BangumiData, link: str = None, proxy: bool = False):
         logger.info(f"Start collecting {data.official_title} Season {data.season}...")
         with SearchTorrent() as st:
             if not link:
                 torrents = st.search_season(data)
             else:
-                torrents = st.get_torrents(link)
-        self.add_season_torrents(data, torrents)
-        logger.info("Completed!")
+                torrents = st.get_torrents(link, _filter="|".join(data.filter))
+            if proxy:
+                torrent_files = [st.get_content(torrent.torrent_link) for torrent in torrents]
+            return self.add_season_torrents(data, torrents, torrent_files=torrent_files)
 
     def subscribe_season(self, data: BangumiData):
         with BangumiDatabase() as db:

src/module/manager/renamer.py

@@ -139,12 +139,12 @@ class Renamer(DownloadClient):
                     if not renamed:
                         logger.warning(f"[Renamer] {subtitle_path} rename failed")
 
-    def rename(self):
+    def rename(self) -> list[Notification]:
         # Get torrent info
         logger.debug("[Renamer] Start rename process.")
         rename_method = settings.bangumi_manage.rename_method
         torrents_info = self.get_torrent_info()
-        renamed_info = []
+        renamed_info: list[Notification] = []
         for info in torrents_info:
             media_list, subtitle_list = self.check_files(info)
             bangumi_name, season = self._path_to_bangumi(info.save_path)

src/module/network/request_contents.py

@@ -42,12 +42,11 @@ class RequestContent(RequestURL):
         _url: str,
         _filter: str = "|".join(settings.rss_parser.filter),
         retry: int = 3,
-    ) -> [TorrentInfo]:
+    ) -> list[TorrentInfo]:
         try:
             soup = self.get_xml(_url, retry)
             torrent_titles, torrent_urls, torrent_homepage = mikan_parser(soup)
-
-            torrents = []
+            torrents: list[TorrentInfo] = []
             for _title, torrent_url, homepage in zip(
                 torrent_titles, torrent_urls, torrent_homepage
             ):

src/module/network/request_url.py

@@ -13,6 +13,7 @@ logger = logging.getLogger(__name__)
 class RequestURL:
     def __init__(self):
         self.header = {"user-agent": "Mozilla/5.0", "Accept": "application/xml"}
+        self._socks5_proxy = False
 
     def get_url(self, url, retry=3):
         try_time = 0
@@ -77,6 +78,7 @@ class RequestURL:
                     "http": url,
                 }
             elif settings.proxy.type == "socks5":
+                self._socks5_proxy = True
                 socks.set_default_proxy(
                     socks.SOCKS5,
                     addr=settings.proxy.host,
@@ -91,4 +93,8 @@ class RequestURL:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._socks5_proxy:
+            socks.set_default_proxy()
+            socket.socket = socks.socksocket
+            self._socks5_proxy = False
         self.session.close()

src/module/notification/notification.py

@@ -10,20 +10,23 @@ from module.database import BangumiDatabase
 logger = logging.getLogger(__name__)
 
 
-def getClient(type=settings.notification.type):
+def getClient(type: str):
     if type.lower() == "telegram":
         return TelegramNotification
     elif type.lower() == "server-chan":
         return ServerChanNotification
     elif type.lower() == "bark":
         return BarkNotification
+    elif type.lower() == "wecom":
+        return WecomNotification
     else:
         return None
 
 
-class PostNotification(getClient()):
+class PostNotification:
     def __init__(self):
-        super().__init__(
+        Notifier = getClient(settings.notification.type)
+        self.notifier = Notifier(
             token=settings.notification.token,
             chat_id=settings.notification.chat_id
         )
@@ -46,12 +49,18 @@ class PostNotification(getClient()):
     def send_msg(self, notify: Notification) -> bool:
         text = self._gen_message(notify)
         try:
-            self.post_msg(text)
+            self.notifier.post_msg(text)
             logger.debug(f"Send notification: {notify.official_title}")
         except Exception as e:
             logger.warning(f"Failed to send notification: {e}")
             return False
 
+    def __enter__(self):
+        self.notifier.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.notifier.__exit__(exc_type, exc_val, exc_tb)
 
 if __name__ == "__main__":
     info = Notification(

src/module/notification/plugin/__init__.py

@@ -1,3 +1,4 @@
 from .bark import BarkNotification
 from .server_chan import ServerChanNotification
-from .telegram import TelegramNotification
+from .telegram import TelegramNotification
+from .wecom import WecomNotification

src/module/notification/plugin/wecom.py

@@ -0,0 +1,35 @@
+import logging
+from module.network import RequestContent
+
+logger = logging.getLogger(__name__)
+
+
+class WecomNotification(RequestContent):
+    """企业微信推送 基于图文消息"""
+
+    def __init__(self, token, chat_id, **kwargs):
+        super().__init__()
+        #Chat_id is used as noti_url in this push tunnel
+        self.notification_url = f"{chat_id}"
+        self.token = token
+
+    def post_msg(self, text: str) -> bool:
+        ##Change message format to match Wecom push better
+        info = text.split("：")
+        print(info)
+        title = "【番剧更新】" + info[1].split("\n")[0].strip()
+        msg = info[2].split("\n")[0].strip()+" "+info[3].split("\n")[0].strip()
+        picurl = info[3].split("\n")[1].strip()
+        #Default pic to avoid blank in message. Resolution:1068*455
+        if picurl == "":
+            picurl = "https://article.biliimg.com/bfs/article/d8bcd0408bf32594fd82f27de7d2c685829d1b2e.png"
+        data = {
+            "key":self.token,                                       
+            "type": "news", 
+            "title": title,                                                              
+            "msg": msg, 
+            "picurl":picurl
+        }
+        resp = self.post_data(self.notification_url, data)
+        logger.debug(f"Wecom notification: {resp.status_code}")
+        return resp.status_code == 200

src/module/parser/analyser/torrent_parser.py

@@ -18,8 +18,8 @@ RULES = [
 ]
 
 SUBTITLE_LANG = {
-    "zh-tw": ["TC", "CHT", "繁", "zh-tw"],
-    "zh": ["SC", "CHS", "简", "zh"],
+    "zh-tw": ["TC", "CHT", "cht", "繁", "zh-tw"],
+    "zh": ["SC", "CHS", "chs", "简", "zh"],
 }