vacuum-wall/lib/wireguard.py

"""WireGuard Manager for Vacuum Wall SSL Proxy Firewall.

Generates wg-quick configurations, manages peers, and controls
the WireGuard tunnel interface. Supports multi-interface mode where
each access class gets its own WireGuard interface.
"""

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 lib.common import deep_merge, load_json, run, run_proc, save_json

logger = logging.getLogger(__name__)

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

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


def _wg_conf_path(ifname: str) -> str:
    """Return the system WG conf path for an interface name."""
    return f"/etc/wireguard/{ifname}.conf"


def _default_class_fields() -> dict[str, Any]:
    """Return the default set of fields for an access class entry."""
    return {
        "name": "",
        "description": "",
        "subnet": None,
        "listen_port": None,
        "lan_access": False,
        "private_key": "",
        "public_key": "",
    }


DEFAULT_CONFIG: dict[str, Any] = {
    "interface": {
        "name": "wg0",
        "listen_port": 51820,
        "private_key": "",
        "public_key": "",
        "addresses": ["10.137.0.1/24"],
        "server_endpoint": "",
        "description": "",
        "post_up": None,
        "post_down": None,
    },
    "access_classes": {
        "full": {
            "name": "Full LAN Access",
            "description": "Peers get full access to internal networks",
            "subnet": "10.137.0.0/24",
            "listen_port": 51820,
            "lan_access": True,
            "private_key": "",
            "public_key": "",
        },
        "internet": {
            "name": "Internet Only",
            "description": "Peers can only reach the internet",
            "subnet": "10.137.1.0/24",
            "listen_port": 51821,
            "lan_access": False,
            "private_key": "",
            "public_key": "",
        },
    },
    "peers": {},
}


# --- Core config persistence ---


def get_config() -> dict[str, Any]:
    """Load the current WireGuard configuration from the JSON store."""
    return deep_merge(deepcopy(DEFAULT_CONFIG), load_json(CONFIG_PATH))


def save_config(cfg: dict[str, Any]) -> None:
    """Persist *cfg* to the JSON store atomically."""
    save_json(CONFIG_PATH, cfg)


# --- Key generation ---


def generate_keypair() -> tuple[str, str]:
    """Generate a WireGuard private/public key pair using ``wg`` CLI."""
    res = run_proc([WG_BIN, "genkey"], sudo=False)
    private_key = res.stdout.strip()
    res2 = run_proc([WG_BIN, "pubkey"], sudo=False, input=private_key)
    public_key = res2.stdout.strip()
    return private_key, public_key


# --- Helpers ---


def _class_interface_name(class_key: str) -> str:
    """Derive interface name for an access class key."""
    return f"wg-{class_key}"


def _class_zone_name(class_key: str) -> str:
    """Derive firewall zone name for an access class key."""
    return f"vpn-{class_key}"


def get_class_interface_name(class_key: str) -> str:
    """Public wrapper for `_class_interface_name`."""
    return _class_interface_name(class_key)


def get_class_zone_name(class_key: str) -> str:
    """Public wrapper for `_class_zone_name`."""
    return _class_zone_name(class_key)


def _class_peers(cfg: dict[str, Any], class_key: str) -> dict[str, Any]:
    """Return peers assigned to a given access class."""
    return {
        name: info
        for name, info in cfg.get("peers", {}).items()
        if isinstance(info, dict) and info.get("access_class") == class_key
    }


# --- wg-<class>.conf generation ---


def generate_class_conf(cfg: dict[str, Any], class_key: str) -> str | None:
    """Render a wg-quick config file for a single access class.

    Returns ``None`` when the class has no peers assigned.
    """
    classes = cfg.get("access_classes", {})
    class_cfg = classes.get(class_key)
    if not class_cfg or not isinstance(class_cfg, dict):
        return None

    peers = _class_peers(cfg, class_key)
    if not peers:
        return None

    ifname = _class_interface_name(class_key)
    subnet = class_cfg.get("subnet")
    if not subnet:
        subnet = "10.137.0.0/24"

    _, prefix = subnet.rsplit("/", 1)
    base = subnet.rsplit(".", 1)[0]
    addr = f"{base}.1/{prefix}"

    listen_port = class_cfg.get("listen_port")
    if not listen_port:
        listen_port = 51820

    private_key = class_cfg.get("private_key", "")
    if not private_key:
        raise ValueError(
            f"Access class '{class_key}' has no private key — generate one first."
        )

    class_iface = {
        "name": ifname,
        "listen_port": listen_port,
        "private_key": private_key,
        "addresses": [addr],
        "post_up": cfg.get("interface", {}).get("post_up"),
        "post_down": cfg.get("interface", {}).get("post_down"),
    }

    tmpl = ENV.get_template("wireguard.conf")
    return tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        interface=class_iface,
        peers=peers,
    )


def generate_conf(cfg: dict[str, Any]) -> str:
    """Render a valid wg-quick config file from *cfg* using Jinja2.

    Legacy single-interface mode — uses the top-level ``interface`` block
    and peers without an ``access_class`` assigned.
    """
    fallback_peers = {
        n: p
        for n, p in cfg.get("peers", {}).items()
        if isinstance(p, dict) and not p.get("access_class")
    }
    tmpl = ENV.get_template("wireguard.conf")
    return tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        interface=cfg["interface"],
        peers=fallback_peers if fallback_peers else {},
    )


# --- Apply / down ---


def apply() -> None:
    """Write the current config to disk and bring tunnels up with wg-quick.

    In multi-interface mode, applies each access class's interface independently.
    Falls back to legacy single-interface mode when no classes have peers.
    """
    cfg = get_config()
    classes = cfg.get("access_classes", {})
    applied = False

    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:
            continue

        ifname = _class_interface_name(class_key)
        conf_path = _wg_conf_path(ifname)
        local_dir = PROJECT_DIR / "data" / "wireguard"
        local_dir.mkdir(parents=True, exist_ok=True)
        local_tmp = local_dir / f"{ifname}.conf.tmp"
        with open(local_tmp, "w") as f:
            f.write(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)
        logger.info("WireGuard tunnel '%s' (class '%s') brought up", ifname, class_key)
        applied = True

    if applied:
        save_config(cfg)
    elif applied is False:
        # Legacy single-interface fallback
        legacy_apply(cfg)


def legacy_apply(cfg: dict[str, Any]) -> None:
    """Legacy single-interface apply."""
    conf_text = generate_conf(cfg)
    save_config(cfg)
    ifname = cfg["interface"]["name"]
    conf_path = _wg_conf_path(ifname)
    local_dir = PROJECT_DIR / "data" / "wireguard"
    local_dir.mkdir(parents=True, exist_ok=True)
    local_tmp = local_dir / f"{ifname}.conf.tmp"
    with open(local_tmp, "w") as f:
        f.write(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)
    logger.info("WireGuard tunnel '%s' brought up", ifname)


def apply_class(class_key: str) -> None:
    """Apply config for a single access class interface."""
    cfg = get_config()
    class_cfg = cfg.get("access_classes", {}).get(class_key)
    if not class_cfg or not isinstance(class_cfg, dict):
        raise ValueError(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_dir = PROJECT_DIR / "data" / "wireguard"
    local_dir.mkdir(parents=True, exist_ok=True)
    local_tmp = local_dir / f"{ifname}.conf.tmp"
    with open(local_tmp, "w") as f:
        f.write(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)
    save_config(cfg)
    logger.info("WireGuard tunnel '%s' (class '%s') brought up", ifname, class_key)


def down() -> None:
    """Bring all WireGuard tunnel interfaces down.

    In multi-interface mode, brings down each class interface with peers.
    """
    cfg = get_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
        if 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

    # Also try 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


def down_class(class_key: str) -> None:
    """Bring down a single access class interface."""
    ifname = _class_interface_name(class_key)
    run([WG_QUICK_BIN, "down", ifname], sudo=True)
    logger.info("WireGuard tunnel '%s' brought down", ifname)


# --- Status ---


def status() -> dict[str, Any]:
    """Query the live tunnel state via ``wg show``.

    In multi-interface mode, collects status for all class interfaces.
    Returns combined status dict keyed by interface name.
    """
    cfg = get_config()
    result: dict[str, Any] = {
        "up": False,
        "interface": {},
        "peers": [],
        "classes": {},
    }

    # Collect per-class status
    classes = cfg.get("access_classes", {})
    for class_key in classes:
        ifname = _class_interface_name(class_key)
        try:
            res = run_proc([WG_BIN, "show", ifname], sudo=True, check=False)
            if res.returncode != 0:
                result["classes"][class_key] = {"up": False, "peers": []}
                continue
            class_status = parse_wg_show_output(res.stdout.strip())
            result["classes"][class_key] = class_status
            if class_status["up"]:
                result["up"] = True
        except Exception:
            result["classes"][class_key] = {"up": False, "peers": []}

    # Legacy single-interface status (still collected for backward compat)
    try:
        ifname = cfg["interface"].get("name", "wg0")
        res = run_proc([WG_BIN, "show", ifname], sudo=True, check=False)
        if res.returncode == 0:
            parsed = parse_wg_show_output(res.stdout.strip())
            result["up"] = parsed["up"]
            result["interface"] = parsed.get("interface", {})
            result["peers"] = parsed.get("peers", [])
    except Exception:
        pass

    return result


def parse_wg_show_output(raw: str) -> dict[str, Any]:
    """Parse ``wg show`` output into structured dict.

    Returns ``{"up", "interface", "peers"}`` where *interface* carries
    ``public_key``, ``listen_port`` and (when present) ``fwmark``.
    """
    result: dict[str, Any] = {
        "up": False,
        "interface": {},
        "peers": [],
    }
    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:"):
            val = line.split(":", 1)[1].strip()
            result["interface"]["listen_port"] = int(val)
            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()
            continue

        if line.startswith("allowed ips:"):
            vals = line.split(":", 1)[1].strip().split(", ")
            current_peer["allowed_ips"] = vals
            continue

        if line.startswith("latest handshake:"):
            current_peer["latest_handshake"] = line.split(":", 1)[1].strip()
            continue

        if line.startswith("transfer:"):
            rest = line.split(":", 1)[1].strip()
            parts = rest.split(", ")
            if parts:
                current_peer["transfer_received"] = parts[0].strip()
            if len(parts) > 1:
                current_peer["transfer_sent"] = parts[1].strip()
            continue

        if line.startswith("persistent-keepalive:"):
            val = line.split(":", 1)[1].strip()
            try:
                current_peer["persistent_keepalive"] = int(val)
            except ValueError:
                current_peer["persistent_keepalive"] = None

    result["peers"] = peers
    return result


# --- Peer management ---

_UNSET = object()


def add_peer(
    name: str,
    endpoint: str | None | object = _UNSET,
    allowed_ips: list[str] | None | object = _UNSET,
    persistent_keepalive: int | None | object = _UNSET,
    preshared_key: str | None = None,
    description: str | None = None,
    access_class: str | None = None,
) -> dict[str, Any]:
    """Add (or update) a peer in the configuration."""
    cfg = get_config()
    peers = cfg.setdefault("peers", {})

    if name in peers:
        peer = peers[name]
        if endpoint is not _UNSET:
            peer["endpoint"] = endpoint
        if allowed_ips is not _UNSET:
            peer["allowed_ips"] = allowed_ips if allowed_ips is not None else []
        if persistent_keepalive is not _UNSET:
            peer["persistent_keepalive"] = persistent_keepalive
        if preshared_key is not None:
            peer["preshared_key"] = preshared_key
        if description is not None:
            peer["description"] = description
        if access_class is not None:
            peer["access_class"] = access_class
        logger.info("WireGuard peer '%s' updated", name)
    else:
        priv, pub = generate_keypair()
        peer = {
            "public_key": pub,
            "private_key": priv,
            "endpoint": endpoint if endpoint is not _UNSET else None,
            "allowed_ips": allowed_ips if allowed_ips is not _UNSET else [],
            "persistent_keepalive": persistent_keepalive
            if persistent_keepalive is not _UNSET
            else None,
            "preshared_key": preshared_key,
            "description": description,
            "access_class": access_class,
        }
        peers[name] = peer
        logger.info("WireGuard peer '%s' added (pubkey=%s...)", name, pub[:16])

    save_config(cfg)
    peer_out = dict(peer)
    peer_out.pop("private_key", None)
    return peer_out


def remove_peer(name: str) -> None:
    """Remove a peer from the configuration by name."""
    cfg = get_config()
    cfg.setdefault("peers", {}).pop(name, None)
    save_config(cfg)
    logger.info("WireGuard peer '%s' removed", name)


def get_peers() -> list[dict[str, Any]]:
    """List all configured peers (from the JSON store, *not* live)."""
    cfg = get_config()
    peers: list[dict[str, Any]] = []
    for name, info in cfg.get("peers", {}).items():
        entry = dict(info)
        entry["name"] = name
        entry.pop("private_key", None)
        peers.append(entry)
    return peers


def get_peer_status() -> list[dict[str, Any]]:
    """Return live peer status from ``wg show`` for all interfaces."""
    st = status()
    all_peers: list[dict[str, Any]] = []
    for _class_key, class_st in st.get("classes", {}).items():
        for p in class_st.get("peers", []):
            merged = dict(p)
            merged["access_class"] = _class_key
            all_peers.append(merged)
    if not all_peers:
        all_peers = st.get("peers", [])
    return all_peers


# --- Client config generation ---


def generate_client_conf(
    peer_name: str,
    server_endpoint: str,
    server_pubkey: str | None = None,
) -> str:
    """Build a client-side wg-quick config snippet for *peer_name*.

    Uses the peer's access class to derive server address from the class
    subnet and the class's listen port.
    """
    cfg = get_config()
    iface = cfg["interface"]
    peer = cfg["peers"].get(peer_name)
    if peer is None:
        raise KeyError(f"Peer '{peer_name}' not found in configuration")

    client_priv = peer.get("private_key", "")
    if not client_priv:
        raise ValueError(
            f"Peer '{peer_name}' has no private key — cannot generate client config."
        )

    # Determine class info for address/port
    access_class = peer.get("access_class")
    sorted_peers = sorted(cfg.get("peers", {}).keys())
    peer_index = sorted_peers.index(peer_name) + 2

    if access_class and access_class in cfg.get("access_classes", {}):
        class_cfg = cfg["access_classes"][access_class]
        subnet = class_cfg.get("subnet", "10.137.0.0/24")
        listen_port = class_cfg.get("listen_port", 51820)
    else:
        subnet = iface.get("addresses", ["10.137.0.1/24"])[0]
        listen_port = iface.get("listen_port", 51820)

    if "/" not in subnet:
        subnet = f"{subnet}/24"
    addr = subnet.rsplit(".", 1)[0]
    prefix = subnet.rsplit("/", 1)[1]
    client_addr = f"{addr}.{peer_index}/{prefix}"

    sk = server_pubkey or iface.get("public_key", "")
    ep = server_endpoint or iface.get("server_endpoint", "")
    if ep and listen_port:
        host = ep.split(":")[0]
        ep = f"{host}:{listen_port}"

    tmpl = ENV.get_template("wireguard-client.conf")
    conf = tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        peer_name=peer_name,
        client_priv=client_priv,
        client_addr=client_addr,
        server_pubkey=sk,
        server_endpoint=ep,
        allowed_ips=peer.get("allowed_ips", ["0.0.0.0/0"]),
        preshared_key=peer.get("preshared_key"),
        persistent_keepalive=peer.get("persistent_keepalive"),
    )
    logger.info("Client config generated for peer '%s'", peer_name)
    return conf


# --- Interface-level setters ---


def set_listen_port(port: int) -> None:
    """Update the server listen port in the stored configuration."""
    if not (1 <= port <= 65535):
        raise ValueError("Listen port must be in range 1..65535")
    cfg = get_config()
    cfg["interface"]["listen_port"] = port
    save_config(cfg)
    logger.info("WireGuard listen port set to %d", port)


def set_post_up(cmd: str | None) -> None:
    """Set (or clear) the PostUp hook command."""
    cfg = get_config()
    cfg["interface"]["post_up"] = cmd
    save_config(cfg)


def set_post_down(cmd: str | None) -> None:
    """Set (or clear) the PostDown hook command."""
    cfg = get_config()
    cfg["interface"]["post_down"] = cmd
    save_config(cfg)


# --- Class key generation ---


def generate_class_keypair(class_key: str) -> tuple[str, str]:
    """Generate a key pair for an access class interface."""
    cfg = get_config()
    classes = cfg.setdefault("access_classes", {})
    if class_key not in classes:
        raise ValueError(f"Access class '{class_key}' not found")
    class_cfg = classes[class_key]
    if class_cfg.get("private_key"):
        return class_cfg["private_key"], class_cfg["public_key"]
    priv, pub = generate_keypair()
    class_cfg["private_key"] = priv
    class_cfg["public_key"] = pub
    save_config(cfg)
    logger.info("Key pair generated for class '%s'", class_key)
    return priv, pub


# --- Initialise ---


def _ensure_class_defaults(cfg: dict[str, Any]) -> None:
    """Ensure access classes have required Phase-2 fields."""
    classes = cfg.setdefault("access_classes", {})
    for _key, c in classes.items():
        if not isinstance(c, dict):
            continue
        for field, default in _default_class_fields().items():
            if field not in c:
                c[field] = default


def _ensure_access_classes(cfg: dict[str, Any]) -> None:
    """Pre-seed default access classes if missing or empty (idempotent).

    Also upgrades existing classes with Phase-2 fields.
    """
    classes = cfg.setdefault("access_classes", {})
    if not classes:
        full_defaults = deepcopy(DEFAULT_CONFIG["access_classes"])
        classes.update(full_defaults)
        return
    _ensure_class_defaults(cfg)


def initialize() -> dict[str, Any]:
    """Perform first-time WireGuard setup."""
    cfg = get_config()

    if cfg["interface"].get("private_key"):
        _ensure_access_classes(cfg)
        save_config(cfg)
        return cfg

    priv, pub = generate_keypair()
    cfg["interface"]["private_key"] = priv
    cfg["interface"]["public_key"] = pub
    _ensure_access_classes(cfg)
    save_config(cfg)
    logger.info("WireGuard initialised (pubkey=%s...)", pub[:16])
    return cfg


__all__ = [
    "DEFAULT_CONFIG",
    "add_peer",
    "apply",
    "apply_class",
    "down",
    "down_class",
    "generate_class_conf",
    "generate_class_keypair",
    "generate_client_conf",
    "generate_conf",
    "generate_keypair",
    "get_class_interface_name",
    "get_class_zone_name",
    "get_config",
    "get_peer_status",
    "get_peers",
    "initialize",
    "parse_wg_show_output",
    "remove_peer",
    "save_config",
    "set_listen_port",
    "set_post_down",
    "set_post_up",
    "status",
]