vacuum-wall/daemon/handlers/wireguard.py

"""WireGuard daemon handler."""

import logging
import os
from pathlib import Path
from typing import Any

from daemon.handlers.common import emit_and_refresh
from daemon.iface import (
    DELETE_WIREGUARD_CLASSES,
    DELETE_WIREGUARD_CLASSES_DOWN,
    DELETE_WIREGUARD_PEERS_REMOVE,
    GET_WIREGUARD_CLASS_STATUS,
    GET_WIREGUARD_CLASSES,
    GET_WIREGUARD_CONFIG,
    GET_WIREGUARD_PEER_STATUS,
    GET_WIREGUARD_PEERS,
    GET_WIREGUARD_STATUS,
    PATCH_WIREGUARD_CLASSES,
    PATCH_WIREGUARD_CONFIG,
    POST_WIREGUARD_APPLY,
    POST_WIREGUARD_CLASS_INIT_KEYS,
    POST_WIREGUARD_CLASSES,
    POST_WIREGUARD_CLASSES_UP,
    POST_WIREGUARD_CONFIG,
    POST_WIREGUARD_DOWN,
    POST_WIREGUARD_GENERATE_CLIENT,
    POST_WIREGUARD_INITIALIZE,
    POST_WIREGUARD_PEERS_ADD,
)
from daemon.server import ConflictError, NotFoundError, registry
from lib.common import deep_merge, run, stamp_applied, strip_apply_meta
from lib.wireguard import (
    _class_interface_name,
    _class_peers,
    _ensure_access_classes,
    _wg_conf_path,
    generate_class_conf,
    generate_class_keypair,
    generate_conf,
    generate_keypair,
)
from lib.wireguard import (
    generate_client_conf as _gen_client_conf,
)
from lib.wireguard import (
    get_config as _get_wireguard_config,
)
from lib.wireguard import (
    get_peers as _get_wireguard_peers,
)
from lib.wireguard import (
    save_config as _save_wireguard_config,
)

logger = logging.getLogger(__name__)

WG_QUICK_BIN = "wg-quick"


def _get_wg_state() -> dict[str, Any]:
    """Return cached WireGuard state, or empty dict if not yet loaded."""
    from lib.state import state as state_store

    wg = state_store.get("wireguard")
    return {} if wg is None else wg


# ---------------------------------------------------------------------------
# Routes


@registry.register(GET_WIREGUARD_CONFIG)
def get_config(_request: Any, _body: Any) -> dict[str, Any]:
    """GET /wireguard/config — return WireGuard config with private key stripped."""
    wg = _get_wg_state()
    if wg:
        return wg.get("config", {})
    cfg = _get_wireguard_config()
    safe = strip_apply_meta(cfg)
    if "interface" in safe:
        safe["interface"] = dict(safe["interface"])
        safe["interface"].pop("private_key", None)
    if "access_classes" in safe:
        for ck, cv in safe["access_classes"].items():
            if isinstance(cv, dict):
                safe["access_classes"][ck] = dict(cv)
                safe["access_classes"][ck].pop("private_key", None)
    return safe


@registry.register(POST_WIREGUARD_CONFIG)
def save_config_handler(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/config — replace config, preserving existing private keys.

    Raises:
        ValueError: When request body is missing.
    """
    if not body:
        raise ValueError("Request body required")

    current = _get_wireguard_config()
    current_key = current.get("interface", {}).get("private_key", "")

    if "interface" in body:
        body = dict(body)
        body["interface"] = dict(body["interface"])
        body["interface"].pop("private_key", None)

    if current_key:
        body.setdefault("interface", {})["private_key"] = current_key

    # Preserve class private keys
    if "access_classes" in body:
        current_classes = current.get("access_classes", {})
        for ck, cv in body.get("access_classes", {}).items():
            if isinstance(cv, dict) and ck in current_classes:
                cur_class_pk = current_classes[ck].get("private_key", "")
                if cur_class_pk and ck in body["access_classes"]:
                    body["access_classes"][ck]["private_key"] = cur_class_pk
    elif "access_classes" not in body:
        body["access_classes"] = current.get("access_classes", {})

    _save_wireguard_config(body)
    emit_and_refresh("wireguard", {"action": "config_saved"})
    return {"config_saved": True}


@registry.register(PATCH_WIREGUARD_CONFIG)
def patch_config(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """PATCH /wireguard/config — deep-merge patch into existing config.

    Raises:
        ValueError: When request body is missing.
    """
    if not body:
        raise ValueError("Request body required")

    if "interface" in body:
        body = dict(body)
        body["interface"] = dict(body["interface"])
        body["interface"].pop("private_key", None)

    # Strip class private keys from patch body
    if "access_classes" in body:
        for _ck, cv in body["access_classes"].items():
            if isinstance(cv, dict):
                cv.pop("private_key", None)

    current = _get_wireguard_config()
    merged = deep_merge(current, body)
    _save_wireguard_config(merged)
    emit_and_refresh("wireguard", {"action": "config_patched"})
    return {"config_saved": True}


@registry.register(POST_WIREGUARD_APPLY)
def apply(_request: Any, _body: Any) -> dict[str, Any]:
    """POST /wireguard/apply — render config, write to disk, bring up tunnels via sudo.

    In multi-interface mode, applies each class with assigned peers
    independently. Falls back to legacy single-interface mode.
    """
    cfg = _get_wireguard_config()
    classes = cfg.get("access_classes", {})
    affected: list[str] = []

    for class_key in classes:
        class_cfg = classes.get(class_key)
        if not class_cfg or not isinstance(class_cfg, dict):
            continue
        if not _class_peers(cfg, class_key):
            continue

        try:
            conf_text = generate_class_conf(cfg, class_key)
            if not conf_text:
                continue
        except ValueError as e:
            logger.warning("Skipping class '%s' during apply: %s", class_key, e)
            continue

        ifname = _class_interface_name(class_key)
        conf_path = _wg_conf_path(ifname)
        local_tmp = Path(f"/run/vacuum-wall/{ifname}.conf.tmp")
        local_tmp.parent.mkdir(exist_ok=True)
        local_tmp.write_text(conf_text)
        os.chmod(local_tmp, 0o600)
        run(["cp", "--", str(local_tmp), conf_path], sudo=True)
        run(["chown", "root:root", conf_path], sudo=True, check=False)
        local_tmp.unlink(missing_ok=True)
        run([WG_QUICK_BIN, "up", ifname], sudo=True, check=False)
        logger.info("WireGuard tunnel '%s' (class '%s') brought up", ifname, class_key)
        affected.append(f"{ifname}.conf")

    if not affected:
        # Legacy single-interface mode
        conf_text = generate_conf(cfg)
        ifname = cfg["interface"]["name"]
        conf_path = _wg_conf_path(ifname)
        local_tmp = Path(f"/run/vacuum-wall/{ifname}.conf.tmp")
        local_tmp.parent.mkdir(exist_ok=True)
        local_tmp.write_text(conf_text)
        os.chmod(local_tmp, 0o600)
        run(["cp", "--", str(local_tmp), conf_path], sudo=True)
        run(["chown", "root:root", conf_path], sudo=True, check=False)
        local_tmp.unlink(missing_ok=True)
        run([WG_QUICK_BIN, "up", ifname], sudo=True, check=False)
        logger.info("WireGuard tunnel '%s' brought up", ifname)

    cfg_after = _get_wireguard_config()
    stamp_applied(cfg_after)
    _save_wireguard_config(cfg_after)
    synced = emit_and_refresh("wireguard", {"action": "config_applied"})
    return {
        "applied": True,
        "synced": synced,
        "interfaces": affected,
    }


@registry.register(POST_WIREGUARD_DOWN)
def down(_request: Any, _body: Any) -> dict[str, Any]:
    """POST /wireguard/down — bring down all WireGuard tunnel interfaces."""
    cfg = _get_wireguard_config()
    classes = cfg.get("access_classes", {})
    for class_key in classes:
        class_cfg = classes.get(class_key)
        if not class_cfg or not isinstance(class_cfg, dict):
            continue
        ifname = _class_interface_name(class_key)
        try:
            run([WG_QUICK_BIN, "down", ifname], sudo=True)
            logger.info("WireGuard tunnel '%s' brought down", ifname)
        except Exception:
            pass

    # Legacy interface (skip if name matches any class interface)
    ifname = cfg["interface"].get("name", "")
    if ifname:
        class_names = {_class_interface_name(k) for k in classes}
        if ifname not in class_names:
            try:
                run([WG_QUICK_BIN, "down", ifname], sudo=True)
                logger.info("WireGuard tunnel '%s' brought down", ifname)
            except Exception:
                pass

    emit_and_refresh("wireguard", {"action": "tunnel_down"})
    return {"down": True}


@registry.register(POST_WIREGUARD_CLASSES_UP)
def class_up(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/classes/<key>/up — bring up a single class's tunnel.

    Raises:
        ValueError: When body or key is missing.
        NotFoundError: When class does not exist.
    """
    if not body:
        raise ValueError("Request body required")
    class_key = body.get("class_key", "").strip()
    if not class_key:
        raise ValueError("'class_key' is required")

    cfg = _get_wireguard_config()
    class_cfg = cfg.get("access_classes", {}).get(class_key)
    if not class_cfg:
        raise NotFoundError(f"Access class '{class_key}' not found")

    conf_text = generate_class_conf(cfg, class_key)
    if not conf_text:
        raise ValueError(f"No peers assigned to class '{class_key}'")

    ifname = _class_interface_name(class_key)
    conf_path = _wg_conf_path(ifname)
    local_tmp = Path(f"/run/vacuum-wall/{ifname}.conf.tmp")
    local_tmp.parent.mkdir(exist_ok=True)
    local_tmp.write_text(conf_text)
    os.chmod(local_tmp, 0o600)
    run(["cp", "--", str(local_tmp), conf_path], sudo=True)
    run(["chown", "root:root", conf_path], sudo=True, check=False)
    local_tmp.unlink(missing_ok=True)
    run([WG_QUICK_BIN, "up", ifname], sudo=True, check=False)
    logger.info("WireGuard class '%s' tunnel '%s' brought up", class_key, ifname)
    emit_and_refresh("wireguard", {"action": "class_up", "class_key": class_key})
    return {"up": True, "interface": ifname}


@registry.register(DELETE_WIREGUARD_CLASSES_DOWN)
def class_down(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """DELETE /wireguard/classes/<key>/down — bring down a single class's tunnel.

    Raises:
        ValueError: When body or key is missing.
    """
    if not body:
        raise ValueError("Request body required")
    class_key = body.get("class_key", "").strip()
    if not class_key:
        raise ValueError("'class_key' is required")

    cfg = _get_wireguard_config()
    if class_key not in cfg.get("access_classes", {}):
        raise NotFoundError(f"Access class '{class_key}' not found")

    ifname = _class_interface_name(class_key)
    try:
        run([WG_QUICK_BIN, "down", ifname], sudo=True)
        logger.info("WireGuard class '%s' tunnel '%s' brought down", class_key, ifname)
    except Exception:
        pass
    emit_and_refresh("wireguard", {"action": "class_down", "class_key": class_key})
    return {"down": True, "interface": ifname}


@registry.register(GET_WIREGUARD_STATUS)
def status(_request: Any, _body: Any) -> dict[str, Any]:
    """GET /wireguard/status — return current WireGuard status from cache."""
    wg = _get_wg_state()
    if wg:
        return wg.get(
            "status", {"up": False, "interface": {}, "peers": [], "classes": {}}
        )
    return {"up": False, "interface": {}, "peers": [], "classes": {}}


@registry.register(POST_WIREGUARD_INITIALIZE)
def initialize(_request: Any, _body: Any) -> dict[str, Any]:
    """POST /wireguard/initialize — generate keypair and store in config (idempotent).

    Also generates key pairs for each access class interface.
    """
    cfg = _get_wireguard_config()
    if cfg["interface"].get("private_key"):
        _ensure_access_classes(cfg)

        # Generate keys for classes that need them (idempotent, saves internally)
        for class_key in cfg.get("access_classes", {}):
            generate_class_keypair(class_key)
        return {"initialized": False, "reason": "already initialized"}

    # Fresh init
    priv, pub = generate_keypair()
    cfg["interface"]["private_key"] = priv
    cfg["interface"]["public_key"] = pub
    _ensure_access_classes(cfg)

    # Generate class keys
    for class_key in cfg.get("access_classes", {}):
        generate_class_keypair(class_key)

    _save_wireguard_config(cfg)
    logger.info("WireGuard initialised (pubkey=%s...)", pub[:16])
    emit_and_refresh("wireguard", {"action": "initialized"})

    safe = dict(cfg)
    safe["interface"] = dict(safe["interface"])
    safe["interface"].pop("private_key", None)
    for ck, cv in safe.get("access_classes", {}).items():
        if isinstance(cv, dict):
            safe["access_classes"][ck] = dict(cv)
            safe["access_classes"][ck].pop("private_key", None)
    return {"initialized": True, "config": safe}


@registry.register(POST_WIREGUARD_CLASS_INIT_KEYS)
def init_class_keys(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/classes/keys/<key> — generate keypair for a class.

    Raises:
        ValueError: When body or key is missing.
        NotFoundError: When class does not exist.
    """
    if not body:
        raise ValueError("Request body required")
    class_key = body.get("class_key", "").strip()
    if not class_key:
        raise ValueError("'class_key' is required")

    cfg = _get_wireguard_config()
    if class_key not in cfg.get("access_classes", {}):
        raise NotFoundError(f"Access class '{class_key}' not found")

    class_cfg = cfg["access_classes"][class_key]
    if class_cfg.get("private_key"):
        return {
            "generated": False,
            "class_key": class_key,
            "reason": "already has keys",
        }

    _, pub = generate_class_keypair(class_key)
    return {"generated": True, "class_key": class_key, "public_key": pub}


@registry.register(POST_WIREGUARD_PEERS_ADD)
def add_peer(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/peers/add — add new peer or update existing one.

    Raises:
        ValueError: When body is missing or name is empty.
    """
    if not body:
        raise ValueError("Request body required")
    name = body.get("name", "").strip()
    if not name:
        raise ValueError("'name' is required")
    cfg = _get_wireguard_config()
    peers = cfg.setdefault("peers", {})
    if name in peers:
        peer = peers[name]
        if "endpoint" in body:
            peer["endpoint"] = body["endpoint"]
        if "allowed_ips" in body:
            peer["allowed_ips"] = body["allowed_ips"]
        if "persistent_keepalive" in body:
            peer["persistent_keepalive"] = body["persistent_keepalive"]
        if "preshared_key" in body:
            peer["preshared_key"] = body["preshared_key"]
        if "description" in body:
            peer["description"] = body["description"]
        if "access_class" in body:
            peer["access_class"] = body["access_class"]
        logger.info("WireGuard peer '%s' updated", name)
        _peer_action = "peer_updated"
    else:
        priv, pub = generate_keypair()
        peers[name] = {
            "public_key": pub,
            "private_key": priv,
            "endpoint": body.get("endpoint"),
            "allowed_ips": body.get("allowed_ips", []),
            "persistent_keepalive": body.get("persistent_keepalive"),
            "preshared_key": body.get("preshared_key"),
            "description": body.get("description", ""),
            "access_class": body.get("access_class"),
        }
        logger.info("WireGuard peer '%s' added", name)
        _peer_action = "peer_added"
    _save_wireguard_config(cfg)
    emit_and_refresh("wireguard", {"action": _peer_action, "peer_name": name})
    peer_out = dict(peers[name])
    peer_out.pop("private_key", None)
    return peer_out


@registry.register(DELETE_WIREGUARD_PEERS_REMOVE)
def remove_peer(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """DELETE /wireguard/peers/remove — remove a peer by name.

    Raises:
        ValueError: When body is missing or name is empty.
        NotFoundError: When peer does not exist.
    """
    if not body:
        raise ValueError("Request body required")
    name = body.get("name", "").strip()
    if not name:
        raise ValueError("'name' is required")
    cfg = _get_wireguard_config()
    peers = cfg.setdefault("peers", {})
    if name not in peers:
        raise NotFoundError(f"Peer '{name}' not found")
    del peers[name]
    _save_wireguard_config(cfg)
    logger.info("WireGuard peer '%s' removed", name)
    emit_and_refresh("wireguard", {"action": "peer_removed", "peer_name": name})
    return {"name": name}


@registry.register(GET_WIREGUARD_PEERS)
def list_peers(_request: Any, _body: Any) -> list[dict[str, Any]]:
    """GET /wireguard/peers — return configured peers with private keys stripped."""
    wg = _get_wg_state()
    if wg:
        return wg.get("peers", [])
    return _get_wireguard_peers()


@registry.register(GET_WIREGUARD_PEER_STATUS)
def get_peer_status(_request: Any, _body: Any) -> list[dict[str, Any]]:
    """GET /wireguard/peer-status — return runtime peer status from cache."""
    wg = _get_wg_state()
    if wg:
        return wg.get("status", {}).get("peers", [])
    return []


@registry.register(GET_WIREGUARD_CLASS_STATUS)
def get_class_status(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """GET /wireguard/classes/<key>/status — return status for a specific class interface.

    Raises:
        NotFoundError: When class does not exist.
    """
    wg = _get_wg_state()
    if wg:
        classes = wg.get("status", {}).get("classes", {})
        if body:
            key = body.get("class_key", "")
            if key and key not in classes:
                raise NotFoundError(f"Access class '{key}' not found")
            return classes.get(key, {"up": False, "interface": {}, "peers": []})
        return classes
    return {"up": False, "interface": {}, "peers": []}


@registry.register(POST_WIREGUARD_GENERATE_CLIENT)
def generate_client_conf(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/generate-client — render client-side WireGuard config for a peer.

    Uses the peer's access class subnet and port to derive the correct
    server address and endpoint port.

    Raises:
        ValueError: When body, name, or server_endpoint is missing.
        NotFoundError: When peer does not exist or has no private key.
    """
    if not body:
        raise ValueError("Request body required")
    name = body.get("name", "").strip()
    if not name:
        raise ValueError("'name' is required")
    server_endpoint = body.get("server_endpoint", "")
    if not server_endpoint:
        raise ValueError("'server_endpoint' is required")
    cfg = _get_wireguard_config()
    if name not in cfg.get("peers", {}):
        raise NotFoundError(f"Peer '{name}' not found")
    peer = cfg["peers"][name]
    if not peer.get("private_key"):
        raise NotFoundError(f"Peer '{name}' has no private key")

    conf = _gen_client_conf(
        peer_name=name,
        server_endpoint=server_endpoint,
        server_pubkey=cfg["interface"].get("public_key"),
    )
    return {"config": conf}


# ---------------------------------------------------------------------------
# Access Classes


@registry.register(GET_WIREGUARD_CLASSES)
def list_classes(_request: Any, _body: Any) -> dict[str, Any]:
    """GET /wireguard/classes — return access classes (with private keys stripped)."""
    cfg = _get_wireguard_config()
    classes = cfg.get("access_classes", {})
    safe: dict[str, Any] = {}
    for k, v in classes.items():
        if isinstance(v, dict):
            entry = dict(v)
            entry.pop("private_key", None)
            safe[k] = entry
        else:
            safe[k] = v
    return safe


@registry.register(POST_WIREGUARD_CLASSES)
def create_class(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """POST /wireguard/classes — create a new access class with phase-2 fields.

    Raises:
        ValueError: When body is missing or key conflicts.
    """
    if not body:
        raise ValueError("Request body required")
    key = body.get("key", "").strip()
    if not key:
        raise ValueError("'key' is required")
    if not key.isalnum() or not key.islower():
        raise ValueError("'key' must be lowercase alphanumeric")
    name = body.get("name", key)
    description = body.get("description", "")
    cfg = _get_wireguard_config()
    classes = cfg.setdefault("access_classes", {})
    if key in classes:
        raise ConflictError(f"Access class '{key}' already exists")

    classes[key] = {
        "name": name,
        "description": description,
        "subnet": body.get("subnet"),
        "listen_port": body.get("listen_port"),
        "lan_access": body.get("lan_access", False),
        "private_key": "",
        "public_key": "",
    }
    _save_wireguard_config(cfg)
    emit_and_refresh("wireguard", {"action": "class_created"})
    out = dict(classes[key])
    out.pop("private_key", None)
    return out


@registry.register(PATCH_WIREGUARD_CLASSES)
def update_class(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """PATCH /wireguard/classes/<key> — update an access class.

    Raises:
        ValueError: When body is missing.
        NotFoundError: When class does not exist.
    """
    if not body:
        raise ValueError("Request body required")
    key = body.get("key", "").strip()
    if not key:
        raise ValueError("'key' is required")
    cfg = _get_wireguard_config()
    classes = cfg.setdefault("access_classes", {})
    if key not in classes:
        raise NotFoundError(f"Access class '{key}' not found")
    class_cfg = classes[key]
    for field in ("name", "description", "subnet", "listen_port", "lan_access"):
        if field in body:
            class_cfg[field] = body[field]
    _save_wireguard_config(cfg)
    emit_and_refresh("wireguard", {"action": "class_updated"})
    out = dict(classes[key])
    out.pop("private_key", None)
    return {"key": key, **out}


@registry.register(DELETE_WIREGUARD_CLASSES)
def delete_class(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    """DELETE /wireguard/classes/<key> — remove an access class.

    Raises:
        ValueError: When body is missing.
        NotFoundError: When class does not exist.
        ConflictError: When peers still reference the class.
    """
    if not body:
        raise ValueError("Request body required")
    key = body.get("key", "").strip()
    if not key:
        raise ValueError("'key' is required")
    cfg = _get_wireguard_config()
    classes = cfg.setdefault("access_classes", {})
    if key not in classes:
        raise NotFoundError(f"Access class '{key}' not found")
    # Check if any peers reference this class
    peers_reusing = [
        name
        for name, info in cfg.get("peers", {}).items()
        if info.get("access_class") == key
    ]
    if peers_reusing:
        raise ConflictError(
            f"Cannot delete class '{key}': {len(peers_reusing)} peer(s) reference it: {', '.join(peers_reusing)}"
        )
    del classes[key]
    _save_wireguard_config(cfg)
    emit_and_refresh("wireguard", {"action": "class_deleted"})
    return {"key": key}