mirror of
https://github.com/g0ldyy/comet.git
synced 2026-01-12 01:16:12 +01:00
g0ldyy
649 lines
20 KiB
Python
649 lines
20 KiB
Python
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import re
|
|
import time
|
|
from asyncio import QueueEmpty
|
|
from collections import defaultdict
|
|
from urllib.parse import unquote
|
|
|
|
import anyio
|
|
import bencodepy
|
|
import orjson
|
|
import xxhash
|
|
from demagnetize.core import Demagnetizer
|
|
from RTN import ParsedData, parse
|
|
from torf import Magnet
|
|
|
|
from comet.core.constants import TORRENT_TIMEOUT
|
|
from comet.core.logger import logger
|
|
from comet.core.models import database, settings
|
|
from comet.utils.parsing import default_dump, is_video
|
|
|
|
TRACKER_PATTERN = re.compile(r"[&?]tr=([^&]+)")
|
|
INFO_HASH_PATTERN = re.compile(r"btih:([a-fA-F0-9]{40}|[a-zA-Z0-9]{32})")
|
|
|
|
|
|
def extract_trackers_from_magnet(magnet_uri: str):
|
|
try:
|
|
trackers = TRACKER_PATTERN.findall(magnet_uri)
|
|
return [unquote(tracker) for tracker in trackers]
|
|
except Exception as e:
|
|
logger.warning(f"Failed to extract trackers from magnet URI: {e}")
|
|
return []
|
|
|
|
|
|
async def download_torrent(session, url: str):
|
|
try:
|
|
async with session.get(
|
|
url, allow_redirects=False, timeout=TORRENT_TIMEOUT
|
|
) as response:
|
|
if response.status == 200:
|
|
return (await response.read(), None, None)
|
|
|
|
location = response.headers.get("Location", "")
|
|
if location:
|
|
match = INFO_HASH_PATTERN.search(location)
|
|
if match:
|
|
info_hash = match.group(1)
|
|
if len(info_hash) == 32:
|
|
info_hash = base64.b16encode(
|
|
base64.b32decode(info_hash)
|
|
).decode("utf-8")
|
|
return (None, info_hash, location)
|
|
return (None, None, None)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Failed to download torrent from {url}: {e} (in most cases, you can ignore this error)"
|
|
)
|
|
return (None, None, None)
|
|
|
|
|
|
demagnetizer = Demagnetizer()
|
|
|
|
|
|
async def get_torrent_from_magnet(magnet_uri: str):
|
|
try:
|
|
magnet = Magnet.from_string(magnet_uri)
|
|
with anyio.fail_after(settings.MAGNET_RESOLVE_TIMEOUT):
|
|
torrent_data = await demagnetizer.demagnetize(magnet)
|
|
if torrent_data:
|
|
return torrent_data.dump()
|
|
except Exception as e:
|
|
logger.warning(f"Failed to get torrent from magnet: {e}")
|
|
return None
|
|
|
|
|
|
def extract_torrent_metadata(content: bytes):
|
|
try:
|
|
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()
|
|
|
|
announce_list = [
|
|
tracker[0].decode() for tracker in torrent_data.get(b"announce-list", [])
|
|
]
|
|
announce = torrent_data.get(b"announce", b"").decode()
|
|
if announce:
|
|
announce_list.append(announce)
|
|
|
|
metadata = {"info_hash": info_hash, "announce_list": announce_list, "files": []}
|
|
|
|
files = info[b"files"] if b"files" in info else [info]
|
|
for idx, file in enumerate(files):
|
|
name = (
|
|
file[b"path"][-1].decode()
|
|
if b"path" in file
|
|
else file[b"name"].decode()
|
|
)
|
|
|
|
if not is_video(name) or "sample" in name.lower():
|
|
continue
|
|
|
|
size = file[b"length"]
|
|
|
|
metadata["files"].append({"index": idx, "name": name, "size": size})
|
|
|
|
return metadata
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to extract torrent metadata: {e}")
|
|
return {}
|
|
|
|
|
|
async def add_torrent(
|
|
info_hash: str,
|
|
seeders: int,
|
|
tracker: str,
|
|
media_id: str,
|
|
search_season: int,
|
|
sources: list,
|
|
file_index: int,
|
|
title: str,
|
|
size: int,
|
|
parsed: ParsedData,
|
|
):
|
|
try:
|
|
seasons_to_process = parsed.seasons if parsed.seasons else [search_season]
|
|
parsed_episodes = parsed.episodes if parsed.episodes else [None]
|
|
|
|
episode_to_insert = parsed_episodes[0] if len(parsed_episodes) == 1 else None
|
|
|
|
for season in seasons_to_process:
|
|
await _upsert_torrent_record(
|
|
{
|
|
"media_id": media_id,
|
|
"info_hash": info_hash,
|
|
"file_index": file_index,
|
|
"season": season,
|
|
"episode": episode_to_insert,
|
|
"title": title,
|
|
"seeders": seeders,
|
|
"size": size,
|
|
"tracker": tracker,
|
|
"sources": orjson.dumps(sources).decode("utf-8"),
|
|
"parsed": orjson.dumps(parsed, default_dump).decode("utf-8"),
|
|
"timestamp": time.time(),
|
|
}
|
|
)
|
|
|
|
additional = ""
|
|
if seasons_to_process:
|
|
additional += f" - S{seasons_to_process[0]:02d}"
|
|
if parsed_episodes and parsed_episodes[0] is not None:
|
|
additional += f"E{parsed_episodes[0]:02d}"
|
|
|
|
logger.log("SCRAPER", f"Added torrent for {media_id} - {title}{additional}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to add torrent for {info_hash}: {e}")
|
|
|
|
|
|
class AddTorrentQueue:
|
|
def __init__(self, max_concurrent: int = 10):
|
|
self.queue = asyncio.Queue()
|
|
self.max_concurrent = max_concurrent
|
|
self.is_running = False
|
|
self.semaphore = asyncio.Semaphore(max_concurrent)
|
|
|
|
async def add_torrent(
|
|
self,
|
|
magnet_url: str,
|
|
seeders: int,
|
|
tracker: str,
|
|
media_id: str,
|
|
search_season: int,
|
|
):
|
|
if not settings.DOWNLOAD_TORRENT_FILES:
|
|
return
|
|
|
|
await self.queue.put((magnet_url, seeders, tracker, media_id, search_season))
|
|
if not self.is_running:
|
|
self.is_running = True
|
|
asyncio.create_task(self._process_queue())
|
|
|
|
async def _process_queue(self):
|
|
while self.is_running:
|
|
try:
|
|
(
|
|
magnet_url,
|
|
seeders,
|
|
tracker,
|
|
media_id,
|
|
search_season,
|
|
) = await self.queue.get()
|
|
|
|
async with self.semaphore:
|
|
try:
|
|
content = await get_torrent_from_magnet(magnet_url)
|
|
if content:
|
|
metadata = extract_torrent_metadata(content)
|
|
for file in metadata["files"]:
|
|
parsed = parse(file["name"])
|
|
|
|
await add_torrent(
|
|
metadata["info_hash"],
|
|
seeders,
|
|
tracker,
|
|
media_id,
|
|
search_season,
|
|
metadata["announce_list"],
|
|
file["index"],
|
|
file["name"],
|
|
file["size"],
|
|
parsed,
|
|
)
|
|
finally:
|
|
self.queue.task_done()
|
|
|
|
except Exception:
|
|
await asyncio.sleep(1)
|
|
|
|
self.is_running = False
|
|
|
|
async def stop(self):
|
|
await self.queue.join()
|
|
self.is_running = False
|
|
|
|
|
|
add_torrent_queue = AddTorrentQueue()
|
|
|
|
|
|
UPDATE_INTERVAL = (
|
|
settings.TORRENT_CACHE_TTL // 2 if settings.TORRENT_CACHE_TTL >= 0 else 31536000
|
|
)
|
|
|
|
|
|
class TorrentUpdateQueue:
|
|
__slots__ = (
|
|
"queue",
|
|
"batch_size",
|
|
"flush_interval",
|
|
"is_running",
|
|
"_lock",
|
|
"_event",
|
|
"upserts",
|
|
"_is_postgresql",
|
|
"_grouped_upserts",
|
|
)
|
|
|
|
def __init__(self, batch_size: int = 1000, flush_interval: float = 5.0):
|
|
self.queue = asyncio.Queue()
|
|
self.batch_size = batch_size
|
|
self.flush_interval = flush_interval
|
|
self.is_running = False
|
|
self._lock = asyncio.Lock()
|
|
self._event = asyncio.Event()
|
|
self.upserts = {}
|
|
self._is_postgresql = settings.DATABASE_TYPE == "postgresql"
|
|
self._grouped_upserts = defaultdict(list)
|
|
|
|
async def add_torrent_info(self, file_info: dict, media_id: str = None):
|
|
await self.queue.put((file_info, media_id))
|
|
self._event.set()
|
|
|
|
if not self.is_running:
|
|
async with self._lock:
|
|
if not self.is_running:
|
|
self.is_running = True
|
|
asyncio.create_task(self._process_queue())
|
|
|
|
async def _process_queue(self):
|
|
last_flush_time = time.time()
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._event.wait(), timeout=self.flush_interval
|
|
)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
self._event.clear()
|
|
|
|
batch_time = time.time()
|
|
items_processed = 0
|
|
while True:
|
|
try:
|
|
file_info, media_id = self.queue.get_nowait()
|
|
self._process_file_info(file_info, media_id, batch_time)
|
|
self.queue.task_done()
|
|
items_processed += 1
|
|
|
|
if len(self.upserts) >= self.batch_size:
|
|
await self._flush_batch()
|
|
last_flush_time = time.time()
|
|
batch_time = last_flush_time
|
|
except QueueEmpty:
|
|
break
|
|
|
|
current_time = time.time()
|
|
|
|
if self.upserts and (
|
|
current_time - last_flush_time >= self.flush_interval
|
|
):
|
|
await self._flush_batch()
|
|
last_flush_time = current_time
|
|
|
|
if self.queue.empty() and not self.upserts:
|
|
break
|
|
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.warning(f"Error in _process_queue: {e}")
|
|
finally:
|
|
if self.upserts:
|
|
await self._flush_batch()
|
|
if self.is_running:
|
|
async with self._lock:
|
|
self.is_running = False
|
|
|
|
async def stop(self):
|
|
self.is_running = False
|
|
self._event.set()
|
|
|
|
shutdown_time = time.time()
|
|
while True:
|
|
try:
|
|
file_info, media_id = self.queue.get_nowait()
|
|
self._process_file_info(file_info, media_id, shutdown_time)
|
|
except QueueEmpty:
|
|
break
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error processing remaining queue items during shutdown: {e}"
|
|
)
|
|
break
|
|
|
|
if self.upserts:
|
|
await self._flush_batch()
|
|
|
|
async def _flush_batch(self):
|
|
if not self.upserts:
|
|
return
|
|
|
|
upserts_to_flush = self.upserts
|
|
self.upserts = {}
|
|
|
|
try:
|
|
grouped = self._grouped_upserts
|
|
for params in upserts_to_flush.values():
|
|
key = _determine_conflict_key(params["season"], params["episode"])
|
|
grouped[key].append(params)
|
|
|
|
for key, rows in grouped.items():
|
|
query = _get_torrent_upsert_query(key)
|
|
try:
|
|
await _execute_batched_upsert(query, rows)
|
|
except Exception as e:
|
|
logger.warning(f"Error processing upsert batch: {e}")
|
|
|
|
total_upserts = len(upserts_to_flush)
|
|
if total_upserts > 0:
|
|
logger.log(
|
|
"SCRAPER",
|
|
f"Upserted {total_upserts} torrents in batch",
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Error in flush_batch: {e}")
|
|
finally:
|
|
grouped.clear()
|
|
|
|
def _process_file_info(
|
|
self, file_info: dict, media_id: str = None, current_time: float = None
|
|
):
|
|
try:
|
|
info_hash = file_info["info_hash"]
|
|
season = file_info["season"]
|
|
episode = file_info["episode"]
|
|
|
|
if current_time is None:
|
|
current_time = time.time()
|
|
|
|
upsert_key = (media_id, info_hash, season, episode)
|
|
|
|
existing = self.upserts.get(upsert_key)
|
|
if existing and existing["timestamp"] >= current_time:
|
|
return
|
|
|
|
params = {
|
|
"info_hash": info_hash,
|
|
"file_index": file_info["index"],
|
|
"season": season,
|
|
"episode": episode,
|
|
"title": file_info["title"],
|
|
"seeders": file_info["seeders"],
|
|
"size": file_info["size"],
|
|
"tracker": file_info["tracker"],
|
|
"sources": orjson.dumps(file_info["sources"]).decode("utf-8"),
|
|
"parsed": orjson.dumps(
|
|
file_info["parsed"], default=default_dump
|
|
).decode("utf-8"),
|
|
"timestamp": current_time,
|
|
"media_id": media_id,
|
|
}
|
|
|
|
if self._is_postgresql:
|
|
params["update_interval"] = UPDATE_INTERVAL
|
|
|
|
params["lock_key"] = _compute_advisory_lock_key(
|
|
media_id, info_hash, season, episode
|
|
)
|
|
|
|
self.upserts[upsert_key] = params
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error processing file info: {e}")
|
|
|
|
|
|
TORRENT_INSERT_TEMPLATE = """
|
|
INSERT INTO torrents (
|
|
media_id,
|
|
info_hash,
|
|
file_index,
|
|
season,
|
|
episode,
|
|
title,
|
|
seeders,
|
|
size,
|
|
tracker,
|
|
sources,
|
|
parsed,
|
|
timestamp
|
|
) VALUES (
|
|
:media_id,
|
|
:info_hash,
|
|
:file_index,
|
|
:season,
|
|
:episode,
|
|
:title,
|
|
:seeders,
|
|
:size,
|
|
:tracker,
|
|
:sources,
|
|
:parsed,
|
|
:timestamp
|
|
)
|
|
"""
|
|
|
|
SQLITE_UPSERT_QUERY = TORRENT_INSERT_TEMPLATE.replace("INSERT", "INSERT OR REPLACE", 1)
|
|
POSTGRES_UPDATE_SET = """
|
|
DO UPDATE SET
|
|
media_id = EXCLUDED.media_id,
|
|
file_index = EXCLUDED.file_index,
|
|
title = EXCLUDED.title,
|
|
seeders = EXCLUDED.seeders,
|
|
size = EXCLUDED.size,
|
|
tracker = EXCLUDED.tracker,
|
|
sources = EXCLUDED.sources,
|
|
parsed = EXCLUDED.parsed,
|
|
timestamp = EXCLUDED.timestamp
|
|
WHERE
|
|
(
|
|
torrents.media_id IS DISTINCT FROM EXCLUDED.media_id OR
|
|
torrents.file_index IS DISTINCT FROM EXCLUDED.file_index OR
|
|
torrents.title IS DISTINCT FROM EXCLUDED.title OR
|
|
torrents.seeders IS DISTINCT FROM EXCLUDED.seeders OR
|
|
torrents.size IS DISTINCT FROM EXCLUDED.size OR
|
|
torrents.tracker IS DISTINCT FROM EXCLUDED.tracker OR
|
|
torrents.sources IS DISTINCT FROM EXCLUDED.sources OR
|
|
torrents.parsed IS DISTINCT FROM EXCLUDED.parsed
|
|
)
|
|
OR
|
|
(
|
|
COALESCE(torrents.timestamp, 0) < (EXCLUDED.timestamp - :update_interval)
|
|
)
|
|
"""
|
|
|
|
POSTGRES_CONFLICT_TARGETS = {
|
|
"series": "(media_id, info_hash, season, episode) WHERE season IS NOT NULL AND episode IS NOT NULL",
|
|
"season_only": "(media_id, info_hash, season) WHERE season IS NOT NULL AND episode IS NULL",
|
|
"episode_only": "(media_id, info_hash, episode) WHERE season IS NULL AND episode IS NOT NULL",
|
|
"none": "(media_id, info_hash) WHERE season IS NULL AND episode IS NULL",
|
|
}
|
|
|
|
_POSTGRES_UPSERT_CACHE: dict[str, str] = {}
|
|
|
|
|
|
def _determine_conflict_key(season, episode):
|
|
if season is not None:
|
|
return "series" if episode is not None else "season_only"
|
|
return "episode_only" if episode is not None else "none"
|
|
|
|
|
|
def _compute_advisory_lock_key(media_id, info_hash, season, episode):
|
|
payload = f"{media_id}|{info_hash}|{season}|{episode}"
|
|
return xxhash.xxh64_intdigest(payload, seed=0) - (1 << 63)
|
|
|
|
|
|
_SQLITE_CHECK_COLS = frozenset(
|
|
["file_index", "title", "seeders", "size", "tracker", "sources", "parsed"]
|
|
)
|
|
|
|
|
|
async def _execute_sqlite_batched_upsert(rows: list[dict]):
|
|
if not rows:
|
|
return
|
|
|
|
info_hashes = {row["info_hash"] for row in rows}
|
|
|
|
if not info_hashes:
|
|
await _execute_standard_sqlite_insert(rows)
|
|
return
|
|
|
|
info_hashes_list = list(info_hashes)
|
|
chunk_size = 900
|
|
existing_rows = []
|
|
|
|
for i in range(0, len(info_hashes_list), chunk_size):
|
|
chunk = info_hashes_list[i : i + chunk_size]
|
|
placeholders = ",".join(f":ih{j}" for j in range(len(chunk)))
|
|
params = {f"ih{j}": ih for j, ih in enumerate(chunk)}
|
|
|
|
chunk_rows = await database.fetch_all(
|
|
f"SELECT media_id, info_hash, season, episode, file_index, title, seeders, size, tracker, sources, parsed, timestamp FROM torrents WHERE info_hash IN ({placeholders})",
|
|
params,
|
|
)
|
|
existing_rows.extend(chunk_rows)
|
|
|
|
existing_map = {
|
|
(row["media_id"], row["info_hash"], row["season"], row["episode"]): row
|
|
for row in existing_rows
|
|
}
|
|
|
|
if not existing_map:
|
|
await _execute_standard_sqlite_insert(rows)
|
|
return
|
|
|
|
to_insert = []
|
|
for row in rows:
|
|
key = (row["media_id"], row["info_hash"], row["season"], row["episode"])
|
|
existing = existing_map.get(key)
|
|
|
|
if not existing:
|
|
to_insert.append(row)
|
|
continue
|
|
|
|
if any(row.get(col) != existing[col] for col in _SQLITE_CHECK_COLS):
|
|
to_insert.append(row)
|
|
continue
|
|
|
|
if existing["timestamp"] < (row["timestamp"] - UPDATE_INTERVAL):
|
|
to_insert.append(row)
|
|
|
|
if to_insert:
|
|
await _execute_standard_sqlite_insert(to_insert)
|
|
|
|
|
|
async def _execute_standard_sqlite_insert(rows: list[dict]):
|
|
keys_to_ignore = {"lock_key", "update_interval"}
|
|
columns = [k for k in rows[0].keys() if k not in keys_to_ignore]
|
|
sanitized_rows = [{k: row[k] for k in columns} for row in rows]
|
|
|
|
query = f"""
|
|
INSERT OR REPLACE INTO torrents ({", ".join(columns)})
|
|
VALUES ({", ".join(f":{col}" for col in columns)})
|
|
"""
|
|
|
|
for attempt in range(5):
|
|
try:
|
|
async with database.transaction():
|
|
await database.execute_many(query, sanitized_rows)
|
|
return
|
|
except Exception:
|
|
if attempt < 4:
|
|
await asyncio.sleep(0.2 * (attempt + 1))
|
|
continue
|
|
raise
|
|
|
|
|
|
async def _execute_batched_upsert(query: str, rows):
|
|
if not rows:
|
|
return
|
|
|
|
if settings.DATABASE_TYPE == "sqlite":
|
|
await _execute_sqlite_batched_upsert(rows)
|
|
return
|
|
|
|
ordered_rows = sorted(rows, key=lambda row: row.get("lock_key") or 0)
|
|
rows_to_insert = []
|
|
|
|
try:
|
|
async with database.transaction():
|
|
for row in ordered_rows:
|
|
lock_key = row.get("lock_key")
|
|
if lock_key is None:
|
|
rows_to_insert.append(row)
|
|
continue
|
|
|
|
acquired = await database.fetch_val(
|
|
"SELECT pg_try_advisory_xact_lock(CAST(:lock_key AS BIGINT))",
|
|
{"lock_key": lock_key},
|
|
)
|
|
if acquired:
|
|
rows_to_insert.append(row)
|
|
|
|
if rows_to_insert:
|
|
sanitized_rows = [
|
|
{key: value for key, value in row.items() if key != "lock_key"}
|
|
for row in rows_to_insert
|
|
]
|
|
|
|
await database.execute_many(query, sanitized_rows)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error executing batched upsert: {e}")
|
|
|
|
|
|
def _get_torrent_upsert_query(conflict_key: str):
|
|
if settings.DATABASE_TYPE == "sqlite":
|
|
return SQLITE_UPSERT_QUERY
|
|
|
|
target = POSTGRES_CONFLICT_TARGETS[conflict_key]
|
|
if conflict_key not in _POSTGRES_UPSERT_CACHE:
|
|
_POSTGRES_UPSERT_CACHE[conflict_key] = (
|
|
TORRENT_INSERT_TEMPLATE + f" ON CONFLICT {target} " + POSTGRES_UPDATE_SET
|
|
)
|
|
return _POSTGRES_UPSERT_CACHE[conflict_key]
|
|
|
|
|
|
async def _upsert_torrent_record(params: dict):
|
|
if settings.DATABASE_TYPE == "sqlite":
|
|
await _execute_sqlite_batched_upsert([params])
|
|
return
|
|
|
|
query = _get_torrent_upsert_query(
|
|
_determine_conflict_key(params.get("season"), params.get("episode"))
|
|
)
|
|
|
|
params["update_interval"] = UPDATE_INTERVAL
|
|
|
|
await database.execute(query, params)
|
|
|
|
|
|
torrent_update_queue = TorrentUpdateQueue()
|