win-dhcpd/dhcpd/lease.py

"""Lease database: in-memory index + JSON persistence."""

from __future__ import annotations

import ipaddress
import json
import logging
import os
import time
from dataclasses import dataclass, asdict, field
from datetime import datetime, timezone
from typing import Dict, List, Optional

log = logging.getLogger("dhcpd.lease")


@dataclass
class Lease:
    mac: str            # "aa:bb:cc:dd:ee:ff"
    ip: str
    hostname: str = ""
    offered_at: float = 0.0
    acknowledged_at: float = 0.0
    expires_at: float = 0.0
    static: bool = False   # True = from static binding
    xid: int = 0
    scope: str = ""

    @property
    def is_expired(self) -> bool:
        return time.time() > self.expires_at

    @property
    def is_active(self) -> bool:
        return self.expires_at > time.time()

    def to_json(self) -> dict:
        return {
            "mac": self.mac,
            "ip": self.ip,
            "hostname": self.hostname,
            "offered_at": self.offered_at,
            "acknowledged_at": self.acknowledged_at,
            "expires_at": self.expires_at,
            "static": self.static,
            "scope": self.scope,
        }


class LeaseDB:
    """Thread-ish safe lease store. (Single-threaded server, no lock needed.)"""

    def __init__(self, db_path: str = "leases.db"):
        self.db_path = db_path
        # mac -> Lease
        self.by_mac: Dict[str, Lease] = {}
        # ip  -> mac
        self.by_ip: Dict[str, str] = {}
        self._load()

    # ── persistence ──────────────────────────────────────────────────────

    def _load(self) -> None:
        if not os.path.exists(self.db_path):
            return
        try:
            with open(self.db_path, "r", encoding="utf-8") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            log.warning("Could not load lease DB: %s", e)
            return
        if isinstance(data, list):
            for rec in data:
                try:
                    lease = Lease(**rec)
                    if lease.is_active or lease.static:
                        self._put(lease)
                except TypeError:
                    continue
            log.info("Loaded %d leases from %s", len(self.by_mac), self.db_path)

    def save(self) -> None:
        try:
            records = [l.to_json() for l in self.by_mac.values()]
            with open(self.db_path, "w", encoding="utf-8") as f:
                json.dump(records, f, indent=2)
        except OSError as e:
            log.error("Failed to save lease DB: %s", e)

    # ── CRUD ─────────────────────────────────────────────────────────────

    def _put(self, lease: Lease) -> None:
        # remove stale entry if any
        if lease.mac in self.by_mac:
            old = self.by_mac[lease.mac]
            if old.ip != lease.ip:
                self.by_ip.pop(old.ip, None)
        if lease.ip in self.by_ip:
            del self.by_mac[self.by_ip[lease.ip]]
        self.by_mac[lease.mac] = lease
        self.by_ip[lease.ip] = lease.mac

    def find_by_mac(self, mac: str) -> Optional[Lease]:
        return self.by_mac.get(mac)

    def find_by_ip(self, ip: str) -> Optional[Lease]:
        mac = self.by_ip.get(ip)
        if mac:
            return self.by_mac.get(mac)
        return None

    def is_ip_taken(self, ip: str) -> bool:
        return ip in self.by_ip

    def add(self, lease: Lease) -> Lease:
        self._put(lease)
        self.save()
        return lease

    def release(self, mac: str) -> bool:
        if mac in self.by_mac:
            lease = self.by_mac[mac]
            self.by_ip.pop(lease.ip, None)
            del self.by_mac[mac]
            self.save()
            log.info("Released lease: %s -> %s", mac, lease.ip)
            return True
        return False

    def extend(self, mac: str, new_expires: float) -> Optional[Lease]:
        lease = self.by_mac.get(mac)
        if lease:
            lease.expires_at = new_expires
            self.save()
            return lease
        return None

    # ── allocation ───────────────────────────────────────────────────────

    def allocate_next(
        self,
        start: str,
        end: str,
        mac: str,
        lease_seconds: int,
        xid: int = 0,
        hostname: str = "",
        scope: str = "",
    ) -> Optional[Lease]:
        """Walk from start_ip to end_ip, return the first free address."""
        now = time.time()
        a_start = ipaddress.ip_address(start)
        a_end = ipaddress.ip_address(end)
        cur = a_start
        while cur <= a_end:
            cand = str(cur)
            if not self.is_ip_taken(cand):
                lease = Lease(
                    mac=mac,
                    ip=cand,
                    hostname=hostname,
                    offered_at=now,
                    expires_at=now + lease_seconds,
                    xid=xid,
                    scope=scope,
                )
                self.add(lease)
                return lease
            cur = cur + 1
        return None

    def list_all(self) -> List[Lease]:
        return sorted(self.by_mac.values(), key=lambda l: l.ip)

    def count(self) -> int:
        return len(self.by_mac)

    def count_active(self) -> int:
        return sum(1 for l in self.by_mac.values() if l.is_active)

    def prune_expired(self) -> int:
        """Remove expired non-static leases. Returns count removed."""
        to_remove = [
            mac for mac, l in self.by_mac.items()
            if not l.static and l.is_expired
        ]
        for mac in to_remove:
            self.release(mac)
        return len(to_remove)