vacuum-wall/daemon/handlers/dnsmasq.py

"""Dnsmasq daemon handler."""

import logging
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, ensure_dirs, load_json, run, run_proc, save_json

logger = logging.getLogger(__name__)

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

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


def _get_state() -> dict[str, Any] | None:
    from lib.state import state as state_store

    return state_store.get("dnsmasq")


def _get_config() -> dict[str, Any]:
    ensure_dirs(CONFIG_DIR, DATA_DIR, FRAGMENTS_DIR)
    raw = load_json(CONFIG_PATH)
    if not raw:
        return deepcopy(DEFAULT_CFG)
    return deep_merge(deepcopy(DEFAULT_CFG), raw)


def _save_config(cfg: dict[str, Any]) -> None:
    ensure_dirs(CONFIG_DIR, DATA_DIR, FRAGMENTS_DIR)
    merged = deep_merge(deepcopy(DEFAULT_CFG), cfg)
    save_json(CONFIG_PATH, merged)


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


def _get_dnsmasq_state() -> dict[str, Any]:
    dm = _get_state()
    if dm is None:
        return {}
    return dm


# ---------------------------------------------------------------------------
# Routes


@registry.register("GET", "/dnsmasq/config")
def get_config(_request: Any, _body: Any) -> dict[str, Any]:
    dm = _get_dnsmasq_state()
    if dm:
        return dm.get("config", {})
    return _get_config()


@registry.register("POST", "/dnsmasq/config")
def save_config_handler(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    _save_config(body)
    refresh_state(["dnsmasq"])
    return {"config_saved": True}


@registry.register("PATCH", "/dnsmasq/config")
def patch_config(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    current = _get_config()
    merged = deep_merge(current, body)
    _save_config(merged)
    refresh_state(["dnsmasq"])
    return {"config_saved": True}


@registry.register("POST", "/dnsmasq/apply")
def apply_config(_request: Any, _body: Any) -> dict[str, Any]:
    cfg = _get_config()
    conf_text = _generate_conf(cfg)
    ensure_dirs(CONFIG_DIR, DATA_DIR, FRAGMENTS_DIR)
    run(["mkdir", "-p", "/etc/dnsmasq.d"], sudo=True)
    run_proc(
        ["tee", DNSMASQ_CONF, "--"],
        sudo=True,
        check=True,
        input=conf_text,
    )
    run(["systemctl", "reload", "dnsmasq"], sudo=True)
    logger.info("dnsmasq config written and reloaded")
    refresh_state(["dnsmasq"])
    return {"applied": True}


@registry.register("GET", "/dnsmasq/status")
def get_status(_request: Any, _body: Any) -> dict[str, Any]:
    dm = _get_dnsmasq_state()
    if dm and "status" in dm:
        return dm["status"]
    return {}


@registry.register("POST", "/dnsmasq/ranges/add")
def set_dhcp_range(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    iface = body.get("interface", "").strip() or ""
    start = body.get("start", "").strip()
    end = body.get("end", "").strip()
    lease_time = body.get("lease_time", "12h")
    if not start or not end:
        raise ValueError("'start' and 'end' are required")
    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 body.get("gateway"):
                ranges[i]["gateway"] = body["gateway"]
            if body.get("dns"):
                ranges[i]["dns"] = body["dns"]
            found = True
            break
    if not found:
        entry: dict[str, Any] = {
            "interface": iface,
            "start": start,
            "end": end,
            "lease_time": lease_time,
        }
        if body.get("gateway"):
            entry["gateway"] = body["gateway"]
        if body.get("dns"):
            entry["dns"] = body["dns"]
        ranges.append(entry)
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"interface": iface, "start": start, "end": end}


@registry.register("DELETE", "/dnsmasq/ranges/remove")
def remove_dhcp_range(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    iface = body.get("interface", "").strip() or ""
    start = body.get("start", "").strip()
    end = body.get("end", "").strip()
    if not start or not end:
        raise ValueError("'start' and 'end' are required")
    cfg = _get_config()
    ranges = cfg["dhcp"]["ranges"]
    before = len(ranges)
    cfg["dhcp"]["ranges"] = [
        r
        for r in ranges
        if not (
            r.get("interface") == iface
            and r.get("start") == start
            and r.get("end") == end
        )
    ]
    if len(cfg["dhcp"]["ranges"]) == before:
        raise NotFoundError(
            f"DHCP range for interface '{iface}' ({start}-{end}) not found"
        )
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"interface": iface, "start": start, "end": end}


@registry.register("GET", "/dnsmasq/leases")
def get_leases(_request: Any, _body: Any) -> list[dict[str, Any]]:
    dm = _get_dnsmasq_state()
    if dm:
        return dm.get("leases", [])
    return []


@registry.register("POST", "/dnsmasq/static-lease/add")
def add_static_lease(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    mac = body.get("mac", "").strip()
    ip = body.get("ip", "").strip()
    hostname = body.get("hostname")
    if not mac or not ip:
        raise ValueError("'mac' and 'ip' are required")
    cfg = _get_config()
    leases = cfg["dhcp"]["static_leases"]
    for i, lease in enumerate(leases):
        if lease["mac"].lower() == mac.lower():
            leases[i].update({"mac": mac, "ip": ip})
            if hostname is not None:
                leases[i]["hostname"] = hostname
            _save_config(cfg)
            refresh_state(["dnsmasq"])
            return {"mac": mac, "ip": ip, "hostname": hostname}
    entry: dict[str, Any] = {"mac": mac, "ip": ip}
    if hostname:
        entry["hostname"] = hostname
    leases.append(entry)
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"mac": mac, "ip": ip, "hostname": hostname}


@registry.register("DELETE", "/dnsmasq/static-lease/remove")
def remove_static_lease(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    mac = body.get("mac", "").strip()
    if not mac:
        raise ValueError("'mac' is required")
    cfg = _get_config()
    leases = cfg["dhcp"]["static_leases"]
    before = len(leases)
    cfg["dhcp"]["static_leases"] = [
        lease for lease in leases if lease["mac"].lower() != mac.lower()
    ]
    if len(cfg["dhcp"]["static_leases"]) == before:
        raise NotFoundError(f"Static lease for MAC '{mac}' not found")
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"mac": mac}


@registry.register("POST", "/dnsmasq/dns-record/add")
def add_dns_record(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    name = body.get("name", "").strip()
    address = body.get("address", "").strip()
    hostname = body.get("hostname")
    if not name or not address:
        raise ValueError("'name' and 'address' are required")
    cfg = _get_config()
    records = cfg["dns"]["custom_records"]
    for i, r in enumerate(records):
        if r["name"] == name:
            records[i].update({"name": name, "address": address})
            if hostname is not None:
                records[i]["hostname"] = hostname
            _save_config(cfg)
            refresh_state(["dnsmasq"])
            return {"name": name, "address": address, "hostname": hostname}
    entry: dict[str, Any] = {"name": name, "address": address}
    if hostname:
        entry["hostname"] = hostname
    records.append(entry)
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"name": name, "address": address, "hostname": hostname}


@registry.register("DELETE", "/dnsmasq/dns-record/remove")
def remove_dns_record(_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()
    records = cfg["dns"]["custom_records"]
    before = len(records)
    cfg["dns"]["custom_records"] = [r for r in records if r["name"] != name]
    if len(cfg["dns"]["custom_records"]) == before:
        raise NotFoundError(f"DNS record '{name}' not found")
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"name": name}


@registry.register("POST", "/dnsmasq/upstreams")
def set_upstreams(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body or "servers" not in body:
        raise ValueError("'servers' is required")
    cfg = _get_config()
    cfg["dns"]["upstreams"] = list(body["servers"])
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"upstreams": cfg["dns"]["upstreams"]}


@registry.register("POST", "/dnsmasq/domain")
def set_domain(_request: Any, body: dict[str, Any] | None) -> dict[str, Any]:
    if not body:
        raise ValueError("Request body required")
    domain = body.get("domain")
    cfg = _get_config()
    cfg["dns"]["domain"] = domain if domain else None
    _save_config(cfg)
    refresh_state(["dnsmasq"])
    return {"domain": cfg["dns"]["domain"]}