package main import ( "fmt" "github.com/joho/godotenv" "golang.captainalm.com/HostPersister/hosts" "log" "os" "os/signal" "strconv" "sync" "syscall" "time" ) var ( buildVersion = "develop" buildDate = "" ) func main() { log.Printf("[Main] Starting up Host Persister #%s (%s)\n", buildVersion, buildDate) y := time.Now() //Hold main thread till safe shutdown exit: wg := &sync.WaitGroup{} wg.Add(1) //Load environment file: err := godotenv.Load() if err != nil { log.Fatalln("Error loading .env file") } //Load ENVs hostsPath := os.Getenv("HOSTS_FILE") sourcePath := os.Getenv("SOURCE_FILE") syncTime, err := strconv.ParseInt(os.Getenv("SYNC_TIME"), 10, 64) if err != nil { log.Println("[Main] Invalid SYNC_TIME; defaulting to one-shot execution.") syncTime = 0 } overwriteMode := os.Getenv("HOSTS_OVERWRITE") == "1" //Load hosts and source files hostsFile, err := hosts.NewHostsFile(hostsPath) if err != nil { log.Fatalln("Failed to load HOSTS_FILE") } sourceFile, err := hosts.NewHostsFile(sourcePath) if err != nil { log.Fatalln("Failed to load SOURCE_FILE") } if syncTime < 1 { //One-shot execution: z := time.Now().Sub(y) log.Printf("[Main] Took '%s' to fully initialize modules\n", z.String()) executePersistence(hostsFile, sourceFile, overwriteMode) } else { //Sync execution: sigs := make(chan os.Signal, 1) signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) //Finished notifier: finNotf := make(chan bool, 1) z := time.Now().Sub(y) log.Printf("[Main] Took '%s' to fully initialize modules\n", z.String()) exec := true syncDur := time.Duration(syncTime) * time.Millisecond go func(exec *bool) { executePersistence(hostsFile, sourceFile, overwriteMode) sleep(syncDur, finNotf) for *exec { err := hostsFile.ReadHostsFile() if err != nil { log.Fatalln("Failed to load HOSTS_FILE") } err = sourceFile.ReadHostsFile() if err != nil { log.Fatalln("Failed to load SOURCE_FILE") } executePersistence(hostsFile, sourceFile, overwriteMode) sleep(syncDur, finNotf) } }(&exec) go func(exec *bool) { <-sigs fmt.Printf("\n") *exec = false finNotf <- true a := time.Now() log.Printf("[Main] Signalling program exit...\n") b := time.Now().Sub(a) log.Printf("[Main] Took '%s' to fully shutdown modules\n", b.String()) wg.Done() }(&exec) wg.Wait() } log.Println("[Main] Goodbye") } func executePersistence(hostsFile *hosts.File, sourceFile *hosts.File, ovrw bool) { for _, entry := range sourceFile.Entries { for _, domain := range entry.Domains { if (!hostsFile.HasDomain(domain)) || ovrw { hostsFile.OverwriteDomainSingleton(domain, entry.IPAddress) } } } err := hostsFile.WriteHostsFile() if err != nil { log.Println("[Main] Error Writing Hosts File.") } } func sleep(d time.Duration, n chan bool) { select { case <-n: break case <-time.After(d): break } }