From ec20feca52a7cf44a89b17661b538f5b91151cc6 Mon Sep 17 00:00:00 2001 From: mhdzumair Date: Thu, 20 Feb 2025 10:12:34 +0530 Subject: [PATCH] refactor: streamline torrent file selection and metadata update process Add support for telegram notification when unable to detect episode files, refactor file selection and update db function --- db/models.py | 2 + streaming_providers/alldebrid/utils.py | 26 +- streaming_providers/debridlink/utils.py | 26 +- streaming_providers/easydebrid/utils.py | 12 +- streaming_providers/offcloud/client.py | 17 +- streaming_providers/offcloud/utils.py | 5 - streaming_providers/parser.py | 397 ++++++++++++++++++----- streaming_providers/pikpak/utils.py | 13 +- streaming_providers/premiumize/utils.py | 17 +- streaming_providers/qbittorrent/utils.py | 17 +- streaming_providers/realdebrid/utils.py | 41 +-- streaming_providers/seedr/utils.py | 14 +- streaming_providers/stremthru/utils.py | 20 +- streaming_providers/torbox/utils.py | 15 +- utils/telegram_bot.py | 10 + 15 files changed, 425 insertions(+), 207 deletions(-) diff --git a/db/models.py b/db/models.py index 216efae..0d3b890 100644 --- a/db/models.py +++ b/db/models.py @@ -123,6 +123,8 @@ class TorrentStreams(Document): torrent_type: Optional[TorrentType] = TorrentType.PUBLIC is_blocked: Optional[bool] = False torrent_file: bytes | None = None + known_file_details: dict[str, Any] | None = None + annotation_requested_at: Optional[datetime] = None @after_event(Insert) async def update_metadata_on_create(self): diff --git a/streaming_providers/alldebrid/utils.py b/streaming_providers/alldebrid/utils.py index 59b94bc..3185dfc 100644 --- a/streaming_providers/alldebrid/utils.py +++ b/streaming_providers/alldebrid/utils.py @@ -1,14 +1,11 @@ from typing import Optional -from fastapi import BackgroundTasks - from db.models import TorrentStreams from db.schemas import UserData from streaming_providers.alldebrid.client import AllDebrid from streaming_providers.exceptions import ProviderException from streaming_providers.parser import ( select_file_index_from_torrent, - update_torrent_streams_metadata, ) @@ -84,7 +81,6 @@ async def wait_for_download_and_get_link( episode, max_retries, retry_interval, - background_tasks, ): torrent_info = await ad_client.wait_for_status( torrent_id, "Ready", max_retries, retry_interval @@ -92,22 +88,14 @@ async def wait_for_download_and_get_link( files_data = {"files": flatten_files(torrent_info["files"])} file_index = await select_file_index_from_torrent( - files_data, - filename, - season, - episode, + torrent_info=files_data, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, + is_filename_trustable=True, ) - if filename is None: - background_tasks.add_task( - update_torrent_streams_metadata, - torrent_stream=stream, - torrent_info=files_data, - file_index=file_index, - season=season, - is_index_trustable=False, - ) - response = await ad_client.create_download_link( files_data["files"][file_index]["link"] ) @@ -121,7 +109,6 @@ async def get_video_url_from_alldebrid( stream: TorrentStreams, filename: Optional[str], user_ip: str, - background_tasks: BackgroundTasks, season: Optional[int] = None, episode: Optional[int] = None, max_retries=5, @@ -146,7 +133,6 @@ async def get_video_url_from_alldebrid( episode, max_retries, retry_interval, - background_tasks, ) diff --git a/streaming_providers/debridlink/utils.py b/streaming_providers/debridlink/utils.py index 5032c7b..6b05014 100644 --- a/streaming_providers/debridlink/utils.py +++ b/streaming_providers/debridlink/utils.py @@ -1,39 +1,31 @@ import asyncio from typing import Optional -from fastapi import BackgroundTasks - from db.models import TorrentStreams from db.schemas import UserData from streaming_providers.debridlink.client import DebridLink from streaming_providers.exceptions import ProviderException from streaming_providers.parser import ( select_file_index_from_torrent, - update_torrent_streams_metadata, ) async def get_download_link( torrent_info: dict, stream: TorrentStreams, - background_tasks, filename: Optional[str], - file_index: Optional[int], episode: Optional[int], season: Optional[int], ) -> str: selected_file_index = await select_file_index_from_torrent( - torrent_info, filename, season, episode + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, + is_filename_trustable=True, + is_index_trustable=True, ) - if filename is None or file_index is None: - background_tasks.add_task( - update_torrent_streams_metadata, - torrent_stream=stream, - torrent_info=torrent_info, - file_index=selected_file_index, - season=season, - is_index_trustable=True, - ) if torrent_info["files"][selected_file_index]["downloadPercent"] != 100: raise ProviderException( @@ -47,9 +39,7 @@ async def get_video_url_from_debridlink( magnet_link: str, user_data: UserData, filename: Optional[str], - file_index: Optional[int], stream: TorrentStreams, - background_tasks: BackgroundTasks, episode: Optional[int], season: Optional[int] = None, max_retries=5, @@ -86,9 +76,7 @@ async def get_video_url_from_debridlink( return await get_download_link( torrent_info, stream, - background_tasks, filename, - file_index, episode, season, ) diff --git a/streaming_providers/easydebrid/utils.py b/streaming_providers/easydebrid/utils.py index c7c2ce6..222561b 100644 --- a/streaming_providers/easydebrid/utils.py +++ b/streaming_providers/easydebrid/utils.py @@ -4,9 +4,9 @@ from typing import Any, Dict, List, Optional, Iterator from db.models import TorrentStreams from db.schemas import UserData +from streaming_providers.easydebrid.client import EasyDebrid from streaming_providers.exceptions import ProviderException from streaming_providers.parser import select_file_index_from_torrent -from streaming_providers.easydebrid.client import EasyDebrid async def get_video_url_from_easydebrid( @@ -14,6 +14,7 @@ async def get_video_url_from_easydebrid( user_data: UserData, filename: str, user_ip: str, + stream: TorrentStreams, season: Optional[int] = None, episode: Optional[int] = None, **kwargs: Any, @@ -33,10 +34,11 @@ async def get_video_url_from_easydebrid( ) file_index = await select_file_index_from_torrent( - torrent_info, - filename, - season, - episode, + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, name_key="filename", ) return torrent_info["files"][file_index]["url"] diff --git a/streaming_providers/offcloud/client.py b/streaming_providers/offcloud/client.py index 7204266..71f6b9c 100644 --- a/streaming_providers/offcloud/client.py +++ b/streaming_providers/offcloud/client.py @@ -4,13 +4,11 @@ from typing import Optional, List import aiohttp - from db.models import TorrentStreams from streaming_providers.debrid_client import DebridClient from streaming_providers.exceptions import ProviderException from streaming_providers.parser import ( select_file_index_from_torrent, - update_torrent_streams_metadata, ) @@ -151,7 +149,6 @@ class OffCloud(DebridClient): filename: Optional[str], season: Optional[int], episode: Optional[int], - background_tasks, ) -> str: if not torrent_info["isDirectory"]: return f"https://{torrent_info['server']}.offcloud.com/cloud/download/{request_id}/{torrent_info['fileName']}" @@ -160,24 +157,14 @@ class OffCloud(DebridClient): files_data = [{"name": path.basename(link), "link": link} for link in links] file_index = await select_file_index_from_torrent( - {"files": files_data}, + torrent_info={"files": files_data}, + torrent_stream=stream, filename=filename, season=season, episode=episode, file_size_callback=self.update_file_sizes, ) - if filename is None: - if "size" not in files_data[0]: - await self.update_file_sizes(files_data) - background_tasks.add_task( - update_torrent_streams_metadata, - torrent_stream=stream, - torrent_info={"files": files_data}, - file_index=file_index, - season=season, - is_index_trustable=False, - ) selected_file_url = links[file_index] return selected_file_url diff --git a/streaming_providers/offcloud/utils.py b/streaming_providers/offcloud/utils.py index 4d92fba..4753c1b 100644 --- a/streaming_providers/offcloud/utils.py +++ b/streaming_providers/offcloud/utils.py @@ -1,8 +1,6 @@ import asyncio from typing import List, Optional -from fastapi import BackgroundTasks - from db.models import TorrentStreams from db.schemas import UserData from streaming_providers.exceptions import ProviderException @@ -14,7 +12,6 @@ async def get_video_url_from_offcloud( magnet_link: str, user_data: UserData, stream: TorrentStreams, - background_tasks: BackgroundTasks, filename: Optional[str] = None, season: Optional[int] = None, episode: Optional[int] = None, @@ -36,7 +33,6 @@ async def get_video_url_from_offcloud( filename, season, episode, - background_tasks, ) if torrent_info["status"] == "error": raise ProviderException( @@ -64,7 +60,6 @@ async def get_video_url_from_offcloud( filename, season, episode, - background_tasks, ) diff --git a/streaming_providers/parser.py b/streaming_providers/parser.py index ed8c18f..ea51a10 100644 --- a/streaming_providers/parser.py +++ b/streaming_providers/parser.py @@ -1,122 +1,353 @@ import logging -from datetime import datetime, timezone -from typing import Any, Optional +from dataclasses import dataclass +from datetime import datetime, timezone, timedelta from os.path import basename +from typing import Any, Optional, TypedDict import PTT from db.models import TorrentStreams, EpisodeFile from streaming_providers.exceptions import ProviderException +from utils.telegram_bot import telegram_notifier from utils.validation_helper import is_video_file +logger = logging.getLogger(__name__) + + +class TorrentFile(TypedDict): + name: str + size: int + + +@dataclass +class FileInfo: + index: int + filename: str + size: int + season: Optional[int] = None + episode: Optional[int] = None + is_video: bool = False + + @classmethod + def from_torrent_file( + cls, + index: int, + file_data: dict[str, Any], + name_key: str = "name", + size_key: str = "size", + ) -> "FileInfo": + filename = basename(file_data[name_key]) + return cls( + index=index, + filename=filename, + size=file_data[size_key], + is_video=is_video_file(filename), + ) + + +class TorrentFileProcessor: + def __init__( + self, + torrent_info: dict[str, Any], + file_key: str = "files", + name_key: str = "name", + size_key: str = "size", + ): + self.torrent_info = torrent_info + self.file_key = file_key + self.name_key = name_key + self.size_key = size_key + self.files = torrent_info[file_key] + self.file_infos = self._process_files() + self._video_files: Optional[list[FileInfo]] = None + self._episodes: Optional[list[EpisodeFile]] = None + + def _process_files(self) -> list[FileInfo]: + """Process all files in the torrent and create FileInfo objects.""" + file_infos = [] + for idx, file in enumerate(self.files): + file_info = FileInfo.from_torrent_file( + idx, file, self.name_key, self.size_key + ) + if file_info.is_video: + # Don't parse season/episode info here - do it only when needed + file_infos.append(file_info) + + return file_infos + + def get_video_files(self) -> list[FileInfo]: + """Get all video files from the torrent, caching the result.""" + if self._video_files is None: + self._video_files = [f for f in self.file_infos if f.is_video] + return self._video_files + + def get_largest_video_file(self) -> Optional[FileInfo]: + """Get the largest video file from the torrent.""" + video_files = self.get_video_files() + return max(video_files, key=lambda x: x.size) if video_files else None + + def _parse_season_episode_info( + self, + file_info: FileInfo, + torrent_title: Optional[str] = None, + default_season: Optional[int] = None, + ) -> tuple[Optional[int], Optional[int]]: + """Parse season and episode information from filename and torrent title.""" + # First try from filename + parsed_data = PTT.parse_title(file_info.filename) + seasons = parsed_data.get("seasons", []) + episodes = parsed_data.get("episodes", []) + + season = seasons[0] if seasons else None + episode = episodes[0] if episodes else None + + # If no season/episode found, try from torrent title + if torrent_title and (not season or not episode): + title_parsed = PTT.parse_title(torrent_title) + title_seasons = title_parsed.get("seasons", []) + title_episodes = title_parsed.get("episodes", []) + + # If only one season in title, use it + if len(title_seasons) == 1: + season = season or title_seasons[0] + + # If only one episode in title, use it + if len(title_episodes) == 1: + episode = episode or title_episodes[0] + + # For season packs without explicit season number + if not season and episode: + season = default_season or 1 + + return season, episode + + def find_specific_episode( + self, season: int, episode: int, torrent_title: Optional[str] = None + ) -> Optional[FileInfo]: + """Find a specific episode in the video files.""" + if self._episodes: + for episode_file in self._episodes: + if ( + episode_file.season_number == season + and episode_file.episode_number == episode + ): + return self.file_infos[episode_file.file_index] + + for file_info in self.get_video_files(): + parsed_season, parsed_episode = self._parse_season_episode_info( + file_info, torrent_title, default_season=season + ) + if parsed_season == season and parsed_episode == episode: + return file_info + return None + + def parse_all_episodes( + self, torrent_title: str, default_season: Optional[int] = None + ) -> list[EpisodeFile]: + """Parse all episode information from video files.""" + episodes = [] + for file_info in self.get_video_files(): + season, episode = self._parse_season_episode_info( + file_info, torrent_title, default_season + ) + if season and episode: + episodes.append( + EpisodeFile( + season_number=season, + episode_number=episode, + filename=file_info.filename, + size=file_info.size, + file_index=file_info.index, + ) + ) + self._episodes = episodes + return episodes + + def find_file_by_name(self, filename: str) -> Optional[FileInfo]: + """Find a file by its filename.""" + for file_info in self.file_infos: + if file_info.filename == filename: + return file_info + return None + async def select_file_index_from_torrent( torrent_info: dict[str, Any], - filename: Optional[str], - season: Optional[int], + torrent_stream: TorrentStreams, + filename: Optional[str] = None, + season: Optional[int] = None, episode: Optional[int] = None, file_key: str = "files", name_key: str = "name", size_key: str = "size", file_size_callback: Optional[callable] = None, + is_filename_trustable: bool = False, + is_index_trustable: bool = False, ) -> int: - """Select the file index from the torrent info.""" + """ + Select the file index from the torrent info with minimal processing. + Only processes all files if initial filename match fails. + """ files = torrent_info[file_key] - if filename: - # Select the file with the matching filename - for index, file in enumerate(files): - if basename(file[name_key]) == filename and is_video_file(filename): - return index - if season and episode: - # Select the file with the matching episode number - for index, file in enumerate(files): - if not is_video_file(file[name_key]): - continue - season_parsed_data = PTT.parse_title(file[name_key]) - found_season = season_parsed_data.get("seasons") - if ( - found_season and season in found_season - ) or episode in season_parsed_data.get("episodes"): - return index - raise ProviderException( - "No matching file available for this torrent", "no_matching_file.mp4" + # Quick filename match without full processing + if filename: + for idx, file in enumerate(files): + file_name = basename(file[name_key]) + if file_name == filename: + return idx + + # Initialize processor for further processing + # get file sizes if callback provided + if file_size_callback: + await file_size_callback(files) + processor = TorrentFileProcessor(torrent_info, file_key, name_key, size_key) + + # Check if there are any video files + video_files = processor.get_video_files() + if not video_files: + if is_filename_trustable: + torrent_stream.is_blocked = True + await torrent_stream.save() + raise ProviderException( + "No valid video files found in torrent. Torrent has been blocked.", + "no_video_file.mp4", + ) + else: + raise ProviderException( + "No valid video files found in torrent", + "no_video_file.mp4", + ) + + # Update metadata if filename is trustable + if is_filename_trustable: + await update_torrent_streams_metadata( + torrent_stream=torrent_stream, + processor=processor, + season=season, + is_index_trustable=is_index_trustable, + is_filename_trustable=is_filename_trustable, ) - if file_size_callback: - # Get the file sizes - await file_size_callback(files) + # Try to find by season/episode + if season and episode: + selected_file = processor.find_specific_episode( + season, episode, torrent_stream.torrent_name + ) + if selected_file: + return selected_file.index + else: + # Found video files but couldn't match season/episode + await _request_annotation(torrent_stream, processor.file_infos) + raise ProviderException( + "Found video files but couldn't match season/episode. " + "Annotation has been requested.", + "episode_not_found.mp4", + ) - # If no file index is provided, select the largest file - largest_file = max(files, key=lambda file_: file_[size_key]) - index = files.index(largest_file) - if is_video_file(largest_file[name_key]): - return index + # Default to largest video file + selected_file = processor.get_largest_video_file() + if selected_file: + return selected_file.index raise ProviderException( - "No matching file available for this torrent", "no_matching_file.mp4" + "No valid video file found in torrent", + "no_video_file.mp4", ) async def update_torrent_streams_metadata( torrent_stream: TorrentStreams, - torrent_info: dict, - file_index: int, + processor: TorrentFileProcessor, season: Optional[int] = None, - file_key: str = "files", - name_key: str = "name", - size_key: str = "size", is_index_trustable: bool = False, -): - files_data = torrent_info[file_key] - if season is None: - torrent_stream.filename = basename(files_data[file_index][name_key]) - if file_index is not None and is_index_trustable: - torrent_stream.file_index = file_index - torrent_stream.updated_at = datetime.now(timezone.utc) - await torrent_stream.save() - logging.info(f"Updated {torrent_stream.id} metadata") + is_filename_trustable: bool = False, +) -> bool: + """ + Update torrent stream metadata using the provided processor instance. + Returns True if update was successful, False if annotation was requested. + """ + # Check if annotation was recently requested + if torrent_stream.annotation_requested_at and datetime.now( + tz=timezone.utc + ) - torrent_stream.annotation_requested_at < timedelta(days=7): + logger.info( + f"Skipping metadata update for {torrent_stream.id} - " + "recent annotation request" + ) + return False - else: - title_parsed_data = PTT.parse_title(torrent_stream.torrent_name) - episodes = [] - for idx, file in enumerate(files_data): - file_parsed_data = PTT.parse_title(file[name_key]) - season_number = file_parsed_data.get("seasons") or None - if ( - season_number is None - and title_parsed_data.get("seasons") - and len(title_parsed_data["seasons"]) == 1 - ): - season_number = title_parsed_data["seasons"][0] - episode_number = file_parsed_data.get("episodes") or None - if ( - episode_number is None - and title_parsed_data.get("episodes") - and len(title_parsed_data["episodes"]) == 1 - ): - episode_number = title_parsed_data["episodes"][0] + try: + if season is None: + # Movie handling + file_info = processor.get_largest_video_file() + if not file_info: + if is_filename_trustable: + torrent_stream.is_blocked = True + await torrent_stream.save() + logger.error(f"No video files found in torrent {torrent_stream.id}") + return False + await _request_annotation(torrent_stream, processor.file_infos) + return False - if not season_number or not episode_number: - continue + torrent_stream.filename = file_info.filename + if is_index_trustable: + torrent_stream.file_index = file_info.index + torrent_stream.size = file_info.size - episodes.append( - EpisodeFile( - season_number=season_number, - episode_number=episode_number, - filename=basename(file[name_key]), - size=file[size_key], - file_index=idx if is_index_trustable else None, - ) + else: + # Series handling + video_files = processor.get_video_files() + if not video_files: + if is_filename_trustable: + torrent_stream.is_blocked = True + await torrent_stream.save() + logger.error(f"No video files found in torrent {torrent_stream.id}") + return False + await _request_annotation(torrent_stream, processor.file_infos) + return False + + episodes = processor.parse_all_episodes( + torrent_stream.torrent_name, default_season=season ) - if not episodes: - return + if not episodes: + await _request_annotation(torrent_stream, processor.file_infos) + return False - torrent_stream.episode_files = episodes - total_size = sum(file[size_key] for file in files_data) - if total_size > torrent_stream.size: - torrent_stream.size = total_size - torrent_stream.updated_at = datetime.now() + torrent_stream.episode_files = episodes + total_size = sum(f.size for f in processor.get_video_files()) + if total_size > torrent_stream.size: + torrent_stream.size = total_size + + torrent_stream.updated_at = datetime.now(tz=timezone.utc) await torrent_stream.save() - logging.info(f"Updated {torrent_stream.id} episode data") + logger.info(f"Successfully updated {torrent_stream.id} metadata") + return True + + except Exception as e: + logger.error( + f"Error updating metadata for {torrent_stream.id}: {str(e)}", exc_info=True + ) + await _request_annotation(torrent_stream, processor.file_infos) + return False + + +async def _request_annotation( + torrent_stream: TorrentStreams, file_infos: list[FileInfo] +): + """Request manual annotation for the torrent stream.""" + file_details = [ + { + "filename": f.filename, + "size": f.size, + } + for f in file_infos + ] + torrent_stream.known_file_details = file_details + torrent_stream.annotation_requested_at = datetime.now(tz=timezone.utc) + await torrent_stream.save() + + # Notify contributors + await telegram_notifier.send_file_annotation_request(torrent_stream.id) + logger.info(f"Requested annotation for {torrent_stream.id}") diff --git a/streaming_providers/pikpak/utils.py b/streaming_providers/pikpak/utils.py index 492aa5f..e02962c 100644 --- a/streaming_providers/pikpak/utils.py +++ b/streaming_providers/pikpak/utils.py @@ -8,11 +8,11 @@ from pikpakapi import PikPakApi, PikpakException from db.config import settings from db.models import TorrentStreams +from db.redis_database import REDIS_ASYNC_CLIENT from db.schemas import UserData from streaming_providers.exceptions import ProviderException from streaming_providers.parser import select_file_index_from_torrent from utils import crypto -from db.redis_database import REDIS_ASYNC_CLIENT async def get_torrent_file_by_info_hash( @@ -123,6 +123,7 @@ async def find_file_in_folder_tree( my_pack_folder_id: str, info_hash: str, filename: str, + stream: TorrentStreams, season: int | None, episode: int | None, ) -> dict | None: @@ -138,7 +139,11 @@ async def find_file_in_folder_tree( files = await get_files_from_folder(pikpak, torrent_file["id"]) file_index = await select_file_index_from_torrent( - {"files": files}, filename, season, episode + torrent_info={"files": files}, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, ) return files[file_index] @@ -273,7 +278,7 @@ async def retrieve_or_download_file( retry_interval: int, ): selected_file = await find_file_in_folder_tree( - pikpak, my_pack_folder_id, info_hash, filename, season, episode + pikpak, my_pack_folder_id, info_hash, filename, stream, season, episode ) if not selected_file: await free_up_space(pikpak, stream.size) @@ -282,7 +287,7 @@ async def retrieve_or_download_file( pikpak, info_hash, max_retries, retry_interval ) selected_file = await find_file_in_folder_tree( - pikpak, my_pack_folder_id, info_hash, filename, season, episode + pikpak, my_pack_folder_id, info_hash, filename, stream, season, episode ) if selected_file is None: raise ProviderException( diff --git a/streaming_providers/premiumize/utils.py b/streaming_providers/premiumize/utils.py index 9aaa30e..9e35325 100644 --- a/streaming_providers/premiumize/utils.py +++ b/streaming_providers/premiumize/utils.py @@ -57,6 +57,7 @@ async def get_video_url_from_premiumize( user_ip: str, stream: TorrentStreams, filename: Optional[str], + season: Optional[int], episode: Optional[int], max_retries=5, retry_interval=5, @@ -75,7 +76,9 @@ async def get_video_url_from_premiumize( if "stream_link" in file_data ] } - return await get_stream_link(torrent_info, filename, episode) + return await get_stream_link( + torrent_info, filename, stream, season, episode + ) # Wait for file selection and then start torrent download torrent_info = await pm_client.wait_for_status( @@ -91,7 +94,7 @@ async def get_video_url_from_premiumize( if "video" in file_data.get("mime_type", "") ] } - return await get_stream_link(torrent_info, filename, episode) + return await get_stream_link(torrent_info, filename, stream, season, episode) async def fetch_downloaded_folder_data( @@ -110,13 +113,17 @@ async def fetch_downloaded_folder_data( async def get_stream_link( torrent_info: dict[str, Any], filename: Optional[str], + stream: TorrentStreams, + season: Optional[int] = None, episode: Optional[int] = None, ) -> str: """Get the stream link from the torrent info.""" selected_file_index = await select_file_index_from_torrent( - torrent_info, - filename, - episode, + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, ) selected_file = torrent_info["files"][selected_file_index] return selected_file["link"] diff --git a/streaming_providers/qbittorrent/utils.py b/streaming_providers/qbittorrent/utils.py index 73657c7..63565f0 100644 --- a/streaming_providers/qbittorrent/utils.py +++ b/streaming_providers/qbittorrent/utils.py @@ -2,7 +2,7 @@ import asyncio from contextlib import asynccontextmanager from datetime import timedelta from os import path -from typing import Optional +from typing import Optional, AsyncGenerator, Any from urllib.parse import urljoin, urlparse, quote from aiohttp import ( @@ -127,6 +127,7 @@ async def find_file_in_folder_tree( webdav: WebDavClient, user_data: UserData, info_hash: str, + stream: TorrentStreams, filename: Optional[str], season: Optional[int], episode: Optional[int], @@ -145,14 +146,20 @@ async def find_file_in_folder_tree( return None selected_file_index = await select_file_index_from_torrent( - {"files": files}, filename, season, episode + torrent_info={"files": files}, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, + is_filename_trustable=True, + is_index_trustable=True, ) selected_file = files[selected_file_index] return selected_file @asynccontextmanager -async def initialize_qbittorrent(user_data: UserData) -> APIClient: +async def initialize_qbittorrent(user_data: UserData) -> AsyncGenerator[APIClient, Any]: async with ClientSession( cookie_jar=CookieJar(unsafe=True), timeout=ClientTimeout(total=15) ) as session: @@ -262,7 +269,7 @@ async def retrieve_or_download_file( retry_interval: int, ): selected_file = await find_file_in_folder_tree( - webdav, user_data, info_hash, filename, season, episode + webdav, user_data, info_hash, stream, filename, season, episode ) if not selected_file: await set_qbittorrent_preferences(qbittorrent, stream.torrent_type) @@ -273,7 +280,7 @@ async def retrieve_or_download_file( qbittorrent, info_hash, play_video_after, max_retries, retry_interval ) selected_file = await find_file_in_folder_tree( - webdav, user_data, info_hash, filename, season, episode + webdav, user_data, info_hash, stream, filename, season, episode ) if not selected_file: raise ProviderException( diff --git a/streaming_providers/realdebrid/utils.py b/streaming_providers/realdebrid/utils.py index aea24d9..3839c8e 100644 --- a/streaming_providers/realdebrid/utils.py +++ b/streaming_providers/realdebrid/utils.py @@ -1,14 +1,11 @@ import asyncio from typing import Optional -from fastapi import BackgroundTasks - from db.models import TorrentStreams from db.schemas import UserData from streaming_providers.exceptions import ProviderException from streaming_providers.parser import ( select_file_index_from_torrent, - update_torrent_streams_metadata, ) from streaming_providers.realdebrid.client import RealDebrid @@ -18,41 +15,29 @@ async def create_download_link( magnet_link: str, torrent_info: dict, filename: Optional[str], - file_index: Optional[int], episode: Optional[int], season: Optional[int], stream: TorrentStreams, - background_tasks, max_retries: int, retry_interval: int, ) -> str: selected_file_index = await select_file_index_from_torrent( - torrent_info, - filename, - season, - episode, - "files", - "path", - "bytes", + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, + file_key="files", + name_key="path", + size_key="bytes", + is_filename_trustable=True, + is_index_trustable=True, ) - if filename is None or file_index is None: - background_tasks.add_task( - update_torrent_streams_metadata, - torrent_stream=stream, - torrent_info=torrent_info, - file_index=selected_file_index, - season=season, - file_key="files", - name_key="path", - size_key="bytes", - is_index_trustable=True, - ) - relevant_file = torrent_info["files"][selected_file_index] selected_files = [file for file in torrent_info["files"] if file["selected"] == 1] - if relevant_file not in selected_files or len(selected_files) != len( + if relevant_file.get("selected") != 1 or len(selected_files) != len( torrent_info["links"] ): await rd_client.delete_torrent(torrent_info["id"]) @@ -89,9 +74,7 @@ async def get_video_url_from_realdebrid( user_data: UserData, user_ip: str, filename: Optional[str], - file_index: Optional[int], stream: TorrentStreams, - background_tasks: BackgroundTasks, max_retries=5, retry_interval=5, episode: Optional[int] = None, @@ -146,11 +129,9 @@ async def get_video_url_from_realdebrid( magnet_link, torrent_info, filename, - file_index, episode, season, stream, - background_tasks, max_retries, retry_interval, ) diff --git a/streaming_providers/seedr/utils.py b/streaming_providers/seedr/utils.py index 392fbb7..dc050fe 100644 --- a/streaming_providers/seedr/utils.py +++ b/streaming_providers/seedr/utils.py @@ -1,11 +1,11 @@ import asyncio +import logging +import math import re from contextlib import asynccontextmanager from datetime import datetime -from typing import Optional, Dict, List, Any, Tuple from enum import Enum -import math -import logging +from typing import Optional, Dict, List, Any, Tuple, AsyncGenerator from aioseedrcc import Seedr from aioseedrcc.exception import SeedrException @@ -30,7 +30,7 @@ def clean_filename(name: Optional[str], replace: str = "") -> str: @asynccontextmanager -async def get_seedr_client(user_data: UserData) -> Seedr: +async def get_seedr_client(user_data: UserData) -> AsyncGenerator[Seedr, Any]: """Context manager that provides a Seedr client instance.""" try: async with Seedr(token=user_data.streaming_provider.token) as seedr: @@ -232,7 +232,11 @@ async def get_video_url_from_seedr( # Get file details folder_content = await get_files_from_folder(seedr, data["id"]) file_index = await select_file_index_from_torrent( - {"files": folder_content}, clean_filename(filename), season, episode + torrent_info={"files": folder_content}, + torrent_stream=stream, + filename=clean_filename(filename), + season=season, + episode=episode, ) selected_file = folder_content[file_index] diff --git a/streaming_providers/stremthru/utils.py b/streaming_providers/stremthru/utils.py index 55c0384..c71d58f 100644 --- a/streaming_providers/stremthru/utils.py +++ b/streaming_providers/stremthru/utils.py @@ -27,16 +27,24 @@ async def add_new_torrent(st_client, magnet_link): async def wait_for_download_and_get_link( - st_client, torrent_id, filename, season, episode, max_retries, retry_interval + st_client, + torrent_id, + filename, + stream, + season, + episode, + max_retries, + retry_interval, ): torrent_info = await st_client.wait_for_status( torrent_id, "downloaded", max_retries, retry_interval ) file_index = await select_file_index_from_torrent( - torrent_info, - filename, - season, - episode, + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, ) response = await st_client.create_download_link( torrent_info["files"][file_index]["link"] @@ -48,6 +56,7 @@ async def get_video_url_from_stremthru( magnet_link: str, user_data: UserData, filename: str, + stream: TorrentStreams, season: int = None, episode: int = None, max_retries=5, @@ -61,6 +70,7 @@ async def get_video_url_from_stremthru( st_client, torrent_id, filename, + stream, season, episode, max_retries, diff --git a/streaming_providers/torbox/utils.py b/streaming_providers/torbox/utils.py index 84a687b..c133c77 100644 --- a/streaming_providers/torbox/utils.py +++ b/streaming_providers/torbox/utils.py @@ -29,7 +29,7 @@ async def get_video_url_from_torbox( and torrent_info["download_present"] is True ): file_id = await select_file_id_from_torrent( - torrent_info, filename, season, episode + torrent_info, filename, stream, season, episode ) response = await torbox_client.create_download_link( torrent_info["id"], @@ -58,7 +58,7 @@ async def get_video_url_from_torbox( torrent_info = await torbox_client.get_available_torrent(info_hash) if torrent_info: file_id = await select_file_id_from_torrent( - torrent_info, filename, season, episode + torrent_info, filename, stream, season, episode ) response = await torbox_client.create_download_link( torrent_info["id"], @@ -129,16 +129,19 @@ async def fetch_downloaded_info_hashes_from_torbox( async def select_file_id_from_torrent( torrent_info: Dict[str, Any], filename: str, + stream: TorrentStreams, season: Optional[int], episode: Optional[int], ) -> int: """Select the file id from the torrent info.""" file_index = await select_file_index_from_torrent( - torrent_info, - filename, - season, - episode, + torrent_info=torrent_info, + torrent_stream=stream, + filename=filename, + season=season, + episode=episode, name_key="short_name", + is_filename_trustable=True, ) return torrent_info["files"][file_index]["id"] diff --git a/utils/telegram_bot.py b/utils/telegram_bot.py index 1236330..e6e1f7c 100644 --- a/utils/telegram_bot.py +++ b/utils/telegram_bot.py @@ -233,6 +233,16 @@ class TelegramNotifier: except Exception as e: logger.error(f"Error sending fallback text message: {e}") + async def send_file_annotation_request(self, info_hash: str): + """Send notification to request episode annotations""" + if not self.enabled: + logger.warning("Telegram notifications are disabled. Check bot token.") + return + + message = f"📝 File Annotation Requested\n\n" f"*Info Hash*: `{info_hash}`\n" + + await self._send_text_only_message(message) + # Create a singleton instance telegram_notifier = TelegramNotifier()