dnshelper/main.go

package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"ekdns/hostfile"
	"ekdns/lockfile"
	"ekdns/platform"
	"ekdns/resolver"
)

// version and copyrightYear are set at build time via -ldflags.
var version = "dev"
var copyrightYear = "2024"

func main() {
	// Banner (always printed to stdout).
	fmt.Printf("ekDNSHelper %s\n", version)
	fmt.Printf("Copyright (c) %s Emberkom LLC\n", copyrightYear)
	fmt.Println("")

	if len(os.Args) < 2 {
		printUsage()
		os.Exit(1)
	}

	switch strings.ToLower(os.Args[1]) {
	case "a", "add":
		os.Exit(runAdd(os.Args[2:]))
	case "d", "del", "delete":
		os.Exit(runDelete(os.Args[2:]))
	default:
		printUsage()
		os.Exit(1)
	}
}

func runAdd(args []string) int {
	fs := flag.NewFlagSet("add", flag.ContinueOnError)
	hostFlag := fs.String("host", "", "Comma separated list of hostnames to add")
	serverFlag := fs.String("server", "", "Resolution mode: omit for smart default, local, gateway, or IP/IP:port")
	timeoutFlag := fs.Int("timeout", 3, "Per-query DNS timeout in seconds")
	verboseFlag := fs.Bool("verbose", false, "Emit per-stage resolution trace to stderr")
	if err := fs.Parse(args); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Validation (FR-006).
	if *hostFlag == "" {
		fmt.Fprintf(os.Stderr, "Error: -host flag is required for the add command\n")
		return 1
	}
	if *timeoutFlag <= 0 {
		fmt.Fprintf(os.Stderr, "Error: -timeout must be a positive integer\n")
		return 1
	}

	mode, err := resolver.ParseServerFlag(*serverFlag)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	config := resolver.QueryConfig{
		Timeout: time.Duration(*timeoutFlag) * time.Second,
		Verbose: *verboseFlag,
	}
	discoverer := platform.NewNetworkDiscoverer()

	hostnames := strings.Split(*hostFlag, ",")

	// Step 1: DNS resolution before lock acquisition (FR-016).
	var entries []hostfile.DNSEntry
	var warnings []string
	for _, h := range hostnames {
		h = strings.TrimSpace(h)
		ips, err := resolver.Resolve(h, mode, config, discoverer)
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("Warning: failed to resolve %s: %v", h, err))
			continue
		}
		for _, ip := range ips {
			entries = append(entries, hostfile.DNSEntry{IP: ip, Hostname: h})
		}
	}

	if len(entries) == 0 && len(warnings) > 0 {
		for _, w := range warnings {
			fmt.Fprintln(os.Stderr, w)
		}
		return 1
	}

	// Step 2: Get hosts file path.
	hostsPath, err := platform.GetHostsFilePath()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 3: Get executable directory for lock and backup storage.
	exePath, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: determining executable path: %v\n", err)
		return 1
	}
	exeDir := filepath.Dir(exePath)

	// Step 4: Acquire lock (FR-016).
	lock, err := lockfile.Acquire(exeDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}
	defer lock.Release()

	// Step 5: Read, modify, write.
	mgr := hostfile.NewManager(hostfile.OSFileSystem{})
	hf, err := mgr.Read(hostsPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	hf.ManagedContent = hostfile.AddEntries(hf.ManagedContent, entries)

	if err := mgr.Write(hf, exeDir); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 6: Output summary (FR-013).
	entryCounts := make(map[string]int)
	for _, e := range entries {
		entryCounts[e.Hostname]++
	}
	for _, h := range hostnames {
		h = strings.TrimSpace(h)
		if count, ok := entryCounts[h]; ok {
			fmt.Printf("Added %d entries for %s\n", count, h)
		}
	}
	for _, w := range warnings {
		fmt.Fprintln(os.Stderr, w)
	}

	if len(warnings) > 0 {
		return 1 // Partial failure (FR-015).
	}
	return 0
}

func runDelete(args []string) int {
	// Check for "all" keyword first — positional argument, not a flag.
	if len(args) > 0 && (strings.ToLower(args[0]) == "all" || strings.ToLower(args[0]) == "a") {
		return runDeleteAll()
	}

	fs := flag.NewFlagSet("delete", flag.ContinueOnError)
	hostFlag := fs.String("host", "", "Comma separated list of hostnames to delete")
	if err := fs.Parse(args); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	if *hostFlag == "" {
		fmt.Fprintf(os.Stderr, "Error: -host flag is required for the delete command\n")
		return 1
	}

	hostnames := strings.Split(*hostFlag, ",")
	for i, h := range hostnames {
		hostnames[i] = strings.TrimSpace(h)
	}

	// Step 1: Get hosts file path.
	hostsPath, err := platform.GetHostsFilePath()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 2: Get executable directory.
	exePath, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: determining executable path: %v\n", err)
		return 1
	}
	exeDir := filepath.Dir(exePath)

	// Step 3: Acquire lock.
	lock, err := lockfile.Acquire(exeDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}
	defer lock.Release()

	// Step 4: Read and parse hosts file.
	mgr := hostfile.NewManager(hostfile.OSFileSystem{})
	hf, err := mgr.Read(hostsPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 5: Remove entries by hostname, counting originals for summary.
	// Count original managed lines per hostname before removal.
	originalCounts := make(map[string]int)
	for _, h := range hostnames {
		for _, line := range hf.ManagedContent {
			if strings.Contains(line, h) {
				originalCounts[h]++
			}
		}
	}

	hf.ManagedContent = hostfile.RemoveByHostname(hf.ManagedContent, hostnames)

	// Step 6: Write updated file.
	if err := mgr.Write(hf, exeDir); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 7: Print per-hostname summary.
	for _, h := range hostnames {
		if count, ok := originalCounts[h]; ok && count > 0 {
			fmt.Printf("Removed %d entries for %s\n", count, h)
		} else {
			fmt.Printf("No managed entries found for %s\n", h)
		}
	}
	return 0
}

func runDeleteAll() int {
	// Step 1: Get hosts file path.
	hostsPath, err := platform.GetHostsFilePath()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 2: Get executable directory.
	exePath, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: determining executable path: %v\n", err)
		return 1
	}
	exeDir := filepath.Dir(exePath)

	// Step 3: Acquire lock.
	lock, err := lockfile.Acquire(exeDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}
	defer lock.Release()

	// Step 4: Read and parse hosts file.
	mgr := hostfile.NewManager(hostfile.OSFileSystem{})
	hf, err := mgr.Read(hostsPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 5: Remove all managed entries.
	entryCount := len(hf.ManagedContent)
	hf.ManagedContent = hostfile.RemoveAll()

	// Step 6: Write updated file.
	if err := mgr.Write(hf, exeDir); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 1
	}

	// Step 7: Print summary.
	if entryCount > 0 {
		fmt.Printf("Removed all managed entries (%d entries removed)\n", entryCount)
	} else {
		fmt.Println("No managed entries found")
	}
	return 0
}

func printUsage() {
	fmt.Println("Updates the local hosts file with DNS entries resolved by this tool.")
	fmt.Println("")
	fmt.Println("Usage:")
	fmt.Println("  ekdns add -host <hostnames> [-server <mode>] [-timeout <seconds>] [-verbose]")
	fmt.Println("  ekdns delete -host <hostnames>")
	fmt.Println("  ekdns delete all")
	fmt.Println("")
	fmt.Println("Flags (add):")
	fmt.Println("  -host      Comma-separated list of hostnames to resolve and add (required)")
	fmt.Println("  -server    Resolution mode (optional):")
	fmt.Println("               (omitted)   Smart default: NS fan-out → authoritative query → parallel A fallback")
	fmt.Println("               local       Query locally configured DNS resolvers only (no public fallback)")
	fmt.Println("               gateway     Query default gateway as DNS server (no fallback)")
	fmt.Println("               <ip>        Query this IP on port 53")
	fmt.Println("               <ip>:<port> Query this IP on the given port")
	fmt.Println("  -timeout   Per-query DNS timeout in seconds (default: 3)")
	fmt.Println("  -verbose   Emit per-stage resolution trace to stderr")
	fmt.Println("")
	fmt.Println("Examples:")
	fmt.Println("  ekdns add -host www.example.com")
	fmt.Println("    Resolve using smart default (NS fan-out + authoritative query).")
	fmt.Println("  ekdns add -host internal.client.local -server local")
	fmt.Println("    Resolve using only locally configured DNS resolvers.")
	fmt.Println("  ekdns add -host www.example.com -server gateway")
	fmt.Println("    Resolve via the default gateway IP.")
	fmt.Println("  ekdns add -host www.example.com -server 8.8.8.8")
	fmt.Println("    Resolve via a specific DNS server.")
	fmt.Println("  ekdns add -host www.example.com -server 8.8.8.8:5353 -timeout 5 -verbose")
	fmt.Println("    Resolve via 8.8.8.8:5353 with 5s timeout and verbose trace.")
	fmt.Println("  ekdns delete -host www.example.com")
	fmt.Println("    Remove all managed entries for www.example.com from the hosts file.")
	fmt.Println("  ekdns delete all")
	fmt.Println("    Remove all managed entries from the hosts file.")
	fmt.Println("")
	fmt.Println("Notes:")
	fmt.Println("  Adding a hostname first removes all existing managed entries for that hostname.")
	fmt.Println("  Only entries added by this tool are ever removed.")
}