Merge pull request #27 from errpunk/feature/win-locate-hosts-via-syst…

…emroot

use %SystemRoot% environment variable for hosts file location on Windows

txeh.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"strings"
@@ -72,7 +73,7 @@ func NewHosts(hc *HostsConfig) (*Hosts, error) {
 	defaultHostsFile := "/etc/hosts"
 
 	if runtime.GOOS == "windows" {
-		defaultHostsFile = `C:\Windows\System32\Drivers\etc\hosts`
+		defaultHostsFile = winDefaultHostsFile()
 	}
 
 	if h.ReadFilePath == "" && h.RawText == nil {
@@ -519,3 +520,15 @@ var localhostIPRegexp = regexp.MustCompile(ipLocalhost)
 func isLocalhost(address string) bool {
 	return localhostIPRegexp.MatchString(address)
 }
+
+// winDefaultHostsFile returns the default hosts file path for Windows
+// it tries to use the SystemRoot environment variable, if that is not set
+// it falls back to C:\Windows\System32\Drivers\etc\hosts
+func winDefaultHostsFile() string {
+	if r := os.Getenv("SystemRoot"); r != "" {
+		return filepath.Join(r, "System32", "drivers", "etc", "hosts")
+	}
+
+	// fallback to C:\
+	return `C:\Windows\System32\drivers\etc\hosts`
+}

txeh_test.go

@@ -1,6 +1,8 @@
 package txeh
 
 import (
+	"os"
+	"runtime"
 	"strings"
 	"testing"
 )
@@ -303,3 +305,43 @@ func TestMethods(t *testing.T) {
 	}
 
 }
+
+func TestWinDefaultHostsFile(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip("Skipping windows test")
+	}
+
+	tc := []struct {
+		Name     string
+		EnvVaule string
+		Expect   string
+	}{
+		{
+			Name:     "Correct SystemRoot environment variable",
+			EnvVaule: "E:\\Windows",
+			Expect:   "E:\\Windows\\System32\\drivers\\etc\\hosts",
+		},
+		{
+			Name:     "SystemRoot with trailing slash",
+			EnvVaule: "E:\\Windows\\",
+			Expect:   "E:\\Windows\\System32\\drivers\\etc\\hosts",
+		},
+		{
+			Name:     "No systemRoot environment variable",
+			EnvVaule: "",
+			Expect:   "C:\\Windows\\System32\\drivers\\etc\\hosts",
+		},
+	}
+	for _, tt := range tc {
+		t.Run(tt.Name, func(t *testing.T) {
+			if err := os.Setenv("SystemRoot", tt.EnvVaule); err != nil {
+				panic(err)
+			}
+			hostFile := winDefaultHostsFile()
+			if hostFile != tt.Expect {
+				t.Fatalf("TestWinDefaultHostsFile failed to get expected path: "+
+					"expect: %v, got: %v", tt.Expect, hostFile)
+			}
+		})
+	}
+}