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, 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": {},
}

_WG_TAGS = {"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", {}),
    )


@registry.register("GET", "/wireguard/config", cache_tags=_WG_TAGS)
def get_config(_request: Any, _body: Any) -> dict[str, Any]:
    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", invalidate=_WG_TAGS)
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)
    return {"config_saved": True}


@registry.register("PATCH", "/wireguard/config", invalidate=_WG_TAGS)
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)
    return {"config_saved": True}


@registry.register("POST", "/wireguard/apply", invalidate=_WG_TAGS)
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"])
    return {"applied": True}


@registry.register("POST", "/wireguard/down", invalidate=_WG_TAGS)
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)
    return {"down": True}


@registry.register("GET", "/wireguard/status", cache_tags=_WG_TAGS)
def status(_request: Any, _body: Any) -> dict[str, Any]:
    cfg = _get_config()
    name = cfg["interface"]["name"]
    result: dict[str, Any] = {"up": False, "interface": {}, "peers": []}
    try:
        res = run_proc([WG_BIN, "show", name], sudo=True, check=False)
        if res.returncode != 0:
            return result
        raw = res.stdout.strip()
    except Exception:
        return result

    current_peer: dict[str, Any] | None = None
    peers: list[dict[str, Any]] = []
    for line in raw.splitlines():
        line = line.strip()
        if not line:
            continue
        if line.startswith("interface:"):
            result["up"] = True
            result["interface"] = {}
            current_peer = None
            continue
        if line.startswith("public key:"):
            result["interface"]["public_key"] = line.split(":", 1)[1].strip()
            continue
        if line.startswith("listening port:"):
            result["interface"]["listen_port"] = int(line.split(":", 1)[1].strip())
            continue
        if line.startswith("fwmark:"):
            result["interface"]["fwmark"] = line.split(":", 1)[1].strip()
            continue
        if line.startswith("peer:"):
            cur_key = line.split(":", 1)[1].strip()
            current_peer = {
                "public_key": cur_key,
                "endpoint": None,
                "allowed_ips": [],
                "latest_handshake": None,
                "transfer_received": 0,
                "transfer_sent": 0,
                "persistent_keepalive": None,
            }
            peers.append(current_peer)
            continue
        if current_peer is None:
            continue
        if line.startswith("endpoint:"):
            current_peer["endpoint"] = line.split(":", 1)[1].strip()
        elif line.startswith("allowed ips:"):
            current_peer["allowed_ips"] = line.split(":", 1)[1].strip().split(", ")
        elif line.startswith("latest handshake:"):
            current_peer["latest_handshake"] = line.split(":", 1)[1].strip()
        elif line.startswith("transfer:"):
            rest = line.split(":", 1)[1].strip().split(", ")
            if rest:
                current_peer["transfer_received"] = rest[0].strip()
            if len(rest) > 1:
                current_peer["transfer_sent"] = rest[1].strip()
        elif line.startswith("persistent-keepalive:"):
            try:
                current_peer["persistent_keepalive"] = int(
                    line.split(":", 1)[1].strip()
                )
            except ValueError:
                current_peer["persistent_keepalive"] = None
    result["peers"] = peers
    return result


@registry.register("POST", "/wireguard/initialize", invalidate=_WG_TAGS)
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])
    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", invalidate=_WG_TAGS)
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)
    peer_out = dict(peers[name])
    peer_out.pop("private_key", None)
    return peer_out


@registry.register("DELETE", "/wireguard/peers/remove", invalidate=_WG_TAGS)
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)
    return {"name": name}


@registry.register("GET", "/wireguard/peers", cache_tags=_WG_TAGS)
def list_peers(_request: Any, _body: Any) -> list[dict[str, Any]]:
    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", cache_tags=_WG_TAGS)
def get_peer_status(_request: Any, _body: Any) -> list[dict[str, Any]]:
    st = status(None, None)
    return st.get("peers", [])


@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}