dnshelper/main.go

package main

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

	"dns-helper/hostfile"
	"dns-helper/lockfile"
	"dns-helper/platform"
	"dns-helper/resolver"
)

func main() {
	// Banner (always printed to stdout).
	fmt.Println("DNSHelper v1.0")
	fmt.Println("Copyright (c) 2024 Emberkom LLC")
	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", "", "DNS resolver to use")
	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 *serverFlag == "" {
		fmt.Fprintf(os.Stderr, "Error: -server flag is required for the add command\n")
		return 1
	}

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

	// Step 1: DNS resolution before lock acquisition (FR-016).
	res := resolver.New()
	var entries []hostfile.DNSEntry
	var warnings []string
	for _, h := range hostnames {
		h = strings.TrimSpace(h)
		ips, err := res.LookupIP(h, server)
		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("This utility will update the local hosts file with DNS entries obtained by the specified DNS server.")
	fmt.Println("Usage: dns-helper [add|delete] -host hostname [-server dns.example.com]")
	fmt.Println("Example:")
	fmt.Println("       dns-helper add -host xyz.acme.com -server dns.example.com")
	fmt.Println("         This will use dns.example.com to find and add all IP addresses for xyz.acme.com to the local hosts file.")
	fmt.Println("       dns-helper delete -host hostname.example.com")
	fmt.Println("         This will delete all entries for hostname.example.com from the local hosts file.")
	fmt.Println("       dns-helper delete all")
	fmt.Println("         This will delete all entries from the local hosts file that were added by this utility.")
	fmt.Println("Note: Adding a hostname will first remove all entries in the hosts file that match the same hostname.")
	fmt.Println("      This utility will only remove entries from the hosts file that it added.")
}