fixes, orval, downloader functioning again
This commit is contained in:
245
backend/thirdparty/downloader/consumers.py
vendored
Normal file
245
backend/thirdparty/downloader/consumers.py
vendored
Normal file
@@ -0,0 +1,245 @@
|
||||
import asyncio
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import zipfile
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import yt_dlp
|
||||
from channels.generic.websocket import AsyncWebsocketConsumer
|
||||
from django.conf import settings
|
||||
from django.core import signing
|
||||
from django.utils.text import slugify
|
||||
|
||||
YDL_BASE = {
|
||||
"quiet": True,
|
||||
"no_check_certificates": True,
|
||||
"js_runtimes": {"node": {}},
|
||||
"remote_components": {"ejs:github"},
|
||||
}
|
||||
|
||||
TOKEN_TTL = 600 # seconds until signed download token expires
|
||||
|
||||
|
||||
class DownloaderConsumer(AsyncWebsocketConsumer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.tmpdir = None
|
||||
|
||||
async def connect(self):
|
||||
await self.accept()
|
||||
|
||||
async def disconnect(self, code):
|
||||
if self.tmpdir and os.path.exists(self.tmpdir):
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
self.tmpdir = None
|
||||
|
||||
async def receive(self, text_data):
|
||||
try:
|
||||
params = json.loads(text_data)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
await self._send({"type": "error", "message": "Invalid JSON"})
|
||||
return
|
||||
|
||||
url = params.get("url")
|
||||
if not url:
|
||||
await self._send({"type": "error", "message": "URL is required"})
|
||||
return
|
||||
|
||||
ext = params.get("ext", "mp4")
|
||||
if not ext or not isinstance(ext, str):
|
||||
await self._send({"type": "error", "message": "Invalid extension"})
|
||||
return
|
||||
|
||||
os.makedirs(settings.DOWNLOADER_TMP_DIR, exist_ok=True)
|
||||
self.tmpdir = tempfile.mkdtemp(prefix="dl_ws_", dir=settings.DOWNLOADER_TMP_DIR)
|
||||
|
||||
await self._send({"type": "status", "message": "Starting…"})
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def progress_hook(d):
|
||||
status = d.get("status")
|
||||
if status == "downloading":
|
||||
percent = (d.get("_percent_str") or "").strip()
|
||||
speed = (d.get("_speed_str") or "").strip()
|
||||
eta = (d.get("_eta_str") or "").strip()
|
||||
downloaded = d.get("downloaded_bytes", 0)
|
||||
total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
|
||||
pct_num = round(downloaded / total * 100, 1) if total else None
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
self._send({
|
||||
"type": "progress",
|
||||
"message": f"Downloading: {percent} at {speed} ETA {eta}",
|
||||
"percent": pct_num,
|
||||
"speed": speed,
|
||||
"eta": eta,
|
||||
}),
|
||||
loop,
|
||||
)
|
||||
elif status == "finished":
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
self._send({"type": "status", "message": "Post-processing…"}),
|
||||
loop,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await asyncio.to_thread(
|
||||
self._run_download,
|
||||
url,
|
||||
ext,
|
||||
params.get("video_quality"),
|
||||
params.get("audio_quality"),
|
||||
params.get("selected_videos"),
|
||||
params.get("subtitles"),
|
||||
bool(params.get("embed_subtitles", False)),
|
||||
bool(params.get("embed_thumbnail", False)),
|
||||
bool(params.get("extract_audio", False)),
|
||||
params.get("cookies"),
|
||||
progress_hook,
|
||||
)
|
||||
except Exception as exc:
|
||||
if self.tmpdir:
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
self.tmpdir = None
|
||||
await self._send({"type": "error", "message": str(exc)})
|
||||
return
|
||||
|
||||
# Sign a token containing the file info — no DB needed.
|
||||
token = signing.dumps(
|
||||
{
|
||||
"file_path": result["file_path"],
|
||||
"filename": result["filename"],
|
||||
"content_type": result["content_type"],
|
||||
"file_size": result["file_size"],
|
||||
"tmpdir": self.tmpdir,
|
||||
},
|
||||
salt="downloader-file-token",
|
||||
)
|
||||
# Consumer no longer owns cleanup — the HTTP serve view will clean up.
|
||||
self.tmpdir = None
|
||||
|
||||
await self._send({
|
||||
"type": "done",
|
||||
"token": token,
|
||||
"filename": result["filename"],
|
||||
"file_size": result["file_size"],
|
||||
})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def _run_download(self, url, ext, video_quality, audio_quality,
|
||||
selected_videos, subtitles, embed_subtitles,
|
||||
embed_thumbnail, extract_audio, cookies, progress_hook):
|
||||
"""Runs synchronously inside a thread pool worker."""
|
||||
tmpdir = self.tmpdir
|
||||
|
||||
# Format selector
|
||||
if video_quality and audio_quality:
|
||||
fmt = f"bv[height<={video_quality}]+ba[abr<={audio_quality}]/b"
|
||||
elif video_quality:
|
||||
fmt = f"bv[height<={video_quality}]+ba/b"
|
||||
elif audio_quality:
|
||||
fmt = f"bv+ba[abr<={audio_quality}]/b"
|
||||
else:
|
||||
fmt = "b/bv+ba"
|
||||
|
||||
ydl_opts = {
|
||||
**YDL_BASE,
|
||||
"format": fmt,
|
||||
"merge_output_format": ext,
|
||||
"max_filesize": settings.DOWNLOADER_MAX_SIZE_BYTES,
|
||||
"postprocessors": [],
|
||||
"progress_hooks": [progress_hook],
|
||||
}
|
||||
|
||||
if cookies:
|
||||
cookie_file = os.path.join(tmpdir, "cookies.txt")
|
||||
with open(cookie_file, "w") as f:
|
||||
f.write(cookies)
|
||||
ydl_opts["cookiefile"] = cookie_file
|
||||
|
||||
if subtitles:
|
||||
if subtitles.lower() == "all":
|
||||
ydl_opts.update({
|
||||
"writesubtitles": True,
|
||||
"writeautomaticsub": True,
|
||||
"subtitleslangs": ["all"],
|
||||
})
|
||||
else:
|
||||
ydl_opts.update({
|
||||
"writesubtitles": True,
|
||||
"subtitleslangs": [l.strip() for l in subtitles.split(",")],
|
||||
})
|
||||
|
||||
if embed_subtitles and subtitles and ext in ("mkv", "mp4"):
|
||||
ydl_opts["postprocessors"].append({"key": "FFmpegEmbedSubtitle"})
|
||||
|
||||
if embed_thumbnail:
|
||||
ydl_opts["writethumbnail"] = True
|
||||
ydl_opts["postprocessors"].append({"key": "EmbedThumbnail"})
|
||||
|
||||
if extract_audio:
|
||||
ydl_opts["postprocessors"].append({
|
||||
"key": "FFmpegExtractAudio",
|
||||
"preferredcodec": ext if ext in ("mp3", "m4a", "opus", "vorbis", "wav") else "mp3",
|
||||
})
|
||||
else:
|
||||
ydl_opts["postprocessors"].append({"key": "FFmpegVideoRemuxer", "preferedformat": ext})
|
||||
|
||||
# Probe to detect playlist
|
||||
probe_opts = {**YDL_BASE, "extract_flat": False}
|
||||
with yt_dlp.YoutubeDL(probe_opts) as ydl:
|
||||
info = ydl.extract_info(url, download=False)
|
||||
|
||||
is_playlist = "entries" in info and info.get("entries") is not None
|
||||
|
||||
if is_playlist:
|
||||
if selected_videos:
|
||||
ydl_opts["playlist_items"] = ",".join(str(n) for n in selected_videos)
|
||||
ydl_opts["outtmpl"] = os.path.join(tmpdir, "%(playlist_index)02d - %(title)s.%(ext)s")
|
||||
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
||||
ydl.download([url])
|
||||
|
||||
zip_path = os.path.join(tmpdir, "playlist.zip")
|
||||
files = [
|
||||
(fn, os.path.join(tmpdir, fn))
|
||||
for fn in os.listdir(tmpdir)
|
||||
if fn not in ("playlist.zip", "cookies.txt")
|
||||
and not fn.startswith(".")
|
||||
and os.path.isfile(os.path.join(tmpdir, fn))
|
||||
]
|
||||
if not files:
|
||||
raise RuntimeError("No files were downloaded from the playlist")
|
||||
|
||||
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for fn, fp in files:
|
||||
zf.write(fp, fn)
|
||||
|
||||
playlist_title = slugify(info.get("title", "playlist"))
|
||||
return {
|
||||
"file_path": zip_path,
|
||||
"filename": f"{playlist_title}.zip",
|
||||
"content_type": "application/zip",
|
||||
"file_size": os.path.getsize(zip_path),
|
||||
}
|
||||
else:
|
||||
ydl_opts["outtmpl"] = os.path.join(tmpdir, "download.%(ext)s")
|
||||
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
||||
info = ydl.extract_info(url, download=True)
|
||||
base = ydl.prepare_filename(info)
|
||||
file_path = base if base.endswith(f".{ext}") else os.path.splitext(base)[0] + f".{ext}"
|
||||
|
||||
safe_title = slugify(info.get("title") or "video")
|
||||
filename = f"{safe_title}.{ext}"
|
||||
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||||
return {
|
||||
"file_path": file_path,
|
||||
"filename": filename,
|
||||
"content_type": content_type,
|
||||
"file_size": os.path.getsize(file_path) if os.path.exists(file_path) else 0,
|
||||
}
|
||||
|
||||
async def _send(self, data: dict):
|
||||
await self.send(text_data=json.dumps(data))
|
||||
Reference in New Issue
Block a user