From e68e80e7ee4f4e70946fa81957fbaa3b8d062bf2 Mon Sep 17 00:00:00 2001 From: mhdzumair Date: Wed, 7 Feb 2024 18:32:34 +0530 Subject: [PATCH] Improve Prowlarr integration and improve database initialization for dramtiq actor --- api/main.py | 3 +- api/task.py | 4 +++ db/database.py | 47 +++++++++++++++++++++--------- scrappers/helpers.py | 2 -- scrappers/prowlarr.py | 61 ++++++++++++++++++++++---------------- utils/torrent.py | 68 +++++++++++++++++++++++-------------------- 6 files changed, 111 insertions(+), 74 deletions(-) diff --git a/api/main.py b/api/main.py index a38bbf9..26b9ac5 100644 --- a/api/main.py +++ b/api/main.py @@ -21,7 +21,7 @@ from fastapi.templating import Jinja2Templates from api.middleware import SecureLoggingMiddleware from db import database, crud, schemas from db.config import settings -from scrappers import tamil_blasters, tamilmv +from scrappers import tamil_blasters, tamilmv, prowlarr from streaming_providers.alldebrid.utils import get_direct_link_from_alldebrid from streaming_providers.debridlink.api import router as debridlink_router from streaming_providers.debridlink.utils import get_direct_link_from_debridlink @@ -77,6 +77,7 @@ async def init_server(): app.state.redis = redis.Redis( connection_pool=redis.ConnectionPool.from_url(settings.redis_url) ) + await prowlarr.init_indexers() @app.on_event("startup") diff --git a/api/task.py b/api/task.py index aec38e8..867adc5 100644 --- a/api/task.py +++ b/api/task.py @@ -1,6 +1,8 @@ # import background actors import asyncio +from db import database +from scrappers import prowlarr from scrappers.helpers import update_torrent_movie_streams_metadata # noqa: F401 from scrappers.prowlarr import parse_and_store_movie_stream_data # noqa: F401 from utils import torrent @@ -9,6 +11,8 @@ from utils import torrent async def async_setup(): # Your async initialization code here await torrent.init_best_trackers() + await prowlarr.init_indexers() + await database.init() asyncio.run(async_setup()) diff --git a/db/database.py b/db/database.py index 8b96f03..68d1b40 100644 --- a/db/database.py +++ b/db/database.py @@ -1,3 +1,6 @@ +import asyncio +import logging + import motor.motor_asyncio from beanie import init_beanie @@ -13,17 +16,33 @@ from db.models import ( async def init(): - # Create Motor client - client = motor.motor_asyncio.AsyncIOMotorClient(settings.mongo_uri) - # Init beanie with the Product document class - await init_beanie( - database=client.get_default_database(), # Note that the database needs to be passed as part of the URI - document_models=[ - MediaFusionMovieMetaData, - MediaFusionSeriesMetaData, - TorrentStreams, - TVStreams, - MediaFusionTVMetaData, - SearchHistory, - ], - ) + retries = 5 + for i in range(retries): + try: + # Create Motor client + client = motor.motor_asyncio.AsyncIOMotorClient(settings.mongo_uri) + # Init beanie with the Product document class + await init_beanie( + database=client.get_default_database(), # Note that the database needs to be passed as part of the URI + document_models=[ + MediaFusionMovieMetaData, + MediaFusionSeriesMetaData, + TorrentStreams, + TVStreams, + MediaFusionTVMetaData, + SearchHistory, + ], + multiprocessing_mode=True, + ) + logging.info("Database initialized successfully.") + break + except Exception as e: + if i < retries - 1: # i is zero indexed + wait_time = 2**i # exponential backoff + logging.warning( + f"Error initializing database: {e}, retrying in {wait_time} seconds..." + ) + await asyncio.sleep(wait_time) + else: + logging.error("Failed to initialize database after several attempts.") + raise e diff --git a/scrappers/helpers.py b/scrappers/helpers.py index c604bdb..2f26789 100644 --- a/scrappers/helpers.py +++ b/scrappers/helpers.py @@ -131,7 +131,6 @@ def get_scrapper_config(site_name: str, get_key: str) -> dict: @dramatiq.actor(time_limit=30 * 60 * 1000) async def update_torrent_movie_streams_metadata(info_hashes: list[str]): """Update torrent streams metadata.""" - await database.init() streams_metadata = await info_hashes_to_torrent_metadata(info_hashes, []) for stream_metadata in streams_metadata: @@ -153,7 +152,6 @@ async def update_torrent_movie_streams_metadata(info_hashes: list[str]): @dramatiq.actor(time_limit=30 * 60 * 1000) async def update_torrent_series_streams_metadata(info_hashes: list[str]): """Update torrent streams metadata.""" - await database.init() streams_metadata = await info_hashes_to_torrent_metadata(info_hashes, []) for stream_metadata in streams_metadata: diff --git a/scrappers/prowlarr.py b/scrappers/prowlarr.py index b40187d..0f1c190 100644 --- a/scrappers/prowlarr.py +++ b/scrappers/prowlarr.py @@ -1,4 +1,5 @@ import asyncio +import json import logging from datetime import datetime, timedelta @@ -6,9 +7,9 @@ import PTN import dramatiq import httpx from redis.asyncio import Redis +from thefuzz import fuzz from torf import Magnet -from db import database from db.config import settings from db.models import TorrentStreams, Season, Episode from scrappers.helpers import ( @@ -213,9 +214,20 @@ async def get_torrent_data_from_prowlarr(download_url: str) -> tuple[dict, bool] async def prowlarr_data_parser(meta_data: dict) -> tuple[dict, bool]: """Parse prowlarr data.""" + if meta_data.get("indexer") in [ + "Torlock", + "YourBittorrent", + "The Pirate Bay", + "BitSearch", + ]: + # For these indexers, the guid is a direct torrent file download link + download_url = meta_data.get("guid") + else: + download_url = meta_data.get("downloadUrl") or meta_data.get("magnetUrl") + try: torrent_data, is_torrent_downloaded = await get_torrent_data_from_prowlarr( - meta_data.get("downloadUrl") or meta_data.get("magnetUrl") + download_url ) except Exception as e: if meta_data.get("infoHash"): @@ -236,12 +248,16 @@ async def prowlarr_data_parser(meta_data: dict) -> tuple[dict, bool]: e, httpx.HTTPError, ): - raise e + return {}, False logging.error( f"Error getting torrent data: {e} {e.__class__.__name__}", exc_info=True ) return {}, False + info_hash = torrent_data.get("info_hash") + if not info_hash: + return {}, False + torrent_data.update( { "seeders": meta_data.get("seeders"), @@ -276,7 +292,7 @@ async def handle_movie_stream_store(info_hash, parsed_data, video_id): prowlarr_catalogs = [ "prowlarr_streams", "prowlarr_movies", - f"{parsed_data.get('source').lower()}_movies", + f"{parsed_data.get('source', '').lower()}_movies", ] if torrent_stream: @@ -365,7 +381,7 @@ async def handle_series_stream_store(info_hash, parsed_data, video_id, season): prowlarr_catalog = [ "prowlarr_streams", "prowlarr_series", - f"{parsed_data.get('source').lower()}_series", + f"{parsed_data.get('source', '').lower()}_series", ] if torrent_stream: @@ -429,22 +445,23 @@ async def parse_and_store_stream( info_hash = parsed_data.get("info_hash", "").lower() torrent_stream, torrent_needed_update = None, False - if not info_hash or parsed_data.get("seeders", 0) == 0: + if not info_hash: logging.warning( - f"Skipping {info_hash} due to missing info_hash or seeders: {parsed_data.get('seeders')}" + f"Skipping {stream_data.get('title')} due to missing info_hash." ) return torrent_stream, torrent_needed_update + title_similarity_ratio = fuzz.ratio( + parsed_data.get("title", "").lower(), title.lower() + ) + if catalog_type == "movie": if ( - not ( - parsed_data.get("title").lower() == title.lower() - and parsed_data.get("year") == year - ) + not (title_similarity_ratio > 80 and parsed_data.get("year") == year) and parsed_data.get("imdb_id") != video_id ): logging.warning( - f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' or year mismatch: '{parsed_data.get('year')}' != '{year}'" + f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' ratio: {title_similarity_ratio} or year mismatch: '{parsed_data.get('year')}' != '{year}'" ) return torrent_stream, torrent_needed_update @@ -452,12 +469,9 @@ async def parse_and_store_stream( info_hash, parsed_data, video_id ) elif catalog_type == "series": - if ( - parsed_data.get("title").lower() != title.lower() - and parsed_data.get("imdb_id") != video_id - ): + if title_similarity_ratio < 80 and parsed_data.get("imdb_id") != video_id: logging.warning( - f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}'" + f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' ratio: {title_similarity_ratio}" ) return torrent_stream, torrent_needed_update @@ -509,9 +523,8 @@ async def parse_and_store_movie_stream_data_actor( title: str, year: str, stream_data: list, -) -> list[TorrentStreams]: - await database.init() - return await parse_and_store_movie_stream_data(video_id, title, year, stream_data) +): + await parse_and_store_movie_stream_data(video_id, title, year, stream_data) async def parse_and_store_series_stream_data( @@ -555,8 +568,6 @@ async def parse_and_store_series_stream_data_actor( title: str, season: int, stream_data: list, -) -> list[TorrentStreams]: - await database.init() - return await parse_and_store_series_stream_data( - video_id, title, season, stream_data - ) +): + await parse_and_store_series_stream_data(video_id, title, season, stream_data) + diff --git a/utils/torrent.py b/utils/torrent.py index 1911b56..c214a13 100644 --- a/utils/torrent.py +++ b/utils/torrent.py @@ -1,4 +1,4 @@ -import io +import hashlib import logging from contextlib import AsyncExitStack, asynccontextmanager from typing import Awaitable, Iterable, AsyncIterator, Optional, TypeVar @@ -6,6 +6,7 @@ from urllib.parse import quote import PTN import anyio +import bencodepy import httpx from anyio import ( create_task_group, @@ -14,7 +15,7 @@ from anyio import ( ) from anyio.streams.memory import MemoryObjectSendStream from demagnetize.core import Demagnetizer -from torf import Magnet, Torrent +from torf import Magnet # remove logging from demagnetize logging.getLogger("demagnetize").setLevel(logging.CRITICAL) @@ -35,50 +36,50 @@ TRACKERS = [ ] -def response_content_to_torrent(content) -> Torrent | None: +def extract_torrent_metadata(content: bytes) -> dict: try: - return Torrent.read_stream(io.BytesIO(content)) - except Exception as e: - logging.error(f"Error occurred: {e}") - return + torrent_data = bencodepy.decode(content) + info = torrent_data[b"info"] + info_encoded = bencodepy.encode(info) + m = hashlib.sha1() + m.update(info_encoded) + info_hash = m.hexdigest() - -def extract_torrent_metadata(content: Torrent | bytes) -> dict: - if isinstance(content, bytes): - torrent = response_content_to_torrent(content) - if not torrent: - return {} - else: - torrent = content - - try: - info_hash = torrent.infohash - total_size = torrent.size + # Extract file size, file list, and announce list + files = info[b"files"] if b"files" in info else [info] + total_size = sum(file[b"length"] for file in files) file_data = [] - for idx, file in enumerate( - torrent.files if torrent.mode == "multifile" else [torrent.name] - ): - filename = str(file).split("/")[-1] + for idx, file in enumerate(files): + filename = ( + file[b"path"][0].decode() + if b"files" in info + else file[b"name"].decode() + ) parsed_data = PTN.parse(filename) file_data.append( { "filename": filename, - "size": file.size if torrent.mode == "multifile" else total_size, + "size": file[b"length"], "index": idx, "season": parsed_data.get("season"), "episode": parsed_data.get("episode"), } ) + + announce_list = [ + tracker[0].decode() for tracker in torrent_data.get(b"announce-list", []) + ] + torrent_name = info.get(b"name", b"").decode() or file_data[0]["filename"] largest_file = max(file_data, key=lambda x: x["size"]) return { - **PTN.parse(torrent.name), + **PTN.parse(torrent_name), "info_hash": info_hash, - "announce_list": sorted(list(torrent.trackers.flat)), + "announce_list": announce_list, "total_size": total_size, "file_data": file_data, - "torrent_name": torrent.name, + "torrent_name": torrent_name, "largest_file": largest_file, } except Exception as e: @@ -134,21 +135,24 @@ async def _acollect_pipe( async def info_hashes_to_torrent_metadata( info_hashes: list[str], trackers: list[str] ) -> list[dict]: - magnets = [ - Magnet(xt=info_hash, tr=trackers or TRACKERS) for info_hash in info_hashes - ] torrents_data = [] demagnetizer = Demagnetizer() async with acollect( - [demagnetizer.demagnetize(magnet) for magnet in magnets], timeout=60 + [ + demagnetizer.demagnetize(Magnet(xt=info_hash, tr=trackers or TRACKERS)) + for info_hash in info_hashes + ], + timeout=60, ) as async_iterator: async for torrent_result in async_iterator: try: if isinstance(torrent_result, Exception): pass else: - torrents_data.append(extract_torrent_metadata(torrent_result)) + torrents_data.append( + extract_torrent_metadata(torrent_result.dump()) + ) except Exception as e: logging.error(f"Error processing torrent: {e}")