144 lines
3.8 KiB
Python
144 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from typing import List
|
|
|
|
|
|
def emit(status, progress, message="", output=""):
|
|
payload = {
|
|
"status": status,
|
|
"progress": progress,
|
|
"message": message,
|
|
}
|
|
if output:
|
|
payload["output"] = output
|
|
print(json.dumps(payload), flush=True)
|
|
|
|
|
|
def run(cmd):
|
|
proc = subprocess.run(cmd, capture_output=True, text=True)
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(proc.stderr.strip() or proc.stdout.strip() or "command failed")
|
|
return proc
|
|
|
|
|
|
def parse_duration(value):
|
|
if value is None:
|
|
return "00:00:10"
|
|
total = int(float(value))
|
|
hours = total // 3600
|
|
minutes = (total % 3600) // 60
|
|
seconds = total % 60
|
|
return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
|
|
|
|
|
|
def format_label(height, ext):
|
|
if height:
|
|
return f"{height}p ({ext})"
|
|
return f"Best ({ext})"
|
|
|
|
|
|
def build_quality_options(formats: List[dict]):
|
|
heights = []
|
|
for item in formats:
|
|
height = item.get("height")
|
|
if isinstance(height, int) and height > 0:
|
|
heights.append(height)
|
|
|
|
unique_heights = sorted(set(heights))
|
|
options = [{"value": "best", "label": "Best available"}]
|
|
for height in unique_heights:
|
|
options.append(
|
|
{
|
|
"value": f"bestvideo[height<={height}]+bestaudio/best[height<={height}]",
|
|
"label": f"Up to {height}p",
|
|
}
|
|
)
|
|
return options
|
|
|
|
|
|
def probe(url):
|
|
cmd = ["yt-dlp", "--dump-single-json", "--no-playlist", url]
|
|
proc = run(cmd)
|
|
payload = json.loads(proc.stdout)
|
|
thumbnail = payload.get("thumbnail") or ""
|
|
duration = payload.get("duration")
|
|
formats = payload.get("formats") or []
|
|
preview = {
|
|
"title": payload.get("title") or "Untitled",
|
|
"thumbnail": thumbnail,
|
|
"durationSeconds": duration or 0,
|
|
"duration": parse_duration(duration),
|
|
"startDefault": "00:00:00",
|
|
"endDefault": parse_duration(duration),
|
|
"qualities": build_quality_options(formats),
|
|
}
|
|
print(json.dumps(preview), flush=True)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--mode", choices=["probe", "download"], default="download")
|
|
parser.add_argument("--url", required=True)
|
|
parser.add_argument("--start", default="00:00:00")
|
|
parser.add_argument("--end", default="00:00:10")
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--quality", default="best")
|
|
args = parser.parse_args()
|
|
|
|
if args.mode == "probe":
|
|
probe(args.url)
|
|
return
|
|
|
|
os.makedirs(os.path.dirname(args.output), exist_ok=True)
|
|
emit("starting", 5, "Resolving media stream")
|
|
|
|
with tempfile.TemporaryDirectory(prefix="aihub-") as tmpdir:
|
|
source_path = os.path.join(tmpdir, "source.%(ext)s")
|
|
download_cmd = [
|
|
"yt-dlp",
|
|
"--no-playlist",
|
|
"-f",
|
|
args.quality,
|
|
"-o",
|
|
source_path,
|
|
args.url,
|
|
]
|
|
run(download_cmd)
|
|
emit("downloaded", 55, "Source downloaded")
|
|
|
|
files = [os.path.join(tmpdir, name) for name in os.listdir(tmpdir)]
|
|
if not files:
|
|
raise RuntimeError("yt-dlp did not produce an output file")
|
|
source_file = sorted(files)[0]
|
|
|
|
ffmpeg_cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-ss",
|
|
args.start,
|
|
"-to",
|
|
args.end,
|
|
"-i",
|
|
source_file,
|
|
"-c",
|
|
"copy",
|
|
args.output,
|
|
]
|
|
emit("cropping", 75, "Cropping requested segment")
|
|
run(ffmpeg_cmd)
|
|
|
|
emit("completed", 100, "Download complete", args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as exc:
|
|
emit("error", 100, str(exc))
|
|
sys.exit(1)
|