yt-downloader/yt_downloader/cli.py

"""Command-line interface for yt-downloader."""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

from . import __version__, vpn
from .config import Config, QUALITY_PRESETS, default_config_path, load_config, write_default_config
from .downloader import cookie_opts, download_video
from .progress import DownloadProgress


def log_setup(verbose: bool) -> None:
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(levelname)-7s %(message)s",
        stream=sys.stderr,
    )


# -- commands ------------------------------------------------------------

def cmd_init(args: argparse.Namespace) -> int:
    path = Path(args.config)
    if path.exists():
        logging.error("Config already exists: %s", path)
        return 1
    write_default_config(path)
    print(f"Wrote config to {path}")
    print("Edit it, then run: yt-downloader download <url>")
    return 0


def _download_argv(args: argparse.Namespace) -> list[str]:
    """Rebuild the `download` argv for the VPN-wrapped subprocess."""
    argv = []
    if getattr(args, "config", None):
        # --config is a parent/global option and must come before the
        # subcommand, or argparse treats it as unrecognized
        argv += ["--config", str(args.config)]
    argv += ["download", args.url]
    if getattr(args, "quality", None):
        argv += ["--quality", args.quality]
    if getattr(args, "format", None):
        argv += ["--format", args.format]
    if getattr(args, "output", None):
        argv += ["--output", str(args.output)]
    return argv


def cmd_download(args: argparse.Namespace) -> int:
    cfg = load_config(args.config)
    if (
        not args.dry_run
        and vpn.configured(cfg)
        and not os.environ.get(vpn.INNER_ENV)
    ):
        try:
            return vpn.run_download_in_ns(cfg, _download_argv(args))
        except RuntimeError as exc:
            logging.error("%s", exc)
            return 1
    if args.dry_run:
        print(f"[dry] would download {args.url}")
        return 0

    quality = args.quality or cfg.quality
    format_override = args.format or cfg.format_override
    out_dir = Path(args.output) if args.output else cfg.out_dir
    extra = cookie_opts(cfg)

    print(f"downloading: {args.url} (quality={quality}, out={out_dir.expanduser()})")
    bar = DownloadProgress()
    path = download_video(args.url, out_dir, quality, format_override, extra, progress=bar)
    if not path:
        logging.error("download failed: %s", args.url)
        return 1
    print(f"saved: {path}")
    return 0


def cmd_vpn(args: argparse.Namespace) -> int:
    cfg = load_config(args.config)
    try:
        if args.vpn_command == "up":
            if vpn.is_up(cfg):
                print("tunnel already up")
            else:
                vpn.up(cfg)
                print("tunnel up")
            if not vpn.wait_for_handshake(cfg):
                logging.error("no handshake — tunnel state:\n%s", vpn._diagnose(cfg))
                vpn.down(cfg)
                logging.error(
                    "tunnel is up but no WireGuard handshake was established — "
                    "if 'transfer' shows sent but 0 received, the endpoint "
                    "146.70.117.130:51820 is unreachable (try a different "
                    "Mullvad server); if 0 sent, the namespace has no route out. "
                    "Downloads will be refused until the tunnel is active."
                )
                return 1
            print("handshake established")
        elif args.vpn_command == "down":
            vpn.down(cfg)
            print("tunnel down")
        elif args.vpn_command == "status":
            print(vpn.status_text(cfg))
    except RuntimeError as exc:
        logging.error("%s", exc)
        return 1
    return 0


# -- entry point ---------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="yt-downloader",
        description="Download a single YouTube video through an isolated "
                    "Mullvad WireGuard tunnel (network namespace).",
    )
    parser.add_argument("--version", action="version", version=__version__)
    parser.add_argument(
        "--config", "-c", type=str, default=str(default_config_path()),
        help="path to config TOML",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")

    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("init", help="write a starter config file")

    p_dl = sub.add_parser("download", help="download a single YouTube video")
    p_dl.add_argument("url", help="YouTube video URL")
    p_dl.add_argument(
        "-q", "--quality", choices=QUALITY_PRESETS, default=None,
        help="video quality preset (default from config)",
    )
    p_dl.add_argument(
        "-f", "--format", default=None,
        help="raw yt-dlp format string, overrides --quality",
    )
    p_dl.add_argument(
        "-o", "--output", type=Path, default=None,
        help="output directory (default from config)",
    )
    p_dl.add_argument(
        "--dry-run", action="store_true",
        help="plan only: don't set up the tunnel and don't download",
    )

    p_vpn = sub.add_parser("vpn", help="manage the isolated Mullvad tunnel used by downloads")
    vsub = p_vpn.add_subparsers(dest="vpn_command", required=True)
    vsub.add_parser("up", help="create the tunnel namespace and bring WireGuard up")
    vsub.add_parser("down", help="tear the tunnel namespace down")
    vsub.add_parser("status", help="show tunnel state")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    log_setup(args.verbose)
    if args.command == "init":
        args.config = args.config if args.config else str(default_config_path())
        return cmd_init(args)
    try:
        if args.command == "download":
            return cmd_download(args)
        if args.command == "vpn":
            return cmd_vpn(args)
    except FileNotFoundError as exc:
        logging.error("%s", exc)
        return 1
    parser.error(f"unknown command {args.command!r}")
    return 2


if __name__ == "__main__":
    sys.exit(main())