variet-agent/tools/nyaa_client.py

"""Nyaa.si RSS 클라이언트 — 토렌트 검색 + Magnet 링크 생성.

RSS Feed: https://nyaa.si/?page=rss&q={query}&c={category}&f={filter}
"""

import httpx
import logging
import re
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from typing import Optional
from urllib.parse import quote

logger = logging.getLogger("variet.tools.nyaa")

RSS_BASE = "https://nyaa.si/"

# Nyaa RSS 네임스페이스
NYAA_NS = "https://nyaa.si/xmlns/nyaa"


@dataclass
class TorrentResult:
    """Nyaa 토렌트 검색 결과."""
    title: str
    torrent_url: str        # .torrent 다운로드 URL
    magnet_link: str        # magnet:?xt=urn:btih:...
    info_hash: str
    size: str
    seeders: int
    leechers: int
    downloads: int
    category: str
    pub_date: str
    view_url: str           # nyaa.si/view/... 페이지 URL

    # 파싱된 정보
    episode: Optional[int] = None
    group: str = ""


def _parse_episode(title: str) -> Optional[int]:
    """제목에서 에피소드 번호 추출.
    예: [ASW] Sousou no Frieren S2 - 07 [1080p ...] → 7
    """
    # 패턴 1: "- 07" 또는 "- 07v2"
    m = re.search(r'\s-\s(\d{1,4})(?:v\d)?(?:\s|\[|$)', title)
    if m:
        return int(m.group(1))
    # 패턴 2: "S02E07"
    m = re.search(r'[Ss]\d{1,2}[Ee](\d{1,4})', title)
    if m:
        return int(m.group(1))
    # 패턴 3: "Episode 07"
    m = re.search(r'[Ee]pisode\s*(\d{1,4})', title)
    if m:
        return int(m.group(1))
    return None


def _parse_group(title: str) -> str:
    """제목에서 릴리스 그룹명 추출. 예: [ASW] → ASW"""
    m = re.match(r'\[([^\]]+)\]', title)
    return m.group(1) if m else ""


class NyaaClient:
    """Nyaa.si RSS 기반 토렌트 검색 클라이언트."""

    def __init__(self, timeout: float = 15.0, default_suffix: str = "ASW HEVC"):
        self._timeout = timeout
        self.default_suffix = default_suffix

    async def search(
        self,
        query: str,
        category: str = "0_0",
        filter_: int = 0,
        use_default_suffix: bool = True,
    ) -> list[TorrentResult]:
        """RSS 기반 토렌트 검색.

        Args:
            query: 검색어
            category: Nyaa 카테고리 (0_0=전체, 1_2=Anime English)
            filter_: 필터 (0=없음, 2=trusted only)
            use_default_suffix: True면 검색어에 default_suffix 자동 추가
        """
        if use_default_suffix and self.default_suffix:
            full_query = f"{query} {self.default_suffix}"
        else:
            full_query = query

        url = f"{RSS_BASE}?page=rss&q={quote(full_query)}&c={category}&f={filter_}"
        logger.info(f"Nyaa RSS 검색: {full_query}")

        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.get(url)
            resp.raise_for_status()

        return self._parse_rss(resp.text)

    def _parse_rss(self, xml_text: str) -> list[TorrentResult]:
        """RSS XML 파싱."""
        root = ET.fromstring(xml_text)
        results = []

        for item in root.findall(".//item"):
            title = item.findtext("title", "")
            link = item.findtext("link", "")
            guid = item.findtext("guid", "")
            pub_date = item.findtext("pubDate", "")

            info_hash = item.findtext(f"{{{NYAA_NS}}}infoHash", "")
            seeders = int(item.findtext(f"{{{NYAA_NS}}}seeders", "0"))
            leechers = int(item.findtext(f"{{{NYAA_NS}}}leechers", "0"))
            downloads = int(item.findtext(f"{{{NYAA_NS}}}downloads", "0"))
            size = item.findtext(f"{{{NYAA_NS}}}size", "")
            category = item.findtext(f"{{{NYAA_NS}}}category", "")

            # Magnet 링크 생성
            magnet = f"magnet:?xt=urn:btih:{info_hash}" if info_hash else ""

            results.append(TorrentResult(
                title=title,
                torrent_url=link,
                magnet_link=magnet,
                info_hash=info_hash,
                size=size,
                seeders=seeders,
                leechers=leechers,
                downloads=downloads,
                category=category,
                pub_date=pub_date,
                view_url=guid,
                episode=_parse_episode(title),
                group=_parse_group(title),
            ))

        logger.info(f"Nyaa 검색 결과: {len(results)}건")
        return results

    async def search_anime(
        self,
        title: str,
        episode: Optional[int] = None,
    ) -> list[TorrentResult]:
        """애니 제목으로 검색. 에피소드 지정 시 필터링."""
        results = await self.search(title)

        if episode is not None:
            results = [r for r in results if r.episode == episode]

        # 시더 수 내림차순 정렬
        results.sort(key=lambda r: r.seeders, reverse=True)
        return results


# ── CLI 진입점 ──
if __name__ == "__main__":
    import sys
    import asyncio

    args = sys.argv[1:]
    client = NyaaClient()

    async def main():
        if not args or args[0] == "search":
            # python tools/nyaa_client.py search "Sousou no Frieren" [--suffix "ASW HEVC"]
            query_parts = []
            suffix = "ASW HEVC"
            i = 1 if args and args[0] == "search" else 0
            while i < len(args):
                if args[i] == "--suffix" and i + 1 < len(args):
                    suffix = args[i + 1]
                    i += 2
                elif args[i] == "--no-suffix":
                    suffix = ""
                    i += 1
                else:
                    query_parts.append(args[i])
                    i += 1

            if not query_parts:
                print("사용법: python tools/nyaa_client.py search \"제목\" [--suffix \"ASW HEVC\"]")
                return

            query = " ".join(query_parts)
            client.default_suffix = suffix
            results = await client.search(query, use_default_suffix=bool(suffix))

            print(f"🔍 Nyaa 검색: '{query}' +'{suffix}' → {len(results)}건")
            for r in results[:20]:
                ep = f" {r.episode}화" if r.episode else ""
                print(f"  [{r.group}] {r.title[:60]}... | {r.size} | S:{r.seeders}{ep}")
                print(f"    magnet: {r.magnet_link[:80]}...")
        else:
            print("사용법: python tools/nyaa_client.py search \"제목\" [--suffix \"ASW HEVC\"]")

    asyncio.run(main())