246 lines
9.0 KiB
Python
246 lines
9.0 KiB
Python
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))
|