MediaFusion/utils/torrent.py

import hashlib
import logging
from contextlib import AsyncExitStack, asynccontextmanager
from datetime import datetime, timezone
from os.path import basename
from typing import Awaitable, Iterable, AsyncIterator, Optional, TypeVar, OrderedDict
from urllib.parse import quote

import PTT
import anyio
import bencodepy
import httpx
from anyio import (
    create_task_group,
    create_memory_object_stream,
    CapacityLimiter,
)
from anyio.streams.memory import MemoryObjectSendStream
from demagnetize.core import Demagnetizer
from torf import Magnet, MagnetError

import utils.runtime_const
from db.config import settings
from utils.parser import is_contain_18_plus_keywords
from utils.runtime_const import TRACKERS
from utils.validation_helper import is_video_file

# remove logging from demagnetize
logging.getLogger("demagnetize").setLevel(logging.CRITICAL)


def extract_torrent_metadata(
    content: bytes, parsed_data: dict = None, is_raise_error: bool = False
) -> dict:
    try:
        torrent_data: OrderedDict = bencodepy.decode(content)
        info = torrent_data[b"info"]
        info_encoded = bencodepy.encode(info)
        m = hashlib.sha1()
        m.update(info_encoded)
        info_hash = m.hexdigest()

        # 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)
        created_at = torrent_data.get(b"creation date", 0)

        announce_list = [
            tracker[0].decode() for tracker in torrent_data.get(b"announce-list", [])
        ]
        torrent_name = info.get(b"name", b"").decode()
        if not torrent_name:
            logging.warning("Torrent name is empty. Skipping")
            if is_raise_error:
                raise ValueError("Torrent name is empty")
            return {}

        metadata = {
            "info_hash": info_hash.lower(),
            "announce_list": announce_list,
            "total_size": total_size,
            "torrent_name": torrent_name,
            "torrent_file": content,
        }
        if parsed_data:
            metadata.update(parsed_data)
        else:
            metadata.update(PTT.parse_title(torrent_name, True))

        if is_contain_18_plus_keywords(torrent_name):
            logging.warning(
                f"Torrent name contains 18+ keywords: {torrent_name}. Skipping"
            )
            if is_raise_error:
                raise ValueError("Torrent name contains 18+ keywords")
            return {}

        if created_at:
            # Convert to UTC datetime
            metadata["created_at"] = datetime.fromtimestamp(created_at, tz=timezone.utc)

        file_data = []
        seasons = set()
        episodes = set()
        for idx, file in enumerate(files):
            full_path = (
                "/".join([p.decode() for p in file[b"path"]])
                if b"files" in info
                else None
            )
            filename = basename(full_path) if full_path else file[b"name"].decode()
            if not is_video_file(filename):
                continue
            if "sample" in filename.lower():
                logging.warning(f"Skipping sample file: {filename}")
                continue
            episode_parsed_data = PTT.parse_title(filename)
            seasons.update(episode_parsed_data.get("seasons", []))
            episodes.update(episode_parsed_data.get("episodes", []))
            season_number = (
                episode_parsed_data["seasons"][0]
                if episode_parsed_data.get("seasons")
                else None
            )
            if (
                season_number is None
                and metadata.get("seasons")
                and len(metadata["seasons"]) == 1
            ):
                season_number = metadata["seasons"][0]
            episode_number = (
                episode_parsed_data["episodes"][0]
                if episode_parsed_data.get("episodes")
                else None
            )
            if (
                episode_number is None
                and metadata.get("episodes")
                and len(metadata["episodes"]) == 1
            ):
                episode_number = metadata["episodes"][0]

            file_data.append(
                {
                    "filename": filename,
                    "size": file[b"length"],
                    "index": idx,
                    "season_number": season_number,
                    "episode_number": episode_number,
                    "episode_title": episode_parsed_data.get("title"),
                }
            )
        if not file_data:
            logging.warning(
                f"No video files found in torrent. Skipping. Found: {files}"
            )
            if is_raise_error:
                raise ValueError("No video files found in torrent")
            return {}

        largest_file = max(file_data, key=lambda x: x["size"])

        metadata.update(
            {
                "largest_file": largest_file,
                "file_data": file_data,
            }
        )

        if not metadata.get("seasons"):
            metadata["seasons"] = list(seasons)
        if not metadata.get("episodes"):
            metadata["episodes"] = list(episodes)

        return metadata
    except Exception as e:
        logging.exception(f"Error occurred: {e}")
        if is_raise_error:
            raise ValueError(f"Failed to extract torrent metadata from torrent: {e}")
        return {}


def convert_info_hash_to_magnet(info_hash: str, trackers: list[str]) -> str:
    magnet_link = f"magnet:?xt=urn:btih:{info_hash}"
    for tracker in set(trackers) or TRACKERS:
        encoded_tracker = quote(tracker, safe="")
        magnet_link += f"&tr={encoded_tracker}"
    return magnet_link


T = TypeVar("T")


@asynccontextmanager
async def acollect(
    coros: Iterable[Awaitable[T]],
    limit: Optional[CapacityLimiter] = None,
    timeout: int = 30,
) -> AsyncIterator[AsyncIterator[T]]:
    async with create_task_group() as tg:
        sender, receiver = create_memory_object_stream[T]()
        async with sender:
            for c in coros:
                tg.start_soon(_acollect_pipe, c, limit, sender.clone(), timeout)
        async with receiver:
            yield receiver


async def _acollect_pipe(
    coro: Awaitable[T],
    limit: Optional[CapacityLimiter],
    sender: MemoryObjectSendStream[T],
    timeout: int,
) -> None:
    async with AsyncExitStack() as stack:
        if limit is not None:
            await stack.enter_async_context(limit)
        await stack.enter_async_context(sender)
        try:
            with anyio.fail_after(timeout):
                value = await coro
            await sender.send(value)
        except Exception as e:
            # Send the exception instead of the value
            await sender.send(e)


async def info_hashes_to_torrent_metadata(
    info_hashes: list[str], trackers: list[str]
) -> list[dict]:
    torrents_data = []

    if not settings.enable_fetching_torrent_metadata_from_p2p:
        logging.info("Fetching torrent metadata from P2P is disabled")
        return torrents_data

    demagnetizer = Demagnetizer()
    async with acollect(
        coros=[
            demagnetizer.demagnetize(Magnet(xt=info_hash, tr=trackers or TRACKERS))
            for info_hash in info_hashes
        ],
        limit=CapacityLimiter(10),
        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.dump())
                    )
            except Exception as e:
                logging.error(f"Error processing torrent: {e}")

    return torrents_data


async def init_best_trackers():
    # get the best trackers from https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt

    try:
        async with httpx.AsyncClient(proxy=settings.requests_proxy_url) as client:
            response = await client.get(
                "https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt",
                timeout=30,
            )
            if response.status_code == 200:
                trackers = [tracker for tracker in response.text.split("\n") if tracker]
                utils.runtime_const.TRACKERS.extend(trackers)
                utils.runtime_const.TRACKERS = list(set(utils.runtime_const.TRACKERS))

                logging.info(
                    f"Loaded {len(trackers)} trackers. Total: {len(utils.runtime_const.TRACKERS)}"
                )
            else:
                logging.error(f"Failed to load trackers: {response.status_code}")
    except (httpx.ConnectTimeout, Exception) as e:
        logging.error(f"Failed to load trackers: {e}")


def parse_magnet(magnet_link: str) -> tuple[str, list[str]]:
    """
    Parse magnet link and return info hash and trackers
    """
    try:
        magnet = Magnet.from_string(magnet_link)
    except MagnetError:
        return "", []
    return magnet.infohash, magnet.tr


def get_info_hash_from_magnet(magnet_link: str) -> str:
    info_hash, _ = parse_magnet(magnet_link)
    return info_hash