vacuum-wall/daemon/handlers/wireguard.py

"""WireGuard daemon handler."""

import logging
import os
from copy import deepcopy
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from jinja2 import Environment, FileSystemLoader

from daemon.server import NotFoundError, refresh_state, registry
from lib.common import deep_merge, load_json, run, run_proc, save_json

logger = logging.getLogger(__name__)

PROJECT_DIR = Path(__file__).resolve().parent.parent.parent
CONFIG_PATH = PROJECT_DIR / "config" / "wireguard" / "config.json"
WG_CONF_PATH = "/etc/wireguard/wg0.conf"
WG_QUICK_BIN = "wg-quick"
WG_BIN = "wg"

ENV = Environment(
    loader=FileSystemLoader(str(PROJECT_DIR / "system")),
    autoescape=False,
    lstrip_blocks=True,
    trim_blocks=True,
)

DEFAULT_CONFIG: dict[str, Any] = {
    "interface": {
        "name": "wg0",
        "listen_port": 51820,
        "private_key": "",
        "public_key": "",
        "addresses": ["10.137.0.1/24"],
        "post_up": None,
        "post_down": None,
    },
    "peers": {},
}


def _get_state() -> dict[str, Any] | None:
    """Retrieve cached WireGuard state from the global state store."""
    from lib.state import state as state_store

    return state_store.get("wireguard")


def _get_config() -> dict[str, Any]:
    """Load and merge the WireGuard config with defaults."""
    return deep_merge(deepcopy(DEFAULT_CONFIG), load_json(CONFIG_PATH))


def _save_config(cfg: dict[str, Any]) -> None:
    """Persist the WireGuard config to disk."""
    save_json(CONFIG_PATH, cfg)


def _generate_conf(cfg: dict[str, Any]) -> str:
    """Render the WireGuard server config file from Jinja template."""
    tmpl = ENV.get_template("wireguard.conf")
    return tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        interface=cfg["interface"],
        peers=cfg.get("peers", {}),
    )


def _get_wg_state() -> dict[str, Any]:
    """Return cached WireGuard state, or empty dict if not yet loaded."""
    wg = _get_state()
    if wg is None:
        return {}
    return 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_config()
    safe = dict(cfg)
    if "interface" in safe:
        safe["interface"] = dict(safe["interface"])
        safe["interface"].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 key.

    Raises:
        ValueError: When request body is missing.
    """
    if not body:
        raise ValueError("Request body required")
    current = _get_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
    _save_config(body)
    refresh_state(["wireguard"])
    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)
    current = _get_config()
    merged = deep_merge(current, body)
    _save_config(merged)
    refresh_state(["wireguard"])
    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 tunnel via sudo."""
    cfg = _get_config()
    conf_text = _generate_conf(cfg)
    _save_config(cfg)
    local_dir = PROJECT_DIR / "data" / "wireguard"
    local_dir.mkdir(parents=True, exist_ok=True)
    local_tmp = local_dir / "wg0.conf.tmp"
    with open(local_tmp, "w") as f:
        f.write(conf_text)
    os.chmod(local_tmp, 0o600)
    run(["cp", "--", str(local_tmp), WG_CONF_PATH], sudo=True)
    run(["chown", "root:root", WG_CONF_PATH], sudo=True, check=False)
    local_tmp.unlink(missing_ok=True)
    run([WG_QUICK_BIN, "up", cfg["interface"]["name"]], sudo=True)
    logger.info("WireGuard tunnel '%s' brought up", cfg["interface"]["name"])
    refresh_state(["wireguard"])
    return {"applied": True}


@registry.register("POST", "/wireguard/down")
def down(_request: Any, _body: Any) -> dict[str, Any]:
    """POST /wireguard/down — bring down the WireGuard tunnel via sudo."""
    cfg = _get_config()
    name = cfg["interface"]["name"]
    run([WG_QUICK_BIN, "down", name], sudo=True)
    logger.info("WireGuard tunnel '%s' brought down", name)
    refresh_state(["wireguard"])
    return {"down": True}


@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": []})
    return {"up": False, "interface": {}, "peers": []}


@registry.register("POST", "/wireguard/initialize")
def initialize(_request: Any, _body: Any) -> dict[str, Any]:
    """POST /wireguard/initialize — generate keypair and store in config (idempotent)."""
    cfg = _get_config()
    if cfg["interface"].get("private_key"):
        return {"initialized": False, "reason": "already initialized"}
    res = run_proc([WG_BIN, "genkey"], sudo=True)
    private_key = res.stdout.strip()
    res2 = run_proc([WG_BIN, "pubkey"], sudo=True, input=private_key)
    public_key = res2.stdout.strip()
    cfg["interface"]["private_key"] = private_key
    cfg["interface"]["public_key"] = public_key
    _save_config(cfg)
    logger.info("WireGuard initialised (pubkey=%s...)", public_key[:16])
    refresh_state(["wireguard"])
    safe = dict(cfg)
    safe["interface"] = dict(safe["interface"])
    safe["interface"].pop("private_key", None)
    return {"initialized": True, "config": safe}


@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_config()
    peers = cfg.setdefault("peers", {})
    allowed_ips = body.get("allowed_ips", [])
    if name in peers:
        peer = peers[name]
        peer["endpoint"] = body.get("endpoint")
        peer["allowed_ips"] = allowed_ips
        peer["persistent_keepalive"] = body.get("persistent_keepalive")
        if body.get("preshared_key") is not None:
            peer["preshared_key"] = body["preshared_key"]
        logger.info("WireGuard peer '%s' updated", name)
    else:
        res = run_proc([WG_BIN, "genkey"], sudo=True)
        priv = res.stdout.strip()
        res2 = run_proc([WG_BIN, "pubkey"], sudo=True, input=priv)
        pub = res2.stdout.strip()
        peers[name] = {
            "public_key": pub,
            "private_key": priv,
            "endpoint": body.get("endpoint"),
            "allowed_ips": allowed_ips,
            "persistent_keepalive": body.get("persistent_keepalive"),
            "preshared_key": body.get("preshared_key"),
        }
        logger.info("WireGuard peer '%s' added", name)
    _save_config(cfg)
    refresh_state(["wireguard"])
    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_config()
    peers = cfg.setdefault("peers", {})
    if name not in peers:
        raise NotFoundError(f"Peer '{name}' not found")
    del peers[name]
    _save_config(cfg)
    logger.info("WireGuard peer '%s' removed", name)
    refresh_state(["wireguard"])
    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", [])
    cfg = _get_config()
    result: list[dict[str, Any]] = []
    for name, info in cfg.get("peers", {}).items():
        entry = dict(info)
        entry["name"] = name
        entry.pop("private_key", None)
        result.append(entry)
    return result


@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("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.

    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_config()
    if name not in cfg.get("peers", {}):
        raise NotFoundError(f"Peer '{name}' not found")
    peer = cfg["peers"][name]
    client_priv = peer.get("private_key", "")
    if not client_priv:
        raise NotFoundError(f"Peer '{name}' has no private key")
    iface = cfg["interface"]
    sorted_peers = sorted(cfg.get("peers", {}).keys())
    peer_index = sorted_peers.index(name) + 2
    srv_addr = iface["addresses"][0] if iface["addresses"] else "10.137.0.1/24"
    addr_part, prefix = srv_addr.rsplit("/", 1)
    prefix_base = addr_part.rsplit(".", 1)[0]
    client_addr = f"{prefix_base}.{peer_index}/{prefix}"
    tmpl = ENV.get_template("wireguard-client.conf")
    conf = tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        peer_name=name,
        client_priv=client_priv,
        client_addr=client_addr,
        server_pubkey=iface.get("public_key", ""),
        server_endpoint=server_endpoint,
        allowed_ips=peer.get("allowed_ips", ["0.0.0.0/0"]),
        preshared_key=peer.get("preshared_key"),
        persistent_keepalive=peer.get("persistent_keepalive"),
    )
    return {"config": conf}