ghost/main.go

package main

import (
	"fmt"
	"net"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"
)

func main() {
	targetIP, domains := parseArgs()
	targetFile := getHostsPath()
	motd()
	process(targetFile, targetIP, domains)
}

func parseArgs() (net.IP, []string) {
	// Check for arguments
	if len(os.Args) < 3 {
		usage()
		os.Exit(1)
	}

	// Parse IP
	rawIP := os.Args[1]
	ip := net.ParseIP(rawIP)
	if ip == nil {
		fmt.Fprintf(os.Stderr, "Not an IP address: %s\n", rawIP)
		usage()
		os.Exit(1)
	}

	// Check for formatting of domain names
	domainExpr, _ := regexp.Compile("^([a-zA-Z0-9]+[a-zA-Z0-9\\-]*[a-zA-Z0-9]+.{0,1})+$")
	for _, a := range os.Args[2:] {
		if !domainExpr.MatchString(a) {
			fmt.Fprintf(os.Stderr, "Doesn't look like a valid domain: %s\n", a)
			usage()
			os.Exit(1)
		}
	}

	// All parsed and correct
	return ip, os.Args[2:]
}

// getHostsPath returns the path to the hosts file depending on the current operating system:
//  /etc/hosts on *NIX,
//  C:\Windows\system32\drivers\etc\hosts on WIN*.
func getHostsPath() (string) {
	if runtime.GOOS == "windows" {
		return "C:\\Windows\\system32\\drivers\\etc\\hosts"
	} else {
		return "/etc/hosts"
	}
}

// process goes over the targetFile and tries to best-fit the domain names for the IP
func process(targetFile string, ip net.IP, domains []string) {
	// Read file
	bytesEtcHosts, readErr := os.ReadFile(targetFile)
	if readErr != nil {
		fmt.Fprintf(os.Stderr, "Failed to read %s: %s\n", targetFile, readErr.Error())
		os.Exit(1)
	}
	etcHosts := string(bytesEtcHosts)

	// Iterate over lines, find the first match for the IP
	lines := strings.Split(etcHosts, "\n")
	ipString := ip.String()
	found := false
	for lPos, l := range lines {
		if strings.HasPrefix(l, ipString) {
			// Matching line, append our domains.

			// Avoid duplicates
			fields := strings.Fields(l)
			newDomainsArray := []string{}
			ignore := false
			for _, d := range domains {
				for _, f := range fields {
					if d == f {
						// this domain name is already in the hosts file, skip
						ignore = true
						break
					}
				}
				if !ignore {
					newDomainsArray = append(newDomainsArray, d)
				}
			}
			newDomains := strings.Join(newDomainsArray, " ")

			// Preserve comments and construct the line again
			hostLine, comment, hasComment := strings.Cut(l, "#")
			if hasComment {
				lines[lPos] = fmt.Sprintf("%s %s # %s", hostLine, newDomains, comment)
			} else {
				lines[lPos] = fmt.Sprintf("%s %s", hostLine, newDomains)
			}
			found = true

			// Inform user
			fmt.Printf("Updated the line for '%s' to include '%s'\n", ipString, newDomains)
		}
	}

	// If a fitting line was not found in the previous for loop, append a new line
	if !found {
		newDomains := strings.Join(domains, " ")
		date := time.Now().Format("02-01-2006")
		newLine := fmt.Sprintf("%s\t%s # Added by ghost on %s", ipString, newDomains, date)
		lines = append(lines, newLine)

		// Inform user
		fmt.Printf("Appended line for '%s' to include '%s'\n", ipString, newDomains)
	}

	// Write out again
	newHosts := strings.Join(lines, "\n")
	writeErr := os.WriteFile(targetFile, []byte(newHosts), 0o644)
	if writeErr != nil {
		fmt.Fprintf(os.Stderr, "Failed to write %s: %s\n", targetFile, writeErr.Error())
		os.Exit(1)
	}
}

// usage prints the usage
func usage() {
	fmt.Fprintf(os.Stderr, "Usage: ghost 10.0.1.1 server1.gh <foo.server2.gh> <www.server2.gh> ...\n")
}

// motd prints the banner at the start, fully easter-egg free
func motd() {
	fmt.Println("+-------------------------+")
	if(time.Now().Month() == time.October) {
		// I mean, it's called 'ghost' after all...
		fmt.Println("👻 🕷️ 🎃 G H O S T 🎃 🕷️ 👻")
		fmt.Println(" your spooooky helper with ")
	} else {
		fmt.Println("         G H O S T         ")
		fmt.Println(" your friendly helper with ")
	}
	fmt.Printf(" %cetc%chosts, written in Go!\n", os.PathSeparator, os.PathSeparator) // technically not fully correct, but the MOTD can only be *that* wide, right?
	fmt.Println("+-------------------------+")
}