zoraxy-dhcp/server.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ---------------------------------------------------------------------------
// Configuration (overridable for testing / custom installs)
// ---------------------------------------------------------------------------

var (
	// Default paths — all under Zoraxy's unified config directory.
	// Override via environment: LEASE_FILE, CONF_FILE, RELOAD_CMD, RELOAD_ARGS
	leasesFile  = envOrDefault("LEASE_FILE", "/opt/zoraxy/conf/dhcp/dnsmasq.leases")
	confFile    = envOrDefault("CONF_FILE", "/opt/zoraxy/conf/dhcp/dnsmasq.conf")
)

func envOrDefault(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

// ---------------------------------------------------------------------------
// dnsmasq interaction
// ---------------------------------------------------------------------------

// parseLeases reads the dnsmasq lease file (default: /opt/zoraxy/conf/dhcp/dnsmasq.leases).
func parseLeases() ([]Lease, error) {
	f, err := os.Open(leasesFile)
	if err != nil {
		return nil, fmt.Errorf("open leases file: %w", err)
	}
	defer f.Close()

	var leases []Lease
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 3 {
			continue
		}
		ts, err := strconv.ParseInt(parts[0], 10, 64)
		if err != nil {
			continue
		}
		lease := Lease{
			Expiry:   time.Unix(ts, 0),
			MAC:      strings.ToLower(parts[1]),
			IP:       parts[2],
		}
		if len(parts) >= 4 {
			lease.Hostname = parts[3]
		}
		if len(parts) >= 5 {
			lease.ClientID = parts[4]
		}
		leases = append(leases, lease)
	}
	return leases, scanner.Err()
}

// dhcpHostPattern matches lines like: dhcp-host=mac,ip,hostname  or dhcp-host=mac,hostname
var dhcpHostPattern = regexp.MustCompile(`^dhcp-host\s*=\s*(.+)$`)

// parsePinnedHosts reads dnsmasq.conf and returns a set of pinned MAC addresses.
func parsePinnedHosts() (map[string]string, error) {
	f, err := os.Open(confFile)
	if err != nil {
		return nil, fmt.Errorf("open config file: %w", err)
	}
	defer f.Close()

	pinned := make(map[string]string)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		matches := dhcpHostPattern.FindStringSubmatch(line)
		if matches == nil {
			continue
		}
		fields := strings.Split(matches[1], ",")
		if len(fields) < 2 {
			continue
		}
		mac := strings.TrimSpace(strings.ToLower(fields[0]))
		hostname := strings.TrimSpace(fields[len(fields)-1])
		pinned[mac] = hostname
	}
	return pinned, scanner.Err()
}

// sanitize strips control characters and newlines from a config value.
func sanitize(s string) string {
	s = strings.ReplaceAll(s, "\n", "")
	s = strings.ReplaceAll(s, "\r", "")
	return strings.TrimSpace(strings.Map(func(r rune) rune {
		if r < 32 && r != '\t' {
			return -1
		}
		return r
	}, s))
}

// pinLease appends a dhcp-host line to dnsmasq.conf. If hostname is provided,
// also appends an address=/hostname/IP line so the name resolves network-wide via dnsmasq DNS.
func pinLease(mac, ip, hostname string) error {
	safeMAC := sanitize(mac)
	safeIP := sanitize(ip)
	safeHost := sanitize(hostname)

	dhcpLine := fmt.Sprintf("dhcp-host=%s,%s,%s", safeMAC, safeIP, safeHost)
	f, err := os.OpenFile(confFile, os.O_APPEND|os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("open config for append: %w", err)
	}
	defer f.Close()

	if _, err := fmt.Fprintln(f, dhcpLine); err != nil {
		return fmt.Errorf("write dhcp-host line: %w", err)
	}

	// If a custom hostname was given, add a DNS A record too.
	if safeHost != "" {
		dnsLine := fmt.Sprintf("address=/%s/%s", safeHost, safeIP)
		if _, err := fmt.Fprintln(f, dnsLine); err != nil {
			return fmt.Errorf("write address line: %w", err)
		}
	}
	return nil
}

// addressPattern matches lines like: address=/hostname/IP
var addressPattern = regexp.MustCompile(`^address\s*=\s*/(.+)/(.+)$`)

// unpinLease removes dhcp-host and associated address= lines matching the given MAC.
func unpinLease(mac string) error {
	macLower := strings.ToLower(mac)

	// First pass: find the IP from the dhcp-host line for this MAC.
	var targetIP string
	f, err := os.Open(confFile)
	if err != nil {
		return fmt.Errorf("open config file: %w", err)
	}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if matches := dhcpHostPattern.FindStringSubmatch(line); matches != nil {
			parts := strings.Split(matches[1], ",")
			if len(parts) > 1 && strings.TrimSpace(strings.ToLower(parts[0])) == macLower {
				targetIP = strings.TrimSpace(parts[1])
				break
			}
		}
	}
	f.Close()

	// Second pass: rebuild config, removing dhcp-host line + matching address= lines.
	f, err = os.Open(confFile)
	if err != nil {
		return fmt.Errorf("open config file: %w", err)
	}

	var kept []string
	removed := false
	scanner = bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)
		// Remove dhcp-host for this MAC.
		if dhcpHostPattern.MatchString(trimmed) {
			matchContent := dhcpHostPattern.FindStringSubmatch(trimmed)[1]
			parts := strings.Split(matchContent, ",")
			if len(parts) > 0 && strings.TrimSpace(strings.ToLower(parts[0])) == macLower {
				removed = true
				continue
			}
		}
		// Remove address= line for the same IP.
		if targetIP != "" && addressPattern.MatchString(trimmed) {
			addrMatches := addressPattern.FindStringSubmatch(trimmed)
			if strings.TrimSpace(addrMatches[2]) == targetIP {
				continue
			}
		}
		kept = append(kept, line)
	}
	f.Close()

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("scan config: %w", err)
	}

	if !removed {
		return nil // nothing to remove, not an error
	}

	// Write back atomically to avoid TOCTOU race.
	tmpFile := confFile + ".tmp"
	out, err := os.Create(tmpFile)
	if err != nil {
		return fmt.Errorf("create temp file: %w", err)
	}
	writeErr := func() error {
		defer out.Close()
		for _, line := range kept {
			if _, err := fmt.Fprintln(out, line); err != nil {
				return fmt.Errorf("write config line: %w", err)
			}
		}
		return nil
	}()
	if writeErr != nil {
		os.Remove(tmpFile)
		return writeErr
	}
	if err := os.Rename(tmpFile, confFile); err != nil {
		os.Remove(tmpFile)
		return fmt.Errorf("rename temp file: %w", err)
	}
	return nil
}

// reloadDnsmasq runs the reload command (e.g. sudo systemctl reload dnsmasq).
func reloadDnsmasq() error {
	// Scrub expired leases before reloading.
	if err := cleanLeaseFile(); err != nil {
		log.Printf("WARNING: failed to clean lease file: %v", err)
	}
	// Touch a trigger file — systemd.path unit watches this and reloads dnsmasq.
	triggerFile := confFile + ".reload"
	f, err := os.Create(triggerFile)
	if err != nil {
		return fmt.Errorf("create trigger file: %w", err)
	}
	f.Close()
	return nil
}

// cleanLeaseFile reads the lease file, removes expired entries, writes back atomically.
func cleanLeaseFile() error {
	leases, err := parseLeases()
	if err != nil {
		return err
	}

	now := time.Now()
	var kept []string
	for _, l := range leases {
		if l.Expiry.After(now) {
			kept = append(kept, fmt.Sprintf("%d %s %s %s %s",
				l.Expiry.Unix(), l.MAC, l.IP, l.Hostname, l.ClientID))
		}
	}

	// If nothing changed, skip the write.
	if len(kept) == len(leases) {
		return nil
	}

	tmpFile := leasesFile + ".tmp"
	f, err := os.Create(tmpFile)
	if err != nil {
		return fmt.Errorf("create temp lease file: %w", err)
	}
	writeErr := func() error {
		defer f.Close()
		for _, line := range kept {
			if _, err := fmt.Fprintln(f, line); err != nil {
				return err
			}
		}
		return nil
	}()
	if writeErr != nil {
		os.Remove(tmpFile)
		return writeErr
	}
	if err := os.Rename(tmpFile, leasesFile); err != nil {
		os.Remove(tmpFile)
		return fmt.Errorf("rename lease file: %w", err)
	}
	return nil
}

// getAllLeases returns merged active (non-expired) and pinned lease info.
func getAllLeases() ([]LeaseEntry, error) {
	activeLeases, err := parseLeases()
	if err != nil {
		return nil, fmt.Errorf("parse leases: %w", err)
	}

	pinned, err := parsePinnedHosts()
	if err != nil {
		return nil, fmt.Errorf("parse pinned: %w", err)
	}

	now := time.Now()
	entries := make([]LeaseEntry, 0, len(activeLeases))
	for _, lease := range activeLeases {
		// Skip expired leases (expiry is in the past and not pinned).
		if lease.Expiry.Before(now) {
			if _, ok := pinned[lease.MAC]; !ok {
				continue
			}
		}
		status := PinStatusActive
		hostname := lease.Hostname
		if pinnedHost, ok := pinned[lease.MAC]; ok {
			status = PinStatusPermanent
			// Use the pinned hostname if one was set (overrides DHCP client hostname).
			if pinnedHost != "" {
				hostname = pinnedHost
			}
		}
		entries = append(entries, LeaseEntry{
			Hostname: hostname,
			IP:       lease.IP,
			MAC:      lease.MAC,
			Expiry:   lease.Expiry,
			Status:   status,
		})
	}
	return entries, nil
}

// ---------------------------------------------------------------------------
// HTTP handlers
// ---------------------------------------------------------------------------

var (
	lastReload  time.Time
	lastReloadMu sync.Mutex
	reloadCooldown = 5 * time.Second
)

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func handleListLeases(w http.ResponseWriter, r *http.Request) {
	entries, err := getAllLeases()
	if err != nil {
		log.Printf("ERROR listing leases: %v", err)
		writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to read lease data"})
		return
	}

	activeCount := len(entries)
	pinnedCount := 0
	for _, e := range entries {
		if e.Status == PinStatusPermanent {
			pinnedCount++
		}
	}

	writeJSON(w, http.StatusOK, LeasesResponse{
		Leases:     entries,
		PinnedCount: pinnedCount,
		ActiveCount: activeCount,
	})
}

func handlePinLease(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "invalid form data: " + err.Error()})
		return
	}
	req := PinRequest{
		MAC:      r.FormValue("mac"),
		IP:       r.FormValue("ip"),
		Hostname: r.FormValue("hostname"),
	}

	if req.MAC == "" || req.IP == "" {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "mac and ip are required"})
		return
	}

	// If already pinned, unpin first so we can re-pin with updated hostname.
	pinned, err := parsePinnedHosts()
	if err != nil {
		log.Printf("ERROR parsing pinned hosts: %v", err)
		writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to read config"})
		return
	}
	if existingHost, exists := pinned[strings.ToLower(req.MAC)]; exists {
		if existingHost == req.Hostname {
			// Same hostname — nothing to change.
			writeJSON(w, http.StatusOK, PinResponse{Success: true, Message: "already pinned"})
			return
		}
		// Hostname changed — unpin old entry, then repin.
		if err := unpinLease(req.MAC); err != nil {
			log.Printf("ERROR unpinning for repin: %v", err)
			writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to update config"})
			return
		}
	}

	if err := pinLease(req.MAC, req.IP, req.Hostname); err != nil {
		log.Printf("ERROR pinning lease: %v", err)
		writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to write config"})
		return
	}
	writeJSON(w, http.StatusOK, PinResponse{Success: true, Message: "pinned successfully"})
}

func handleUnpinLease(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "invalid form data: " + err.Error()})
		return
	}
	req := PinRequest{
		MAC:      r.FormValue("mac"),
		IP:       r.FormValue("ip"),
		Hostname: r.FormValue("hostname"),
	}

	if req.MAC == "" {
		writeJSON(w, http.StatusBadRequest, ErrorResponse{Error: "mac is required"})
		return
	}

	if err := unpinLease(req.MAC); err != nil {
		log.Printf("ERROR unpinning lease: %v", err)
		writeJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "failed to update config"})
		return
	}
	writeJSON(w, http.StatusOK, PinResponse{Success: true, Message: "unpinned successfully"})
}

func handleReload(w http.ResponseWriter, r *http.Request) {
	lastReloadMu.Lock()
	elapsed := time.Since(lastReload)
	if elapsed < reloadCooldown {
		lastReloadMu.Unlock()
		writeJSON(w, http.StatusTooManyRequests, PinResponse{
			Success: false,
			Message: fmt.Sprintf("rate limited — wait %.0fs", reloadCooldown.Seconds()-elapsed.Seconds()),
		})
		return
	}
	lastReload = time.Now()
	lastReloadMu.Unlock()

	if err := reloadDnsmasq(); err != nil {
		log.Printf("ERROR reloading dnsmasq: %v", err)
		writeJSON(w, http.StatusOK, PinResponse{
			Success: false,
			Message: "reload failed — check server logs",
		})
		return
	}
	writeJSON(w, http.StatusOK, PinResponse{Success: true, Message: "dnsmasq reloaded"})
}