forked from alfred/HostPersister
133 lines
2.9 KiB
Go
133 lines
2.9 KiB
Go
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
|
|
}
|
|
}
|