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))