vacuum-wall/lib/dnsmasq.py

"""
dnsmasq configuration manager for Vacuum Wall SSL proxy firewall.

Generates /etc/dnsmasq.d/vacuum-wall.conf and manages DHCP range,
static leases, and custom DNS records through sudo.
"""

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

from jinja2 import Environment, FileSystemLoader

PROJECT_DIR = Path(__file__).resolve().parent.parent
CONFIG_DIR = PROJECT_DIR / "config" / "dnsmasq"
DATA_DIR = PROJECT_DIR / "data" / "dnsmasq"
CONFIG_PATH = CONFIG_DIR / "config.json"
FRAGMENTS_DIR = DATA_DIR / "fragments"
DNSMASQ_CONF = "/etc/dnsmasq.d/vacuum-wall.conf"
LEASE_FILE = "/var/lib/dnsmasq/dnsmasq.leases"

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

# --- defaults ---
DEFAULT_CFG: dict[str, Any] = {
    "dhcp": {
        "ranges": [],
        "static_leases": [],
    },
    "dns": {
        "upstreams": ["8.8.8.8", "1.1.1.1"],
        "domain": None,
        "custom_records": [],
    },
}

# ───────── helpers ───────────────────────────────────────────────────


def _ensure_dirs() -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    FRAGMENTS_DIR.mkdir(parents=True, exist_ok=True)


def _sudo(*cmd: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        ["sudo", *list(cmd)],
        capture_output=True,
        text=True,
        check=True,
    )


def _load_json(path: Path) -> dict:
    if not path.exists():
        return {}
    with open(path) as f:
        return json.load(f)


def _save_json(path: Path, data: dict) -> None:
    _ensure_dirs()
    with open(path, "w") as f:
        json.dump(data, f, indent=4)


def _deep_merge(base: dict, overrides: dict) -> dict:
    result = deepcopy(base)
    for k, v in overrides.items():
        if k in result and isinstance(result[k], dict) and isinstance(v, dict):
            result[k] = _deep_merge(result[k], v)
        else:
            result[k] = deepcopy(v)
    return result


# ───────── config lifecycle ──────────────────────────────────────────


def get_config() -> dict:
    """Load current dnsmasq config from JSON state file."""
    _ensure_dirs()
    raw = _load_json(CONFIG_PATH)
    if not raw:
        return deepcopy(DEFAULT_CFG)
    return _deep_merge(deepcopy(DEFAULT_CFG), raw)


def save_config(cfg: dict) -> None:
    """Persist config to JSON (does NOT touch on-disk dnsmasq config)."""
    _ensure_dirs()
    merged = _deep_merge(deepcopy(DEFAULT_CFG), cfg)
    _save_json(CONFIG_PATH, merged)


def apply_config() -> None:
    """Write generated config to disk via sudo tee, then reload dnsmasq."""
    cfg = get_config()
    conf_text = generate_conf(cfg)

    _ensure_dirs()
    _sudo("mkdir", "-p", "/etc/dnsmasq.d")
    subprocess.run(
        ["sudo", "tee", DNSMASQ_CONF, "--"],
        input=conf_text,
        capture_output=True,
        text=True,
        check=True,
    )
    _sudo("systemctl", "reload", "dnsmasq")


# ───────── config generation ─────────────────────────────────────────


def generate_conf(cfg: dict) -> str:
    """Render a complete dnsmasq.conf text block from the config dict."""
    dhcp_cfg = cfg.get("dhcp", {})
    dns_cfg = cfg.get("dns", {})

    interfaces = [
        r["interface"] for r in dhcp_cfg.get("ranges", []) if "interface" in r
    ]

    tmpl = ENV.get_template("dnsmasq.conf")
    return tmpl.render(
        timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
        interfaces=interfaces,
        dhcp=dhcp_cfg,
        dns=dns_cfg,
        fragments_dir=str(FRAGMENTS_DIR) if FRAGMENTS_DIR.exists() else None,
    )


# ───────── dhcp management ───────────────────────────────────────────


def set_dhcp_range(
    iface: str,
    start: str,
    end: str,
    lease_time: str = "12h",
    gateway: str | None = None,
    dns: str | None = None,
) -> None:
    """Add or replace the DHCP range for a given interface."""
    cfg = get_config()
    ranges = cfg["dhcp"]["ranges"]

    found = False
    for i, r in enumerate(ranges):
        if r.get("interface") == iface:
            ranges[i] = {
                "interface": iface,
                "start": start,
                "end": end,
                "lease_time": lease_time,
            }
            if gateway:
                ranges[i]["gateway"] = gateway
            if dns:
                ranges[i]["dns"] = dns
            found = True
            break

    if not found:
        entry: dict[str, Any] = {
            "interface": iface,
            "start": start,
            "end": end,
            "lease_time": lease_time,
        }
        if gateway:
            entry["gateway"] = gateway
        if dns:
            entry["dns"] = dns
        ranges.append(entry)

    save_config(cfg)


def add_static_lease(mac: str, ip: str, hostname: str | None = None) -> None:
    """Add (or update) a static DHCP lease by MAC address."""
    cfg = get_config()
    leases = cfg["dhcp"]["static_leases"]

    for i, lease in enumerate(leases):
        if lease["mac"].lower() == mac.lower():
            leases[i] = {"mac": mac, "ip": ip}
            if hostname:
                leases[i]["hostname"] = hostname
            save_config(cfg)
            return

    entry: dict[str, Any] = {"mac": mac, "ip": ip}
    if hostname:
        entry["hostname"] = hostname
    leases.append(entry)
    save_config(cfg)


def remove_static_lease(mac: str) -> None:
    """Remove a static DHCP lease by MAC address."""
    cfg = get_config()
    cfg["dhcp"]["static_leases"] = [
        lease
        for lease in cfg["dhcp"]["static_leases"]
        if lease["mac"].lower() != mac.lower()
    ]
    save_config(cfg)


# ───────── dns record management ─────────────────────────────────────


def add_dns_record(name: str, address: str, hostname: str | None = None) -> None:
    """Add or update a custom DNS A record."""
    cfg = get_config()
    records = cfg["dns"]["custom_records"]

    for i, r in enumerate(records):
        if r["name"] == name:
            records[i] = {"name": name, "address": address}
            if hostname:
                records[i]["hostname"] = hostname
            save_config(cfg)
            return

    entry: dict[str, Any] = {"name": name, "address": address}
    if hostname:
        entry["hostname"] = hostname
    records.append(entry)
    save_config(cfg)


def remove_dns_record(name: str) -> None:
    """Remove a custom DNS record by name."""
    cfg = get_config()
    cfg["dns"]["custom_records"] = [
        r for r in cfg["dns"]["custom_records"] if r["name"] != name
    ]
    save_config(cfg)


# ───────── lease table ───────────────────────────────────────────────


def _parse_lease_line(line: str) -> dict[str, Any] | None:
    """Parse one line from dnsmasq.leases into a dict."""
    line = line.strip()
    if not line or line.startswith("#"):
        return None
    parts = line.split()
    if len(parts) < 3:
        return None

    try:
        ts = datetime.fromtimestamp(int(parts[0]), tz=UTC)
    except (ValueError, OSError):
        ts = None

    return {
        "expires_at": ts,
        "mac": parts[1],
        "ip": parts[2],
        "hostname": parts[3] if len(parts) > 3 else "",
        "interface": parts[4] if len(parts) > 4 else "",
    }


def get_lease_table() -> list[dict]:
    """Read and parse the current dnsmasq lease file."""
    leases: list[dict] = []
    try:
        result = _sudo("cat", LEASE_FILE)
        for entry in map(_parse_lease_line, result.stdout.splitlines()):
            if entry is not None:
                leases.append(entry)
    except subprocess.CalledProcessError:
        pass
    return leases


# ───────── upstream / domain helpers ─────────────────────────────────


def set_upstreams(servers: list[str]) -> None:
    """Set the list of upstream DNS forwarders."""
    cfg = get_config()
    cfg["dns"]["upstreams"] = list(servers)
    save_config(cfg)


def set_domain(domain: str | None) -> None:
    """Set (or clear) the local DNS domain."""
    cfg = get_config()
    cfg["dns"]["domain"] = domain if domain else None
    save_config(cfg)


# ───────── status / info ─────────────────────────────────────────────


def get_status() -> dict:
    """Return service status, config summary, and current lease count."""
    cfg = get_config()

    # dnsmasq process check
    try:
        proc = subprocess.run(
            ["sudo", "systemctl", "is-active", "dnsmasq"],
            capture_output=True,
            text=True,
        )
        active = proc.stdout.strip() == "active"
    except Exception:
        active = False

    # config on disk
    conf_exists = os.path.isfile(DNSMASQ_CONF)
    if conf_exists:
        try:
            with open(DNSMASQ_CONF) as f:
                conf_on_disk = f.read()
        except PermissionError:
            conf_on_disk = ""
    else:
        conf_on_disk = ""

    # current expected config
    expected = generate_conf(cfg)

    leases = get_lease_table()

    return {
        "service_active": active,
        "config_file_exists": conf_exists,
        "config_in_sync": conf_on_disk == expected,
        "dhcp_ranges": len(cfg["dhcp"]["ranges"]),
        "static_leases": len(cfg["dhcp"]["static_leases"]),
        "custom_dns_records": len(cfg["dns"]["custom_records"]),
        "upstreams": cfg["dns"]["upstreams"],
        "domain": cfg["dns"].get("domain"),
        "active_leases": len(leases),
        "leases": leases,
    }