zoraxy-dhcp/server.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"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")
	reloadCmd   = envOrDefault("RELOAD_CMD", "sudo")
	reloadArgs  = strings.Fields(envOrDefault("RELOAD_ARGS", "systemctl reload dnsmasq"))
)

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.
func pinLease(mac, ip, hostname string) error {
	line := fmt.Sprintf("dhcp-host=%s,%s,%s", sanitize(mac), sanitize(ip), sanitize(hostname))
	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, line); err != nil {
		return fmt.Errorf("write dhcp-host line: %w", err)
	}
	return nil
}

// unpinLease removes dhcp-host lines matching the given MAC from dnsmasq.conf.
func unpinLease(mac string) error {
	macLower := strings.ToLower(mac)
	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)
		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 // skip this line
			}
		}
		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 {
	cmd := exec.Command(reloadCmd, reloadArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("reload failed: %w — %s", err, string(output))
	}
	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
		if _, ok := pinned[lease.MAC]; ok {
			status = PinStatusPermanent
		}
		entries = append(entries, LeaseEntry{
			Hostname: lease.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
	}

	// Check if already pinned.
	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 _, exists := pinned[strings.ToLower(req.MAC)]; exists {
		writeJSON(w, http.StatusConflict, PinResponse{
			Success: false,
			Message: fmt.Sprintf("%s is already pinned", req.MAC),
		})
		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"})
}