diff --git a/.env-sample b/.env-sample index af7c718..a760033 100644 --- a/.env-sample +++ b/.env-sample @@ -15,5 +15,4 @@ GET_TORRENT_TIMEOUT=5 # maximum time to obtain the torrent info hash in seconds ZILEAN_URL=None # for DMM search - https://github.com/iPromKnight/zilean CUSTOM_HEADER_HTML=None # only set it if you know what it is PROXY_DEBRID_STREAM=False # Proxy Debrid Streams (very useful to use your debrid service on multiple IPs at same time) -PROXY_DEBRID_STREAM_PASSWORD=CHANGE_ME # Secret password to enter on configuration page to prevent people from abusing your debrid stream proxy -PROXY_DEBRID_STREAM_BYTES_PER_CHUNK=102400 # 10MB per chunks \ No newline at end of file +PROXY_DEBRID_STREAM_PASSWORD=CHANGE_ME # Secret password to enter on configuration page to prevent people from abusing your debrid stream proxy \ No newline at end of file diff --git a/comet/api/stream.py b/comet/api/stream.py index 28d1066..43e97ab 100644 --- a/comet/api/stream.py +++ b/comet/api/stream.py @@ -3,9 +3,11 @@ import json import time import aiohttp +import httpx from fastapi import APIRouter, Request from fastapi.responses import RedirectResponse, StreamingResponse +from starlette.background import BackgroundTask from RTN import Torrent, parse, sort_torrents, title_match from comet.debrid.manager import getDebrid @@ -310,43 +312,41 @@ async def playback(b64config: str, hash: str, index: str): return RedirectResponse(download_link, status_code=302) - @streams.get("/{b64config}/playback/{hash}/{index}") async def playback(request: Request, b64config: str, hash: str, index: str): config = config_check(b64config) if not config: return - + async with aiohttp.ClientSession() as session: debrid = getDebrid(session, config) download_link = await debrid.generate_download_link(hash, index) - if download_link is None: return - + proxy = ( debrid.proxy if config["debridService"] == "alldebrid" else None ) # proxy is not needed to proxy realdebrid stream - + if ( settings.PROXY_DEBRID_STREAM and settings.PROXY_DEBRID_STREAM_PASSWORD == config["debridStreamProxyPassword"] ): - - async def stream_content(headers: dict): - async with aiohttp.ClientSession() as session: - async with session.get( - download_link, headers=headers, proxy=proxy - ) as response: - while True: - chunk = await response.content.read( - settings.PROXY_DEBRID_STREAM_BYTES_PER_CHUNK - ) - if not chunk: - break - yield chunk - + class Streamer: + def __init__(self): + self.response = None + + async def stream_content(self, headers: dict): + async with httpx.AsyncClient(proxy=proxy) as client: + async with client.stream("GET", download_link, headers=headers) as self.response: + async for chunk in self.response.aiter_raw(): + yield chunk + + async def close(self): + if self.response is not None: + await self.response.aclose() + range = None range_header = request.headers.get("range") if range_header: @@ -357,21 +357,20 @@ async def stream_content(headers: dict): range = f"bytes={start}-{end}" async with await session.get( - download_link, headers={"Range": f"bytes={start}-{end}"}, proxy=proxy + download_link, headers={"Range": range}, proxy=proxy ) as response: - await session.close() - if response.status == 206: + streamer = Streamer() + return StreamingResponse( - stream_content({"Range": range}), + streamer.stream_content({"Range": range}), status_code=206, headers={ "Content-Range": response.headers["Content-Range"], "Content-Length": response.headers["Content-Length"], "Accept-Ranges": "bytes", - }, + }, background=BackgroundTask(await streamer.close()) ) - return - - return RedirectResponse(download_link, status_code=302) + + return RedirectResponse(download_link, status_code=302) \ No newline at end of file diff --git a/comet/utils/models.py b/comet/utils/models.py index cdb1e1b..7122ae6 100644 --- a/comet/utils/models.py +++ b/comet/utils/models.py @@ -28,7 +28,6 @@ class AppSettings(BaseSettings): CUSTOM_HEADER_HTML: Optional[str] = None PROXY_DEBRID_STREAM: Optional[bool] = False PROXY_DEBRID_STREAM_PASSWORD: Optional[str] = "CHANGE_ME" - PROXY_DEBRID_STREAM_BYTES_PER_CHUNK: Optional[int] = 102400 settings = AppSettings() diff --git a/poetry.lock b/poetry.lock index 57ee1f5..95dd7be 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2279,4 +2279,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "24f9bb0c3506c9138b89c473f41e1633b6c41b05274bc4d3dc67c4688a4a1b82" +content-hash = "1310bf542d25a317cf35b3373902cce46c9a743a2883035e4bcb78e37a09d1e4" diff --git a/pyproject.toml b/pyproject.toml index 535a36b..45a90ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ databases = "*" aiosqlite = "*" pydantic-settings = "*" bencode-py = "*" +httpx = "*" [tool.poetry.group.dev.dependencies] isort = "*"