replace existing entry from hosts file

hortonworks · Jan 25, 2017 · 6c8996d · 6c8996d
1 parent b2f1a39
commit 6c8996d
Show file tree

Hide file tree

Showing 2 changed files with 205 additions and 11 deletions.
diff --git a/saltboot/hostname.go b/saltboot/hostname.go
@@ -3,6 +3,7 @@ package saltboot
 import (
 	"io/ioutil"
 	"log"
+	"os"
 	"strings"
 )
 
@@ -42,19 +43,22 @@ func ensureIpv6Resolvable(customDomain string) error {
 	} else {
 		domain = customDomain
 	}
-	updateIpv6HostName(hostname, domain)
+	updateIpv6HostName(hostname, domain, HOST_FILE_NAME, getIpv4Address, ioutil.ReadFile, ioutil.WriteFile)
 
 	return nil
 }
 
-func updateIpv6HostName(hostName string, domain string) error {
+func updateIpv6HostName(hostName string, domain string, file string,
+	getIpv4Address func() (string, error),
+	readFile func(filename string) ([]byte, error),
+	writeFile func(filename string, data []byte, perm os.FileMode) error) error {
 	log.Printf("[updateIpv6HostName] hostName: %s, domain: %s", hostName, domain)
-	b, err := ioutil.ReadFile(HOST_FILE_NAME)
+	b, err := readFile(file)
 	if err != nil {
 		return err
 	}
-	hostfile := string(b)
-	log.Printf("[updateIpv6HostName] original hostfile: %s", hostfile)
+	hostsFile := string(b)
+	log.Printf("[updateIpv6HostName] original hosts file: %s", hostsFile)
 	address, err := getIpv4Address()
 	if err != nil {
 		return err
@@ -64,14 +68,20 @@ func updateIpv6HostName(hostName string, domain string) error {
 	}
 	ipv6hostString := address + " " + hostName + domain + " " + hostName
 	log.Printf("[updateIpv6HostName] ipv6hostString: %s", ipv6hostString)
-	if !strings.Contains(hostfile, address) {
-		hostfile = hostfile + "\n" + ipv6hostString
-		log.Printf("[updateIpv6HostName] updated hostfile: %s", hostfile)
-		err = ioutil.WriteFile(HOST_FILE_NAME, []byte(hostfile), 0644)
-		if err != nil {
-			return err
+
+	lines := strings.Split(hostsFile, "\n")
+	var filteredLines = make([]string, 0)
+	for _, line := range lines {
+		if !strings.Contains(line, address) {
+			filteredLines = append(filteredLines, line)
 		}
 	}
+	hostsFile = strings.Join(filteredLines, "\n") + "\n" + ipv6hostString
+	log.Printf("[updateIpv6HostName] updated hosts file: %s", hostsFile)
+	err = writeFile(file, []byte(hostsFile), 0644)
+	if err != nil {
+		return err
+	}
 
 	return nil
 }
diff --git a/saltboot/hostname_test.go b/saltboot/hostname_test.go
@@ -0,0 +1,184 @@
+package saltboot
+
+import (
+	"os"
+	"testing"
+)
+
+func TestHostsFileWriteRemoveExistingIp(t *testing.T) {
+	getIpv4Address := func() (string, error) {
+		return "10.0.0.1", nil
+	}
+
+	readFile := func(filename string) ([]byte, error) {
+		hostsFile := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.1 hostname-1.compute.internal hostname-1`
+		return []byte(hostsFile), nil
+	}
+
+	var result string
+	writeFile := func(filename string, data []byte, perm os.FileMode) error {
+		result = string(data)
+		return nil
+	}
+
+	updateIpv6HostName("hostname-1", "example.com", "hosts", getIpv4Address, readFile, writeFile)
+
+	expected := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.1 hostname-1.example.com hostname-1`
+
+	if expected != result {
+		t.Errorf("Invalid hostname replacement, %s != %s", expected, result)
+	}
+}
+
+func TestHostsFileWriteRemoveExistingIpNotLastLine(t *testing.T) {
+	getIpv4Address := func() (string, error) {
+		return "10.0.0.1", nil
+	}
+
+	readFile := func(filename string) ([]byte, error) {
+		hostsFile := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.1 hostname-1.compute.internal hostname-1
+10.0.0.2 hostname-2.compute.internal hostname-2`
+		return []byte(hostsFile), nil
+	}
+
+	var result string
+	writeFile := func(filename string, data []byte, perm os.FileMode) error {
+		result = string(data)
+		return nil
+	}
+
+	updateIpv6HostName("hostname-1", "example.com", "hosts", getIpv4Address, readFile, writeFile)
+
+	expected := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2.compute.internal hostname-2
+10.0.0.1 hostname-1.example.com hostname-1`
+
+	if expected != result {
+		t.Errorf("Invalid hostname replacement, %s != %s", expected, result)
+	}
+}
+
+func TestHostsFileWriteRemoveExistingIpMiddleLastLine(t *testing.T) {
+	getIpv4Address := func() (string, error) {
+		return "10.0.0.1", nil
+	}
+
+	readFile := func(filename string) ([]byte, error) {
+		hostsFile := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2.compute.internal hostname-2
+10.0.0.1 hostname-1.compute.internal hostname-1
+10.0.0.3 hostname-3.compute.internal hostname-3`
+		return []byte(hostsFile), nil
+	}
+
+	var result string
+	writeFile := func(filename string, data []byte, perm os.FileMode) error {
+		result = string(data)
+		return nil
+	}
+
+	updateIpv6HostName("hostname-1", "example.com", "hosts", getIpv4Address, readFile, writeFile)
+
+	expected := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2.compute.internal hostname-2
+10.0.0.3 hostname-3.compute.internal hostname-3
+10.0.0.1 hostname-1.example.com hostname-1`
+
+	if expected != result {
+		t.Errorf("Invalid hostname replacement, %s != %s", expected, result)
+	}
+}
+
+func TestHostsFileWriteIpNotPresent(t *testing.T) {
+	getIpv4Address := func() (string, error) {
+		return "10.0.0.1", nil
+	}
+
+	readFile := func(filename string) ([]byte, error) {
+		hostsFile := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2
+10.0.0.3 hostname-3`
+		return []byte(hostsFile), nil
+	}
+
+	var result string
+	writeFile := func(filename string, data []byte, perm os.FileMode) error {
+		result = string(data)
+		return nil
+	}
+
+	updateIpv6HostName("hostname-1", "example.com", "hosts", getIpv4Address, readFile, writeFile)
+
+	expected := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2
+10.0.0.3 hostname-3
+10.0.0.1 hostname-1.example.com hostname-1`
+
+	if expected != result {
+		t.Errorf("Invalid hostname replacement, %s != %s", expected, result)
+	}
+}
+
+func TestHostsFileWriteExistingWithDefaultDomain(t *testing.T) {
+	getIpv4Address := func() (string, error) {
+		return "10.0.0.1", nil
+	}
+
+	readFile := func(filename string) ([]byte, error) {
+		hostsFile := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.1 hostname-1.compute.internal hostname-1
+10.0.0.2 hostname-2.compute.internal hostname-2
+10.0.0.3 hostname-3.compute.internal hostname-3`
+		return []byte(hostsFile), nil
+	}
+
+	var result string
+	writeFile := func(filename string, data []byte, perm os.FileMode) error {
+		result = string(data)
+		return nil
+	}
+
+	updateIpv6HostName("hostname-1", ".compute.internal", "hosts", getIpv4Address, readFile, writeFile)
+
+	expected := `
+127.0.0.1	localhost
+255.255.255.255	broadcasthost
+::1             localhost
+10.0.0.2 hostname-2.compute.internal hostname-2
+10.0.0.3 hostname-3.compute.internal hostname-3
+10.0.0.1 hostname-1.compute.internal hostname-1`
+
+	if expected != result {
+		t.Errorf("Invalid hostname replacement, %s != %s", expected, result)
+	}
+}