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:
    from lib.state import state as state_store

    return state_store.get("wireguard")


def _get_config() -> dict[str, Any]:
    return deep_merge(deepcopy(DEFAULT_CONFIG), load_json(CONFIG_PATH))


def _save_config(cfg: dict[str, Any]) -> None:
    save_json(CONFIG_PATH, cfg)


def _generate_conf(cfg: dict[str, Any]) -> str:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]:
    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]]:
    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]]:
    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]:
    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}