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

import json
import os
import subprocess
from datetime import UTC, datetime
from pathlib import Path

from jinja2 import Environment, FileSystemLoader

PROJECT_DIR = Path(__file__).resolve().parent.parent
CONFIG_PATH = str(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,
)


# --- Helpers ---


def _run(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
    """Run a command via sudo and return the completed process."""
    return subprocess.run(
        ["sudo", *cmd],
        capture_output=True,
        text=True,
        check=check,
    )


def _ensure_dir(path: str) -> None:
    """Create parent directories for *path* if they don't exist."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)


def _default_config() -> dict:
    """Return the skeleton config with no keys and no peers."""
    return {
        "interface": {
            "name": "wg0",
            "listen_port": 51820,
            "private_key": "",
            "public_key": "",
            "addresses": ["10.137.0.1/24"],
            "post_up": None,
            "post_down": None,
        },
        "peers": {},
    }


# --- Core config persistence ---


def get_config() -> dict:
    """Load the current WireGuard configuration from the JSON store.

    Returns the full config dict.  If the file does not exist or is
    unreadable, returns the default (empty) config skeleton.
    """
    try:
        with open(CONFIG_PATH) as f:
            cfg = json.load(f)
        # Backfill keys that might be missing from older snapshots.
        defaults = _default_config()
        cfg.setdefault("interface", defaults["interface"])
        cfg["interface"].setdefault("name", defaults["interface"]["name"])
        cfg["interface"].setdefault("listen_port", defaults["interface"]["listen_port"])
        cfg["interface"].setdefault("private_key", defaults["interface"]["private_key"])
        cfg["interface"].setdefault("public_key", defaults["interface"]["public_key"])
        cfg["interface"].setdefault("addresses", defaults["interface"]["addresses"])
        cfg["interface"].setdefault("post_up", defaults["interface"]["post_up"])
        cfg["interface"].setdefault("post_down", defaults["interface"]["post_down"])
        cfg.setdefault("peers", {})
        return cfg
    except (FileNotFoundError, json.JSONDecodeError):
        return _default_config()


def save_config(cfg: dict) -> None:
    """Persist *cfg* to the JSON store atomically.

    Writes to a temporary file in the same directory and then renames
    to avoid partial reads on crash.
    """
    _ensure_dir(CONFIG_PATH)
    tmp = CONFIG_PATH + ".tmp"
    with open(tmp, "w") as f:
        json.dump(cfg, f, indent=4)
        f.write("\n")
    os.replace(tmp, CONFIG_PATH)


# --- Key generation ---


def generate_keypair() -> tuple[str, str]:
    """Generate a WireGuard private/public key pair using ``wg`` CLI.

    Returns:
        ``(private_key, public_key)`` as two 43-character base64 strings.
    """
    res = _run([WG_BIN, "genkey"])
    private_key = res.stdout.strip()
    res2 = _run([WG_BIN, "pubkey"], input=private_key)
    public_key = res2.stdout.strip()
    return private_key, public_key


# --- wg0.conf generation ---


def generate_conf(cfg: dict) -> str:
    """Render a valid wg-quick config file from *cfg* using Jinja2."""
    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", {}),
    )


# --- Apply / down ---


def apply() -> None:
    """Write the current config to disk and bring the tunnel up with wg-quick."""
    cfg = get_config()
    conf_text = generate_conf(cfg)
    save_config(cfg)  # ensure latest state persisted

    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])
    _run(["chown", "root:root", WG_CONF_PATH], check=False)
    local_tmp.unlink(missing_ok=True)

    _run([WG_QUICK_BIN, "up", cfg["interface"]["name"]])


def down() -> None:
    """Bring the WireGuard tunnel interface down."""
    cfg = get_config()
    name = cfg["interface"]["name"]
    _run([WG_QUICK_BIN, "down", name])


# --- Status ---


def status() -> dict:
    """Query the live tunnel state via ``wg show``.

    Returns a dict with keys:
      - ``up`` (bool) - whether the interface is currently up.
      - ``interface`` (dict) - name, public key, listen port, fwmark.
      - ``peers`` (list[dict]) - per-peer status from ``wg show wg0``.
    """
    cfg = get_config()
    name = cfg["interface"]["name"]
    result = {
        "up": False,
        "interface": {},
        "peers": [],
    }

    try:
        proc = _run([WG_BIN, "show", name], check=False)
        if proc.returncode != 0:
            return result

        raw = proc.stdout.strip()
    except Exception:
        return result

    # Parse the wg show output.
    # Format (multi-section, separated by blank lines or interleaved):
    #   interface:
    #     public key: ...
    #     listening port: ...
    #   peer: <key>
    #     endpoint: ...
    #     allowed ips: ...
    #     latest handshake: ...
    #     transfer: ...
    #     persistent-keepalive: ...
    current_peer = None
    peers: list[dict] = []

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


def add_peer(
    name: str,
    endpoint: str | None = None,
    allowed_ips: list[str] | None = None,
    persistent_keepalive: int | None = None,
    preshared_key: str | None = None,
) -> dict:
    """Add (or update) a peer in the configuration.

    If the peer has no public key yet, one will be generated
    together with a matching private key (useful for client provi-
    sioning).  The returned dict mirrors the stored peer record
    with an additional ``private_key`` field so the caller can
    distribute the client credentials.

    Args:
        name: Human-readable identifier (dict key in config).
        endpoint: e.g. ``203.0.113.1:51820``.
        allowed_ips: CIDR list, e.g. ``["0.0.0.0/0"]``.
        persistent_keepalive: Interval in seconds (or ``None``).
        preshared_key: Optional PSK (base64 string).

    Returns:
        The peer dict as stored, plus ``private_key`` for client use.
    """
    cfg = get_config()
    peers = cfg.setdefault("peers", {})
    allowed_ips = allowed_ips or []

    if name in peers:
        peer = peers[name]
        peer["endpoint"] = endpoint
        peer["allowed_ips"] = allowed_ips
        peer["persistent_keepalive"] = persistent_keepalive
        if preshared_key is not None:
            peer["preshared_key"] = preshared_key
    else:
        # Generate a key pair for the new peer.
        priv, pub = generate_keypair()
        peer = {
            "public_key": pub,
            "private_key": priv,  # stored so we can hand it to the client
            "endpoint": endpoint,
            "allowed_ips": allowed_ips,
            "persistent_keepalive": persistent_keepalive,
            "preshared_key": preshared_key,
        }
        peers[name] = peer

    save_config(cfg)

    # Return a copy that includes the private key (safe — used for provisioning).
    peer_out = dict(peer)
    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)


def get_peers() -> list[dict]:
    """List all configured peers (from the JSON store, *not* live).

    Returns a list of dicts.  Each dict includes ``name`` and all
    stored fields **except** ``private_key`` (not exposed here).
    """
    cfg = get_config()
    peers = []
    for name, info in cfg.get("peers", {}).items():
        entry = dict(info)
        entry["name"] = name
        # Strip private key from the public listing.
        entry.pop("private_key", None)
        peers.append(entry)
    return peers


def get_peer_status() -> list[dict]:
    """Return live peer status from ``wg show``.

    Each element contains:
      - ``public_key``, ``endpoint``, ``allowed_ips``,
        ``latest_handshake``, ``transfer_received``,
        ``transfer_sent``, ``persistent_keepalive``.
    """
    st = status()
    return st.get("peers", [])


# --- Client config generation ---


def generate_client_conf(
    peer_name: str,
    server_endpoint: str,
    server_pubkey: str,
) -> str:
    """Build a client-side wg-quick config snippet for *peer_name*."""
    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."
        )

    sorted_peers = sorted(cfg.get("peers", {}).keys())
    peer_index = sorted_peers.index(peer_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")
    return 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=server_pubkey,
        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"),
    )


# --- Interface-level setters ---


def set_listen_port(port: int) -> None:
    """Update the server listen port in the stored configuration.

    Does **not** hot-reload; call :func:`apply` afterwards to
    activate the change.
    """
    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)


def set_post_up(cmd: str | None) -> None:
    """Set (or clear) the PostUp hook command.

    The command is passed verbatim to the generated wg0.conf.
    """
    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)


# --- Initialise ---


def initialize() -> dict:
    """Perform first-time WireGuard setup.

    Generates a fresh server key pair, writes the initial config
    to disk, and returns the full config dict.

    Call this once at appliance bootstrapping time.  It will
    **not** overwrite an existing config that already has a
    non-empty private key.
    """
    cfg = get_config()

    if cfg["interface"].get("private_key"):
        # Already initialised — return existing config.
        return cfg

    priv, pub = generate_keypair()
    cfg["interface"]["private_key"] = priv
    cfg["interface"]["public_key"] = pub
    save_config(cfg)
    return cfg


# --- Utility: parse wg show into structured peer map ---


def _parse_wg_show(output: str) -> dict:
    """Internal parser for ``wg show`` multiline output.

    Returns a dict keyed by peer public key with parsed values.
    Used internally; ``status()`` is the public interface.
    """
    peers: dict = {}
    current = None

    for line in output.splitlines():
        line = line.strip()
        if line.startswith("peer:"):
            key = line.split(":", 1)[1].strip()
            current = {"_key": key}
            peers[key] = current
            continue

        if current is None:
            continue

        if line.startswith("endpoint:"):
            val = line.split(":", 1)[1].strip()
            current["endpoint"] = val
        elif line.startswith("allowed ips:"):
            current["allowed_ips"] = line.split(":", 1)[1].strip()
        elif line.startswith("latest handshake:"):
            current["latest_handshake"] = line.split(":", 1)[1].strip()
        elif line.startswith("transfer:"):
            current["transfer_raw"] = line.split(":", 1)[1].strip()
        elif line.startswith("persistent-keepalive:"):
            current["persistent_keepalive"] = line.split(":", 1)[1].strip()

    return peers