cli.py

#!/usr/bin/env python3
"""
proxdeploy — CLI entrypoint for Proxmox VM deployment.
"""

from __future__ import annotations

import ipaddress
import logging
import os
import re
import sys
import time
from pathlib import Path
from typing import Any, Mapping, NoReturn, Optional

import click
import yaml

from proxmox.api import ProxmoxAuthError, ProxmoxClient, ProxmoxAPIError
from proxmox.cloudinit import CloudInitConfig, apply_cloudinit
from proxmox.config_validate import (
    client_retry_settings,
    client_timeout_tuple,
    validate_config,
)
from proxmox.exceptions import ConfigurationError, ProxDeployError, SSHProvisionError
from proxmox.guest import list_guest_ipv4s
from proxmox.tasks import TaskFailedError, TaskTimeoutError, wait_for_task
from proxmox.template import ValidatedTemplate, validate_template
from proxmox.vm import (
    VMCreateResult,
    VMCreateSpec,
    create_kvm,
    destroy_vm,
    get_vm_status,
    list_vms,
    start_vm,
    stop_vm,
    vm_exists,
)
from provision.ssh import run_remote_bash_script, wait_for_ssh_port

logger = logging.getLogger(__name__)

DEFAULT_CONFIG = Path("config.yaml")
LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"

HOST_LIKE = re.compile(r"^[A-Za-z0-9.\[\]:%_\-]+$")


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

def _no_color_env() -> bool:
    return bool(os.environ.get("NO_COLOR", "").strip())


def _use_color(ctx: Optional[click.Context]) -> bool:
    if _no_color_env():
        return False
    if ctx and isinstance(ctx.obj, dict) and ctx.obj.get("no_color"):
        return False
    return sys.stdout.isatty()


def _say(message: str, *, fg: Optional[str] = None, err: bool = False, bold: bool = False) -> None:
    ctx = click.get_current_context(silent=True)
    use = _use_color(ctx)
    if use and (fg or bold):
        click.secho(message, fg=fg, bold=bold, err=err)
    else:
        click.echo(message, err=err)


def _say_success(message: str) -> None:
    ctx = click.get_current_context(silent=True)
    if _use_color(ctx):
        click.secho(f"  ✓ {message}", fg="green")
    else:
        click.echo(f"  OK: {message}")


def _say_step(message: str) -> None:
    _say(message, fg="cyan", bold=False)


def _say_warn(message: str) -> None:
    _say(f"  Warning: {message}", fg="yellow", err=True)


def _fail(ctx: click.Context, exc: ProxDeployError) -> NoReturn:
    logger.error("%s", exc)
    msg = exc.user_message()
    if _use_color(ctx):
        click.secho(f"Error: {msg}", fg="red", err=True)
    else:
        click.echo(f"Error: {msg}", err=True)
    raise SystemExit(exc.exit_code)


# ---------------------------------------------------------------------------
# YAML / password helpers
# ---------------------------------------------------------------------------

def _load_yaml(path: Path) -> Any:
    try:
        text = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise click.ClickException(f"Cannot read {path}: {exc}") from exc
    try:
        return yaml.safe_load(text)
    except yaml.YAMLError as exc:
        raise click.ClickException(f"Invalid YAML in {path}: {exc}") from exc


def _resolve_password(cfg: Mapping[str, Any]) -> str:
    env_pw = os.environ.get("PROXDEPLOY_PASSWORD")
    if env_pw:
        return env_pw
    px = cfg.get("proxmox") or {}
    pw = px.get("password") or ""
    return str(pw)


def _build_spec(cfg: Mapping[str, Any], tpl: ValidatedTemplate) -> tuple[str, VMCreateSpec]:
    px = cfg.get("proxmox") or {}
    defaults = cfg.get("defaults") or {}
    if not isinstance(defaults, dict):
        raise ConfigurationError("config.defaults must be a mapping when present")

    node = px.get("node")
    if not node:
        raise ConfigurationError("config.proxmox.node is required")

    storage = tpl.storage or defaults.get("storage")
    if not storage:
        raise ConfigurationError(
            "storage must be set in template or defaults.storage in config.yaml"
        )

    bridge = tpl.bridge or str(defaults.get("bridge", "vmbr0"))
    ostype = tpl.ostype or str(defaults.get("ostype", "l26"))

    spec = VMCreateSpec(
        name=tpl.name,
        memory=tpl.memory,
        cores=tpl.cores,
        disk_gb=tpl.disk,
        bridge=bridge,
        storage=str(storage),
        vmid=tpl.vmid,
        ostype=ostype,
    )
    return str(node), spec


def _resolve_ssh_password(validated: ValidatedTemplate, cli_password: Optional[str]) -> Optional[str]:
    if cli_password:
        return cli_password
    env = os.environ.get("PROXDEPLOY_SSH_PASSWORD")
    if env:
        return env
    return validated.ssh_password


# ---------------------------------------------------------------------------
# CLI input validators
# ---------------------------------------------------------------------------

def _validate_ssh_host(_ctx: click.Context, _param: click.Parameter, value: Optional[str]) -> Optional[str]:
    if value is None or (isinstance(value, str) and not value.strip()):
        return None
    h = value.strip()
    if len(h) > 253:
        raise click.BadParameter("SSH host must be at most 253 characters.")
    try:
        ipaddress.ip_address(h)
        return h
    except ValueError:
        pass
    if not HOST_LIKE.fullmatch(h):
        raise click.BadParameter(
            "SSH host must be an IP address or hostname (use letters, digits, . : - _ % [ ] only)."
        )
    return h


def _validate_ssh_user(_ctx: click.Context, _param: click.Parameter, value: str) -> str:
    u = (value or "").strip()
    if not u:
        raise click.BadParameter("SSH user must not be empty.")
    if len(u) > 128:
        raise click.BadParameter("SSH user is too long (max 128 characters).")
    return u


# ---------------------------------------------------------------------------
# Build authenticated client from config
# ---------------------------------------------------------------------------

def _build_client(config: Mapping[str, Any]) -> tuple[ProxmoxClient, str, str]:
    """Return (client, api_host, node) from validated config. Authenticates."""
    proxmox_cfg = config["proxmox"]
    api_host = proxmox_cfg["host"]
    verify_ssl = proxmox_cfg["verify_ssl"]
    username = proxmox_cfg["username"]
    node = proxmox_cfg["node"]

    password = _resolve_password(config)
    if not password:
        raise ConfigurationError(
            "Proxmox password missing: set PROXDEPLOY_PASSWORD or proxmox.password in config.yaml"
        )

    if not verify_ssl:
        try:
            import urllib3
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        except ImportError:
            pass

    timeout_pair = client_timeout_tuple(proxmox_cfg)
    retries, backoff_base, max_wait = client_retry_settings(proxmox_cfg)
    client = ProxmoxClient(
        api_host,
        verify_ssl=verify_ssl,
        timeout=timeout_pair,
        max_retries=retries,
        retry_backoff_base=backoff_base,
        retry_max_wait=max_wait,
    )
    client.authenticate(str(username), password)
    return client, api_host, node


def _load_config(config_path: Path) -> Mapping[str, Any]:
    cfg_raw = _load_yaml(config_path)
    if not isinstance(cfg_raw, dict):
        raise ConfigurationError("config root must be a mapping")
    return validate_config(cfg_raw)


# ---------------------------------------------------------------------------
# SSH provisioning
# ---------------------------------------------------------------------------

def _provision_via_ssh(
    px_client: ProxmoxClient,
    node: str,
    vmid: int,
    *,
    ssh_host: Optional[str],
    ssh_port: int,
    ssh_user: str,
    script_path: Path,
    ssh_key: Optional[Path],
    ssh_password: Optional[str],
    timeout_sec: int,
    trust_host_key: bool,
) -> None:
    deadline = time.monotonic() + timeout_sec
    while time.monotonic() < deadline:
        if ssh_host:
            hosts = [ssh_host]
        else:
            hosts = list_guest_ipv4s(px_client, node, vmid)
            if not hosts:
                logger.info(
                    "Waiting for guest IPv4 from QEMU agent (vmid=%s). "
                    "Use --ssh-host to skip agent discovery.",
                    vmid,
                )

        if hosts:
            for host in hosts:
                port_deadline = min(time.monotonic() + 30.0, deadline)
                if port_deadline <= time.monotonic():
                    break
                try:
                    wait_for_ssh_port(host, ssh_port, deadline=port_deadline, poll=2.0)
                except SSHProvisionError:
                    logger.debug("SSH port not yet open on %s", host)
                    continue
                try:
                    run_remote_bash_script(
                        host, ssh_port, ssh_user, script_path,
                        key_path=ssh_key, password=ssh_password,
                        trust_host_key=trust_host_key,
                    )
                    _say_success(f"Provisioning finished on {host}")
                    return
                except SSHProvisionError as exc:
                    logger.warning("SSH or script failed on %s: %s", host, exc)
                    continue

        time.sleep(5.0)

    raise SSHProvisionError(
        f"Timed out after {timeout_sec}s waiting for SSH/script on "
        f"{ssh_host or 'guest (QEMU agent discovery)'}"
    )


# ---------------------------------------------------------------------------
# Task polling helper
# ---------------------------------------------------------------------------

def _poll_task(ctx: click.Context, client: ProxmoxClient, node: str, upid: Optional[str], label: str) -> None:
    """If ``upid`` is a UPID string, poll until done."""
    if not upid:
        return
    _say_step(f"  Waiting for task: {label}...")
    try:
        wait_for_task(client, node, upid, timeout=600.0, poll=2.0)
    except TaskTimeoutError as exc:
        _fail(ctx, exc)
    except TaskFailedError as exc:
        _fail(ctx, exc)


# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

def _configure_logging(*, debug: bool) -> None:
    level = logging.DEBUG if debug else logging.INFO
    try:
        logging.basicConfig(level=level, format=LOG_FORMAT, force=True)
    except TypeError:
        logging.basicConfig(level=level, format=LOG_FORMAT)

    logging.getLogger("urllib3").setLevel(logging.WARNING)
    logging.getLogger("urllib3.connectionpool").setLevel(logging.WARNING)
    logging.getLogger("paramiko").setLevel(logging.WARNING)


# ---------------------------------------------------------------------------
# Shared CLI options
# ---------------------------------------------------------------------------

_config_option = click.option(
    "-c", "--config", "config_path",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    default=DEFAULT_CONFIG, show_default=True,
    help="Path to config.yaml",
)


# ---------------------------------------------------------------------------
# CLI group
# ---------------------------------------------------------------------------

@click.group()
@click.pass_context
@click.option("--no-color", is_flag=True, help="Disable ANSI colors (honors NO_COLOR env).")
@click.option("-v", "--verbose", is_flag=True, help="Verbose logging (same as --debug).")
@click.option("--debug", is_flag=True, help="DEBUG: log API requests and sanitized responses.")
def main(ctx: click.Context, no_color: bool, verbose: bool, debug: bool) -> None:
    """Deploy KVM VMs on Proxmox VE using YAML config and templates."""
    ctx.ensure_object(dict)
    ctx.obj["no_color"] = bool(no_color)
    _configure_logging(debug=verbose or debug)


# ===== DEPLOY =============================================================

@main.command("deploy")
@click.pass_context
@_config_option
@click.option("-t", "--template", "template_path",
              type=click.Path(path_type=Path, exists=True, dir_okay=False),
              required=True, help="Deployment template YAML")
@click.option("--dry-run", is_flag=True, help="Validate and show plan; do not create anything.")
@click.option("--ssh-key", type=click.Path(path_type=Path, exists=True, dir_okay=False),
              default=None, help="SSH private key for guest provisioning")
@click.option("--ssh-user", "--user", "ssh_user", default="root", show_default=True,
              callback=_validate_ssh_user, help="Linux user for SSH provisioning")
@click.option("--ssh-host", default=None, callback=_validate_ssh_host,
              help="Guest SSH address (default: discover via QEMU agent)")
@click.option("--ssh-port", default=22, show_default=True, type=click.IntRange(1, 65535),
              help="Guest SSH port")
@click.option("--ssh-password", default=None, help="SSH password (prefer PROXDEPLOY_SSH_PASSWORD)")
@click.option("--ssh-timeout", default=600, show_default=True, type=click.IntRange(30, 86400),
              help="Max seconds to wait for SSH and script")
@click.option("--ssh-strict-host-key", is_flag=True,
              help="Reject unknown SSH host keys")
def deploy_cmd(
    ctx: click.Context,
    config_path: Path,
    template_path: Path,
    dry_run: bool,
    ssh_key: Optional[Path],
    ssh_user: str,
    ssh_host: Optional[str],
    ssh_port: int,
    ssh_password: Optional[str],
    ssh_timeout: int,
    ssh_strict_host_key: bool,
) -> None:
    """Create a VM from config + template, then start it."""
    try:
        config = _load_config(config_path)

        tpl_raw = _load_yaml(template_path)
        if not isinstance(tpl_raw, dict):
            raise ConfigurationError("template root must be a mapping")
        validated = validate_template(tpl_raw, template_path=template_path)

        node, spec = _build_spec(config, validated)
        storage = spec.storage

        _say_step(f"Target API: {config['proxmox']['host']}  (node: {node})")
        _say_step(f"Template:   {template_path.name}  →  VM name {spec.name!r}")
        _say_step(f"Resources:  {spec.memory} MB RAM, {spec.cores} cores, {spec.disk_gb} GB disk")
        if validated.os_image:
            _say_step(f"OS image:   {validated.os_image} (used as disk source reference)")
        has_cloudinit = bool(validated.ci_user or validated.ci_ssh_keys or validated.ci_password)
        if has_cloudinit:
            _say_step(f"Cloud-init: user={validated.ci_user or 'default'}, "
                       f"ssh_keys={len(validated.ci_ssh_keys)}, "
                       f"ip={validated.ci_ip_config or 'dhcp'}")
        if validated.script:
            _say_step(f"Script:     {validated.script}")

        if dry_run:
            _say_success("Dry-run complete — plan is valid. No changes made.")
            return

        client, api_host, _ = _build_client(config)

        existing = vm_exists(client, node, spec.name)
        if existing:
            _say_warn(f"A VM named {spec.name!r} already exists (vmid={existing}). Creating anyway.")

        result: VMCreateResult = create_kvm(client, node, spec)
        vmid = result.vmid
        _poll_task(ctx, client, node, result.upid, f"VM {vmid} creation")
        _say_success(f"Created VM ID {vmid}")

        if validated.os_image:
            from proxmox.cloudinit import import_disk_image
            _say_step(f"  Setting disk image: {validated.os_image}")
            try:
                import_disk_image(client, node, vmid, storage, validated.os_image)
            except ProxmoxAPIError as exc:
                _say_warn(f"os_image attach failed: {exc}. Continuing without it.")

        if has_cloudinit:
            _say_step("  Applying cloud-init config...")
            ci_keys: list[str] = list(validated.ci_ssh_keys)
            if ssh_key and ssh_key.with_suffix(".pub").is_file():
                pub = ssh_key.with_suffix(".pub").read_text(encoding="utf-8").strip()
                if pub and pub not in ci_keys:
                    ci_keys.append(pub)
            ci = CloudInitConfig(
                user=validated.ci_user,
                password=validated.ci_password,
                ssh_authorized_keys=ci_keys,
                ip_config=validated.ci_ip_config or "ip=dhcp",
                nameserver=validated.ci_nameserver,
                searchdomain=validated.ci_searchdomain,
            )
            try:
                apply_cloudinit(client, node, vmid, storage, ci)
                _say_success("Cloud-init config applied")
            except ProxmoxAPIError as exc:
                _say_warn(f"Cloud-init apply failed: {exc}")

        start_upid = start_vm(client, node, vmid)
        _poll_task(ctx, client, node, start_upid, f"VM {vmid} start")
        _say_success(f"Started VM {vmid} ({spec.name})")

        if validated.script:
            script_path = validated.script_resolved_path
            if script_path is None:
                raise ConfigurationError("internal: script set but script_resolved_path is missing")
            pw = _resolve_ssh_password(validated, ssh_password)
            if not ssh_key and not pw:
                raise ConfigurationError(
                    "Template defines script: provide --ssh-key and/or "
                    "PROXDEPLOY_SSH_PASSWORD (or --ssh-password / template ssh_password)"
                )
            dest = ssh_host or (validated.ssh_host if validated.ssh_host else None)
            trust = not ssh_strict_host_key
            _say_step(
                f"  SSH provisioning: {ssh_user}@{dest or '<agent IP>'}:{ssh_port} → {script_path.name}"
            )
            _provision_via_ssh(
                client, node, vmid,
                ssh_host=dest, ssh_port=ssh_port, ssh_user=ssh_user,
                script_path=script_path, ssh_key=ssh_key, ssh_password=pw,
                timeout_sec=ssh_timeout, trust_host_key=trust,
            )

    except (ConfigurationError, ProxmoxAuthError, ProxmoxAPIError, SSHProvisionError,
            TaskFailedError, TaskTimeoutError) as exc:
        _fail(ctx, exc)


# ===== LIST ================================================================

@main.command("list")
@click.pass_context
@_config_option
@click.option("--format", "output_format", type=click.Choice(["table", "json"]),
              default="table", show_default=True, help="Output format")
def list_cmd(ctx: click.Context, config_path: Path, output_format: str) -> None:
    """List all QEMU VMs on the configured node."""
    try:
        config = _load_config(config_path)
        client, api_host, node = _build_client(config)

        _say_step(f"Listing VMs on {api_host} node={node}")
        vms = list_vms(client, node)

        if not vms:
            _say("No VMs found.")
            return

        if output_format == "json":
            import json
            click.echo(json.dumps(vms, indent=2, default=str))
            return

        header = f"{'VMID':<8} {'NAME':<24} {'STATUS':<12} {'MEM(MB)':<10} {'CPUS':<6} {'PID':<8}"
        sep = "-" * len(header)
        _say(sep, bold=True)
        _say(header, bold=True)
        _say(sep, bold=True)
        for vm in sorted(vms, key=lambda v: int(v.get("vmid", 0))):
            vmid = vm.get("vmid", "?")
            name = vm.get("name", "")
            status = vm.get("status", "unknown")
            mem = vm.get("maxmem", 0)
            mem_mb = int(mem) // (1024 * 1024) if isinstance(mem, (int, float)) and mem > 1024 else mem
            cpus = vm.get("cpus", vm.get("maxcpu", "?"))
            pid = vm.get("pid", "-")
            status_color = {"running": "green", "stopped": "red"}.get(str(status))
            line = f"{vmid:<8} {str(name):<24} "
            ctx_obj = click.get_current_context(silent=True)
            if _use_color(ctx_obj) and status_color:
                click.echo(line, nl=False)
                click.secho(f"{str(status):<12}", fg=status_color, nl=False)
                click.echo(f" {str(mem_mb):<10} {str(cpus):<6} {str(pid):<8}")
            else:
                _say(f"{vmid:<8} {str(name):<24} {str(status):<12} {str(mem_mb):<10} {str(cpus):<6} {str(pid):<8}")

    except (ConfigurationError, ProxmoxAuthError, ProxmoxAPIError) as exc:
        _fail(ctx, exc)


# ===== DESTROY =============================================================

@main.command("destroy")
@click.pass_context
@_config_option
@click.option("--vmid", required=True, type=int, help="VM ID to destroy")
@click.option("--force", is_flag=True, help="Skip confirmation prompt")
@click.option("--purge", is_flag=True, default=True, show_default=True,
              help="Also remove disks and firewall rules")
def destroy_cmd(ctx: click.Context, config_path: Path, vmid: int, force: bool, purge: bool) -> None:
    """Stop and destroy a VM by ID."""
    try:
        config = _load_config(config_path)
        client, api_host, node = _build_client(config)

        _say_step(f"Target: {api_host} node={node}, VM {vmid}")

        try:
            status = get_vm_status(client, node, vmid)
        except ProxmoxAPIError:
            _fail(ctx, ProxmoxAPIError(f"VM {vmid} not found on node {node}"))

        vm_name = status.get("name", "<unknown>")
        vm_status = status.get("status", "unknown")
        _say_step(f"  VM {vmid}: name={vm_name!r}, status={vm_status}")

        if not force:
            confirmed = click.confirm(
                f"Destroy VM {vmid} ({vm_name!r}) on node {node}? This cannot be undone",
                abort=True,
            )

        if vm_status == "running":
            _say_step(f"  Stopping VM {vmid}...")
            stop_upid = stop_vm(client, node, vmid)
            _poll_task(ctx, client, node, stop_upid, f"VM {vmid} stop")
            _say_success(f"VM {vmid} stopped")

        _say_step(f"  Destroying VM {vmid} (purge={purge})...")
        del_upid = destroy_vm(client, node, vmid, purge=purge)
        _poll_task(ctx, client, node, del_upid, f"VM {vmid} destroy")
        _say_success(f"VM {vmid} ({vm_name}) destroyed")

    except click.Abort:
        _say("Aborted.", fg="yellow")
    except (ConfigurationError, ProxmoxAuthError, ProxmoxAPIError,
            TaskFailedError, TaskTimeoutError) as exc:
        _fail(ctx, exc)


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------

def run() -> None:
    """Console script entrypoint."""
    try:
        main(prog_name="proxdeploy")
    except click.exceptions.Exit as exc:
        code = exc.exit_code
        if code is None:
            code = 0
        sys.exit(code)
    except click.ClickException as exc:
        logger.error("%s", exc)
        ctx = click.get_current_context(silent=True)
        if ctx and _use_color(ctx):
            click.secho(f"Error: {exc}", fg="red", err=True)
        else:
            click.echo(f"Error: {exc}", err=True)
        sys.exit(1)


if __name__ == "__main__":
    run()