WeBrake/app.py

#!/usr/bin/env python3
"""
WeBrake — self-hosted HandBrake web UI for Alpine LXC.
Repo: https://gogs.av2x.dev/av2x/WeBrake

Serves a single-page UI, accepts media uploads, runs HandBrakeCLI jobs with
live JSON progress, exposes every HandBrake flag (structured groups + raw
passthrough), detects GPU/hardware encoders, and serves finished files back
for download.

Secrets live in token.json (generated by install.sh, never committed).
"""

import json
import os
import re
import secrets
import shlex
import shutil
import subprocess
import threading
import time
import uuid
from functools import wraps
from pathlib import Path

from flask import (Flask, Response, abort, jsonify, request,
                   send_file, send_from_directory)

# --------------------------------------------------------------------------
# Paths & config
# --------------------------------------------------------------------------
APP_DIR    = Path(__file__).resolve().parent
STATIC_DIR = APP_DIR / "static"
DATA_DIR   = Path(os.environ.get("WEBRAKE_DATA", "/var/lib/webrake"))
UPLOAD_DIR = DATA_DIR / "uploads"
OUTPUT_DIR = DATA_DIR / "output"
JOBS_DIR   = DATA_DIR / "jobs"
TOKEN_FILE = Path(os.environ.get("WEBRAKE_TOKENS", str(APP_DIR / "token.json")))

for d in (UPLOAD_DIR, OUTPUT_DIR, JOBS_DIR):
    d.mkdir(parents=True, exist_ok=True)

HANDBRAKE = shutil.which("HandBrakeCLI") or "/usr/bin/HandBrakeCLI"
VERSION   = "1.0.1"

def load_tokens():
    """token.json is created by install.sh and is never part of the repo."""
    if TOKEN_FILE.exists():
        try:
            return json.loads(TOKEN_FILE.read_text())
        except Exception:
            pass
    # First run without installer: generate one locally so the app still boots.
    tok = {
        "api_token": secrets.token_urlsafe(32),
        "secret_key": secrets.token_urlsafe(32),
        "require_auth": False,
        "bots": {},
    }
    try:
        TOKEN_FILE.write_text(json.dumps(tok, indent=2))
        os.chmod(TOKEN_FILE, 0o600)
    except Exception:
        pass
    return tok

TOKENS = load_tokens()

app = Flask(__name__, static_folder=None)
app.secret_key = TOKENS.get("secret_key", secrets.token_urlsafe(32))
app.config["MAX_CONTENT_LENGTH"] = 512 * 1024 * 1024 * 1024  # 512 GB ceiling

def auth_required(fn):
    """If require_auth is on in token.json, every API call must carry a token."""
    @wraps(fn)
    def wrapper(*args, **kwargs):
        if TOKENS.get("require_auth"):
            supplied = (request.headers.get("X-API-Token")
                        or request.args.get("token", ""))
            valid = {TOKENS.get("api_token")} | set(
                (TOKENS.get("bots") or {}).values())
            if supplied not in valid or not supplied:
                abort(401, description="Missing or invalid API token")
        return fn(*args, **kwargs)
    return wrapper

# --------------------------------------------------------------------------
# Capabilities / GPU detection
# --------------------------------------------------------------------------
_caps_cache = None

def detect_capabilities():
    """Probe HandBrakeCLI + /dev for hardware encode paths (LXC GPU passthru)."""
    global _caps_cache
    if _caps_cache:
        return _caps_cache

    caps = {
        "handbrake_found": os.path.exists(HANDBRAKE) or shutil.which("HandBrakeCLI") is not None,
        "handbrake_version": None,
        "encoders": [],
        "hw_encoders": [],
        "devices": {"dri": [], "nvidia": []},
        "vaapi": False, "qsv": False, "nvenc": False, "vce": False,
    }

    dri = Path("/dev/dri")
    if dri.exists():
        caps["devices"]["dri"] = sorted(p.name for p in dri.iterdir())
    for n in ("/dev/nvidia0", "/dev/nvidiactl"):
        if os.path.exists(n):
            caps["devices"]["nvidia"].append(os.path.basename(n))

    if caps["handbrake_found"]:
        try:
            out = subprocess.run([HANDBRAKE, "--version"], capture_output=True,
                                 text=True, timeout=20)
            m = re.search(r"HandBrake\s+([\w.\-]+)", out.stdout + out.stderr)
            caps["handbrake_version"] = m.group(1) if m else "unknown"
        except Exception:
            pass
        try:
            out = subprocess.run([HANDBRAKE, "--help"], capture_output=True,
                                 text=True, timeout=30)
            help_txt = out.stdout + out.stderr
            # Encoder ids appear indented in the -e/--encoder section
            enc = set(re.findall(
                r"^\s{6,}((?:x26[45]|mpeg[24]|VP[89]|svt_av1|theora|ffv1|"
                r"qsv_\w+|nvenc_\w+|vce_\w+|vaapi_\w+|mf_\w+)[\w]*)\s*$",
                help_txt, re.M))
            caps["encoders"] = sorted(enc)
            caps["hw_encoders"] = sorted(e for e in enc if re.match(
                r"^(qsv|nvenc|vce|vaapi|mf)_", e))
            caps["qsv"]   = any(e.startswith("qsv_") for e in enc)
            caps["nvenc"] = any(e.startswith("nvenc_") for e in enc)
            caps["vce"]   = any(e.startswith("vce_") for e in enc)
            caps["vaapi"] = any(e.startswith("vaapi_") for e in enc) or (
                bool(caps["devices"]["dri"]))
        except Exception:
            pass

    _caps_cache = caps
    return caps

# --------------------------------------------------------------------------
# Job model
# --------------------------------------------------------------------------
JOBS: dict = {}
JOBS_LOCK = threading.Lock()
QUEUE_EVENT = threading.Event()

def persist_job(job):
    slim = {k: v for k, v in job.items() if k != "proc"}
    try:
        (JOBS_DIR / f"{job['id']}.json").write_text(json.dumps(slim, indent=2))
    except Exception:
        pass

def load_persisted_jobs():
    for f in sorted(JOBS_DIR.glob("*.json")):
        try:
            j = json.loads(f.read_text())
            if j.get("status") in ("queued", "running", "scanning"):
                j["status"] = "failed"
                j["error"] = "Interrupted by server restart"
            JOBS[j["id"]] = j
        except Exception:
            continue

load_persisted_jobs()

def new_job(kind, filename, src_path, options):
    job = {
        "id": uuid.uuid4().hex[:12],
        "kind": kind,                      # "encode"
        "filename": filename,
        "src": str(src_path),
        "out": None,
        "options": options,
        "status": "queued",                # queued|running|done|failed|cancelled
        "progress": 0.0,
        "fps": None, "fps_avg": None, "eta": None, "pass": None,
        "log_tail": [],
        "error": None,
        "created": time.time(),
        "started": None, "finished": None,
        "cmd": None,
        "proc": None,
    }
    with JOBS_LOCK:
        JOBS[job["id"]] = job
    persist_job(job)
    QUEUE_EVENT.set()
    return job

# --------------------------------------------------------------------------
# HandBrake command construction — full flag surface
# --------------------------------------------------------------------------
# Structured options map 1:1 onto HandBrakeCLI flags. Anything not covered
# structurally can be supplied verbatim through options["extra_args"], so the
# complete HandBrake feature set is reachable from the UI.
FLAG_MAP = {
    # General / container
    "preset":              ("--preset", str),
    "preset_import_file":  ("--preset-import-file", str),
    "format":              ("--format", str),
    "optimize":            ("--optimize", bool),
    "align_av":            ("--align-av", bool),
    "inline_parameter_sets": ("--inline-parameter-sets", bool),
    "markers":             ("--markers", bool),
    "no_markers":          ("--no-markers", bool),
    # Source
    "title":               ("--title", str),
    "min_duration":        ("--min-duration", str),
    "main_feature":        ("--main-feature", bool),
    "chapters":            ("--chapters", str),
    "angle":               ("--angle", str),
    "previews":            ("--previews", str),
    "start_at_preview":    ("--start-at-preview", str),
    "start_at":            ("--start-at", str),
    "stop_at":             ("--stop-at", str),
    # Video
    "encoder":             ("--encoder", str),
    "encoder_preset":      ("--encoder-preset", str),
    "encoder_tune":        ("--encoder-tune", str),
    "encoder_profile":     ("--encoder-profile", str),
    "encoder_level":       ("--encoder-level", str),
    "quality":             ("--quality", str),
    "vb":                  ("--vb", str),
    "two_pass":            ("--two-pass", bool),
    "turbo":               ("--turbo", bool),
    "rate":                ("--rate", str),
    "cfr":                 ("--cfr", bool),
    "vfr":                 ("--vfr", bool),
    "pfr":                 ("--pfr", bool),
    "encopts":             ("--encopts", str),
    "enable_hw_decoding":  ("--enable-hw-decoding", str),
    "disable_hw_decoding": ("--disable-hw-decoding", bool),
    # Audio
    "audio_lang_list":     ("--audio-lang-list", str),
    "all_audio":           ("--all-audio", bool),
    "first_audio":         ("--first-audio", bool),
    "audio":               ("--audio", str),
    "aencoder":            ("--aencoder", str),
    "audio_copy_mask":     ("--audio-copy-mask", str),
    "audio_fallback":      ("--audio-fallback", str),
    "ab":                  ("--ab", str),
    "aq":                  ("--aq", str),
    "ac":                  ("--ac", str),
    "mixdown":             ("--mixdown", str),
    "normalize_mix":       ("--normalize-mix", str),
    "arate":               ("--arate", str),
    "drc":                 ("--drc", str),
    "gain":                ("--gain", str),
    "adither":             ("--adither", str),
    "aname":               ("--aname", str),
    # Picture
    "width":               ("--width", str),
    "height":              ("--height", str),
    "crop":                ("--crop", str),
    "crop_mode":           ("--crop-mode", str),
    "maxWidth":            ("--maxWidth", str),
    "maxHeight":           ("--maxHeight", str),
    "non_anamorphic":      ("--non-anamorphic", bool),
    "auto_anamorphic":     ("--auto-anamorphic", bool),
    "loose_anamorphic":    ("--loose-anamorphic", bool),
    "custom_anamorphic":   ("--custom-anamorphic", bool),
    "display_width":       ("--display-width", str),
    "keep_display_aspect": ("--keep-display-aspect", bool),
    "pixel_aspect":        ("--pixel-aspect", str),
    "modulus":             ("--modulus", str),
    "color_matrix":        ("--color-matrix", str),
    # Filters
    "comb_detect":         ("--comb-detect", "optval"),
    "deinterlace":         ("--deinterlace", "optval"),
    "decomb":              ("--decomb", "optval"),
    "detelecine":          ("--detelecine", "optval"),
    "hqdn3d":              ("--hqdn3d", "optval"),
    "nlmeans":             ("--nlmeans", "optval"),
    "nlmeans_tune":        ("--nlmeans-tune", str),
    "chroma_smooth":       ("--chroma-smooth", "optval"),
    "chroma_smooth_tune":  ("--chroma-smooth-tune", str),
    "unsharp":             ("--unsharp", "optval"),
    "unsharp_tune":        ("--unsharp-tune", str),
    "lapsharp":            ("--lapsharp", "optval"),
    "lapsharp_tune":       ("--lapsharp-tune", str),
    "deblock":             ("--deblock", "optval"),
    "deblock_tune":        ("--deblock-tune", str),
    "rotate":              ("--rotate", "optval"),
    "grayscale":           ("--grayscale", bool),
    "pad":                 ("--pad", str),
    "colorspace":          ("--colorspace", str),
    # Subtitles
    "subtitle_lang_list":  ("--subtitle-lang-list", str),
    "all_subtitles":       ("--all-subtitles", bool),
    "first_subtitle":      ("--first-subtitle", bool),
    "subtitle":            ("--subtitle", str),
    "subtitle_forced":     ("--subtitle-forced", "optval"),
    "subtitle_burned":     ("--subtitle-burned", "optval"),
    "subtitle_default":    ("--subtitle-default", "optval"),
    "subname":             ("--subname", str),
    "native_language":     ("--native-language", str),
    "native_dub":          ("--native-dub", bool),
    "srt_file":            ("--srt-file", str),
    "srt_codeset":         ("--srt-codeset", str),
    "srt_offset":          ("--srt-offset", str),
    "srt_lang":            ("--srt-lang", str),
    "srt_default":         ("--srt-default", "optval"),
    "srt_burn":            ("--srt-burn", "optval"),
    "ssa_file":            ("--ssa-file", str),
    "ssa_offset":          ("--ssa-offset", str),
    "ssa_lang":            ("--ssa-lang", str),
    "ssa_default":         ("--ssa-default", "optval"),
    "ssa_burn":            ("--ssa-burn", "optval"),
}

EXT_FOR_FORMAT = {"av_mp4": ".mp4", "av_mkv": ".mkv", "av_webm": ".webm"}

def build_cmd(job):
    o = job["options"] or {}
    src = Path(job["src"])
    fmt = o.get("format") or "av_mkv"
    ext = EXT_FOR_FORMAT.get(fmt, ".mkv")
    default_stem = re.sub(r"^[0-9a-f]{8}_", "", src.stem)  # drop upload prefix
    out_name = (o.get("output_name") or default_stem) + ext
    out_name = re.sub(r"[^\w.\- ()\[\]]", "_", out_name)
    out_path = OUTPUT_DIR / f"{job['id']}_{out_name}"
    job["out"] = str(out_path)

    cmd = [HANDBRAKE, "--json", "-i", str(src), "-o", str(out_path)]

    for key, (flag, typ) in FLAG_MAP.items():
        if key not in o:
            continue
        val = o[key]
        if typ is bool:
            if val:
                cmd.append(flag)
        elif typ == "optval":
            # Filter-style flags: True enables with defaults, a string passes settings
            if val is True or val == "":
                cmd.append(flag)
            elif val:
                cmd.append(f"{flag}={val}")
        else:
            if val not in (None, ""):
                cmd += [flag, str(val)]

    extra = (o.get("extra_args") or "").strip()
    if extra:
        cmd += shlex.split(extra)
    return cmd

# --------------------------------------------------------------------------
# Worker — runs one HandBrakeCLI job at a time, parses --json progress
# --------------------------------------------------------------------------
def parse_progress_blocks(job, stream):
    """HandBrakeCLI --json emits 'Progress: { ... }' blocks; brace-balance them."""
    buf, depth, capturing = [], 0, False
    for raw in iter(stream.readline, ""):
        line = raw.rstrip("\n")
        job["log_tail"].append(line)
        if len(job["log_tail"]) > 400:
            del job["log_tail"][:200]

        if not capturing:
            if line.strip().startswith("Progress:"):
                capturing = True
                brace_part = line.split("Progress:", 1)[1]
                buf = [brace_part]
                depth = brace_part.count("{") - brace_part.count("}")
                if depth == 0 and "{" in brace_part:
                    _apply_progress(job, "".join(buf))
                    capturing = False
            continue

        buf.append(line)
        depth += line.count("{") - line.count("}")
        if depth <= 0 and any("{" in b for b in buf):
            _apply_progress(job, "\n".join(buf))
            capturing, buf, depth = False, [], 0

def _apply_progress(job, text):
    try:
        data = json.loads(text)
    except Exception:
        return
    state = data.get("State")
    if state == "WORKING":
        w = data.get("Working", {})
        job["progress"] = round(float(w.get("Progress", 0)) * 100, 2)
        job["fps"] = round(float(w.get("Rate", 0)), 2)
        job["fps_avg"] = round(float(w.get("RateAvg", 0)), 2)
        job["eta"] = int(w.get("ETASeconds", 0)) or None
        job["pass"] = w.get("Pass")
        job["status"] = "running"
    elif state == "SCANNING":
        s = data.get("Scanning", {})
        job["progress"] = round(float(s.get("Progress", 0)) * 100, 2)
        job["status"] = "scanning"
    elif state == "MUXING":
        job["status"] = "running"
        job["progress"] = max(job["progress"], 99.0)
    elif state == "WORKDONE":
        err = data.get("WorkDone", {}).get("Error", 0)
        if err not in (0, "0", None):
            job["error"] = f"HandBrake error code {err}"

def worker_loop():
    while True:
        QUEUE_EVENT.wait(timeout=5)
        QUEUE_EVENT.clear()
        while True:
            with JOBS_LOCK:
                pending = [j for j in JOBS.values() if j["status"] == "queued"]
                pending.sort(key=lambda j: j["created"])
                job = pending[0] if pending else None
                if job:
                    job["status"] = "scanning"
                    job["started"] = time.time()
            if not job:
                break
            run_job(job)

def run_job(job):
    try:
        cmd = build_cmd(job)
        job["cmd"] = " ".join(shlex.quote(c) for c in cmd)
        persist_job(job)
        env = dict(os.environ)
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT, text=True,
                                bufsize=1, env=env)
        job["proc"] = proc
        parse_progress_blocks(job, proc.stdout)
        rc = proc.wait()
        job["proc"] = None
        if job["status"] == "cancelled":
            pass
        elif rc == 0 and not job["error"] and Path(job["out"]).exists():
            job["status"] = "done"
            job["progress"] = 100.0
        else:
            job["status"] = "failed"
            job["error"] = job["error"] or f"HandBrakeCLI exited with code {rc}"
    except FileNotFoundError:
        job["status"] = "failed"
        job["error"] = "HandBrakeCLI not found — is the handbrake package installed?"
    except Exception as e:
        job["status"] = "failed"
        job["error"] = str(e)
    finally:
        job["finished"] = time.time()
        persist_job(job)

threading.Thread(target=worker_loop, daemon=True).start()

# --------------------------------------------------------------------------
# Routes
# --------------------------------------------------------------------------
@app.route("/")
def index():
    return send_from_directory(STATIC_DIR, "index.html")

@app.route("/api/version")
def api_version():
    # Intentionally unauthenticated: the UI probes this to learn whether a
    # token is required before it can show the token panel.
    caps = detect_capabilities()
    return jsonify({"webrake": VERSION,
                    "handbrake": caps.get("handbrake_version"),
                    "auth_required": bool(TOKENS.get("require_auth"))})

@app.route("/api/capabilities")
@auth_required
def api_capabilities():
    return jsonify(detect_capabilities())

@app.route("/api/presets")
@auth_required
def api_presets():
    try:
        out = subprocess.run([HANDBRAKE, "--preset-list"], capture_output=True,
                             text=True, timeout=30)
        presets, category = [], None
        for line in (out.stdout + out.stderr).splitlines():
            m_cat = re.match(r"^([A-Z][\w /&\-]+)/\s*$", line.strip())
            if m_cat:
                category = m_cat.group(1)
                continue
            m_p = re.match(r"^\s{4}(\S.*\S|\S)\s*$", line)
            if m_p and category and not line.strip().startswith(("+", "-")):
                name = m_p.group(1)
                if not name.endswith(":") and len(line) - len(line.lstrip()) == 4:
                    presets.append({"category": category, "name": name})
        return jsonify({"presets": presets})
    except Exception as e:
        return jsonify({"presets": [], "error": str(e)})

@app.route("/api/upload", methods=["POST"])
@auth_required
def api_upload():
    f = request.files.get("file")
    if not f or not f.filename:
        abort(400, description="No file supplied")
    safe = re.sub(r"[^\w.\- ()\[\]]", "_", os.path.basename(f.filename))
    dest = UPLOAD_DIR / f"{uuid.uuid4().hex[:8]}_{safe}"
    f.save(dest)
    return jsonify({"upload_id": dest.name, "filename": safe,
                    "size": dest.stat().st_size})

@app.route("/api/uploads")
@auth_required
def api_uploads():
    items = []
    for p in sorted(UPLOAD_DIR.iterdir()):
        if p.is_file():
            items.append({"upload_id": p.name,
                          "filename": p.name.split("_", 1)[-1],
                          "size": p.stat().st_size,
                          "mtime": p.stat().st_mtime})
    return jsonify({"uploads": items})

@app.route("/api/uploads/<upload_id>", methods=["DELETE"])
@auth_required
def api_delete_upload(upload_id):
    p = UPLOAD_DIR / os.path.basename(upload_id)
    if p.exists():
        p.unlink()
        return jsonify({"deleted": True})
    abort(404)

@app.route("/api/scan/<upload_id>")
@auth_required
def api_scan(upload_id):
    """Title/track discovery via HandBrakeCLI --scan --json."""
    p = UPLOAD_DIR / os.path.basename(upload_id)
    if not p.exists():
        abort(404, description="Upload not found")
    try:
        out = subprocess.run(
            [HANDBRAKE, "--json", "-i", str(p), "--scan", "--title", "0"],
            capture_output=True, text=True, timeout=300)
        text = out.stdout
        m = re.search(r"JSON Title Set:\s*(\{.*)", text, re.S)
        if not m:
            return jsonify({"error": "Scan produced no title data",
                            "log": text[-2000:]}), 500
        blob, depth, end = m.group(1), 0, 0
        for i, ch in enumerate(blob):
            if ch == "{": depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    end = i + 1
                    break
        return jsonify(json.loads(blob[:end]))
    except subprocess.TimeoutExpired:
        return jsonify({"error": "Scan timed out"}), 504
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/api/jobs", methods=["GET"])
@auth_required
def api_jobs():
    with JOBS_LOCK:
        jobs = [{k: v for k, v in j.items() if k not in ("proc", "log_tail")}
                for j in JOBS.values()]
    jobs.sort(key=lambda j: j["created"], reverse=True)
    return jsonify({"jobs": jobs})

@app.route("/api/jobs", methods=["POST"])
@auth_required
def api_create_job():
    body = request.get_json(force=True, silent=True) or {}
    upload_id = body.get("upload_id")
    if not upload_id:
        abort(400, description="upload_id required")
    src = UPLOAD_DIR / os.path.basename(upload_id)
    if not src.exists():
        abort(404, description="Upload not found")
    options = body.get("options") or {}
    job = new_job("encode", src.name.split("_", 1)[-1], src, options)
    return jsonify({"job_id": job["id"]}), 201

@app.route("/api/jobs/<job_id>")
@auth_required
def api_job(job_id):
    job = JOBS.get(job_id)
    if not job:
        abort(404)
    slim = {k: v for k, v in job.items() if k != "proc"}
    return jsonify(slim)

@app.route("/api/jobs/<job_id>/log")
@auth_required
def api_job_log(job_id):
    job = JOBS.get(job_id)
    if not job:
        abort(404)
    return Response("\n".join(job.get("log_tail", [])), mimetype="text/plain")

@app.route("/api/jobs/<job_id>/cancel", methods=["POST"])
@auth_required
def api_cancel(job_id):
    job = JOBS.get(job_id)
    if not job:
        abort(404)
    if job["status"] in ("queued", "running", "scanning"):
        job["status"] = "cancelled"
        proc = job.get("proc")
        if proc:
            try:
                proc.terminate()
            except Exception:
                pass
        persist_job(job)
    return jsonify({"status": job["status"]})

@app.route("/api/jobs/<job_id>", methods=["DELETE"])
@auth_required
def api_delete_job(job_id):
    job = JOBS.get(job_id)
    if not job:
        abort(404)
    if job["status"] in ("running", "scanning"):
        abort(409, description="Cancel the job before deleting it")
    with JOBS_LOCK:
        JOBS.pop(job_id, None)
    try:
        (JOBS_DIR / f"{job_id}.json").unlink(missing_ok=True)
        if job.get("out") and Path(job["out"]).exists():
            Path(job["out"]).unlink()
    except Exception:
        pass
    return jsonify({"deleted": True})

@app.route("/api/download/<job_id>")
@auth_required
def api_download(job_id):
    job = JOBS.get(job_id)
    if not job or job["status"] != "done" or not job.get("out"):
        abort(404, description="No finished output for this job")
    out = Path(job["out"])
    if not out.exists():
        abort(410, description="Output file no longer exists")
    name = out.name.split("_", 1)[-1]
    return send_file(out, as_attachment=True, download_name=name)

@app.errorhandler(400)
@app.errorhandler(401)
@app.errorhandler(404)
@app.errorhandler(409)
@app.errorhandler(410)
def json_error(err):
    return jsonify({"error": getattr(err, "description", str(err))}), err.code

if __name__ == "__main__":
    host = os.environ.get("WEBRAKE_HOST", "0.0.0.0")
    port = int(os.environ.get("WEBRAKE_PORT", "8090"))
    detect_capabilities()
    app.run(host=host, port=port, threaded=True)