Feature/save dns records to hosts

commands/dns-records.go

@@ -1,9 +1,11 @@
 package commands
 
 import (
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"strings"
 
 	"github.com/bitly/go-simplejson"
@@ -17,12 +19,61 @@ type DNSRecords struct {
 	BaseCommand
 }
 
+// DNSRecord is the struct for a single DNS entry
+type DNSRecord struct {
+	Id      string // nolint
+	Name    string
+	Image   string
+	IPs     []string
+	TTL     int64
+	Aliases []string
+}
+
+// DNSRecordsList is an array of DNSRecords
+type DNSRecordsList []*DNSRecord
+
+const (
+	unixHostsPreamble  = "##+++ added by rig"
+	unixHostsPostamble = "##--- end rig additions"
+)
+
+func (record *DNSRecord) String() string {
+	result := ""
+	for _, ip := range record.IPs {
+		result += fmt.Sprintf("%s\t%s.%s.%s\n", ip, record.Name, record.Image, "vm")
+		// attach any aliases too
+		for _, a := range record.Aliases {
+			result += fmt.Sprintf("%s\t%s\n", ip, a)
+		}
+	}
+	return result
+}
+
+// String converts a list of DNSRecords to a formatted string
+func (hosts DNSRecordsList) String() string {
+	result := ""
+	for _, host := range hosts {
+		result += host.String()
+	}
+	return result
+}
+
 // Commands returns the operations supported by this command
 func (cmd *DNSRecords) Commands() []cli.Command {
 	return []cli.Command{
 		{
-			Name:   "dns-records",
-			Usage:  "List all DNS records for running containers",
+			Name:  "dns-records",
+			Usage: "List all DNS records for running containers",
+			Flags: []cli.Flag{
+				cli.BoolFlag{
+					Name:  "save",
+					Usage: "Save the DNS records to /etc/hosts or {FIX insert Windows desscription}",
+				},
+				cli.BoolFlag{
+					Name:  "remove",
+					Usage: "Remove the DNS records from /etc/hosts or {FIX insert Windows desscription}",
+				},
+			},
 			Before: cmd.Before,
 			Action: cmd.Run,
 		},
@@ -31,27 +82,31 @@ func (cmd *DNSRecords) Commands() []cli.Command {
 
 // Run executes the `rig dns-records` command
 func (cmd *DNSRecords) Run(c *cli.Context) error {
+	// Don't require rig to be started to remove records
+	if c.Bool("remove") {
+		if c.Bool("save") {
+			return cmd.Failure("--remove and --save are mutually exclusive", "COMMAND-ERROR", 13)
+		}
+		// TODO The VM might have to be up for Windows
+		return cmd.removeDNSRecords()
+	}
 
 	records, err := cmd.LoadRecords()
 	if err != nil {
 		return cmd.Failure(err.Error(), "COMMAND-ERROR", 13)
 	}
 
-	for _, record := range records {
-		for _, ip := range record["IPs"].([]interface{}) {
-			fmt.Printf("%s\t%s.%s.%s\n", ip, record["Name"], record["Image"], "vm")
-			// attach any aliases too
-			for _, a := range record["Aliases"].([]interface{}) {
-				fmt.Printf("%s\t%s\n", ip, a)
-			}
-		}
+	if c.Bool("save") {
+		return cmd.saveDNSRecords(records)
 	}
 
+	printDNSRecords(records)
+
 	return cmd.Success("")
 }
 
 // LoadRecords retrieves the records from DNSDock and processes/return them
-func (cmd *DNSRecords) LoadRecords() ([]map[string]interface{}, error) {
+func (cmd *DNSRecords) LoadRecords() ([]*DNSRecord, error) {
 	ip, err := util.Command("docker", "inspect", "--format", "{{.NetworkSettings.IPAddress}}", "dnsdock").Output()
 	if err != nil {
 		return nil, fmt.Errorf("failed to discover dnsdock IP address: %s", err)
@@ -74,11 +129,142 @@ func (cmd *DNSRecords) LoadRecords() ([]map[string]interface{}, error) {
 	}
 
 	dnsdockMap, _ := js.Map()
-	records := []map[string]interface{}{}
-	for id, value := range dnsdockMap {
-		record := value.(map[string]interface{})
-		record["Id"] = id
+	records := make([]*DNSRecord, 0, 20)
+	for id, rawValue := range dnsdockMap {
+		// Cast rawValue to its actual type
+		value := rawValue.(map[string]interface{})
+		ttl, _ := value["TTL"].(json.Number).Int64()
+		record := &DNSRecord{
+			Id:    id,
+			Name:  value["Name"].(string),
+			Image: value["Image"].(string),
+			TTL:   ttl,
+		}
+		record.IPs = make([]string, 0, 10)
+		for _, ip := range value["IPs"].([]interface{}) {
+			record.IPs = append(record.IPs, ip.(string))
+		}
+		record.Aliases = make([]string, 0, 10)
+		for _, alias := range value["Aliases"].([]interface{}) {
+			record.Aliases = append(record.Aliases, alias.(string))
+		}
 		records = append(records, record)
 	}
 	return records, nil
 }
+
+func printDNSRecords(records []*DNSRecord) {
+	for _, record := range records {
+		fmt.Print(record)
+	}
+}
+
+// Write the records to /etc/hosts or FIX Windows?
+func (cmd *DNSRecords) saveDNSRecords(records []*DNSRecord) error {
+	if util.IsMac() || util.IsLinux() {
+		return cmd.saveDNSRecordsUnix(records)
+	} else if util.IsWindows() {
+		return cmd.saveDNSRecordsWindows(records)
+	}
+	return cmd.Success("Not implemented")
+}
+
+// ⚠ Administrative privileges needed...
+
+func (cmd *DNSRecords) saveDNSRecordsUnix(records []*DNSRecord) error {
+	// Both of these are []string
+	oldHostEntries := util.LoadFile("/etc/hosts")
+	newHostEntries := stripDNS(oldHostEntries)
+	// records.String does the formatting, so convert both to a string
+	oldHosts := strings.Join(oldHostEntries, "\n")
+	newHosts := strings.Join(newHostEntries, "\n") + "\n" +
+		unixHostsPreamble + "\n" +
+		DNSRecordsList(records).String() +
+		unixHostsPostamble + "\n"
+	if oldHosts == newHosts {
+		return cmd.Success("No changes made")
+	}
+	return cmd.writeEtcHosts(newHosts)
+}
+
+func (cmd *DNSRecords) saveDNSRecordsWindows(records []*DNSRecord) error {
+	return cmd.Failure("Not Implemented", "COMMAND-ERROR", 13)
+}
+
+func (cmd *DNSRecords) removeDNSRecords() error {
+	if util.IsMac() || util.IsLinux() {
+		return cmd.removeDNSRecordsUnix()
+	} else if util.IsWindows() {
+		return cmd.removeDNSRecordsWindows()
+	}
+	return cmd.Success("Not implemented")
+}
+
+func (cmd *DNSRecords) removeDNSRecordsUnix() error {
+	oldHostsEntries := util.LoadFile("/etc/hosts")
+	newHostsEntries := stripDNS(oldHostsEntries)
+	oldHosts := strings.Join(oldHostsEntries, "\n")
+	newHosts := strings.Join(newHostsEntries, "\n")
+	if oldHosts == newHosts {
+		return cmd.Success("No changes made")
+	}
+	return cmd.writeEtcHosts(newHosts)
+}
+
+func (cmd *DNSRecords) removeDNSRecordsWindows() error {
+	return cmd.Failure("Not Implemented", "COMMAND-ERROR", 13)
+}
+
+// Save a new version of /etc/hosts, arg is the full text to save
+func (cmd *DNSRecords) writeEtcHosts(hostsText string) error {
+	// Make sure it ends in a newline
+	if hostsText[len(hostsText)-1] != '\n' {
+		hostsText += "\n"
+	}
+	// Write new version to a temp file
+	tmpfile, err := ioutil.TempFile("", "rig-hosts")
+	if err != nil {
+		return cmd.Failure("Unable to create hosts tempfile: "+err.Error(), "COMMAND-ERROR", 13)
+	}
+	tmpname := tmpfile.Name()
+	defer os.Remove(tmpname)
+	if _, err := tmpfile.Write([]byte(hostsText)); err != nil {
+		return cmd.Failure("Unable to write hosts tempfile: ("+tmpname+") "+err.Error(), "COMMAND-ERROR", 13)
+	}
+	if err := tmpfile.Close(); err != nil {
+		return cmd.Failure("Unable to close hosts tempfile: ("+tmpname+") "+err.Error(), "COMMAND-ERROR", 13)
+	}
+	// mv it into place. This is safer than trying to write /etc/hosts on the fly.
+	if err := util.EscalatePrivilege(); err != nil {
+		return cmd.Failure("Unable to obtain privileges to replace: "+err.Error(), "COMMAND-ERROR", 13)
+	}
+	if err := util.Command("sudo", "mv", "-f", tmpname, "/etc/hosts").Run(); err != nil {
+		return cmd.Failure("Unable to replace /etc/hosts: "+err.Error(), "COMMAND-ERROR", 13)
+	}
+	return cmd.Success("/etc/hosts updated")
+}
+
+// Remove a section of the hosts file we previously added
+func stripDNS(hosts []string) []string {
+	const (
+		looking = iota
+		found
+	)
+	results := make([]string, 0, 1000)
+	state := looking
+	for _, host := range hosts {
+		switch state {
+		case looking:
+			if host == unixHostsPreamble {
+				state = found
+			} else {
+				results = append(results, host)
+			}
+		case found:
+			if host == unixHostsPostamble {
+				state = looking
+			}
+		}
+	}
+	return results
+}

commands/doctor.go

@@ -135,9 +135,9 @@ func (cmd *Doctor) Run(c *cli.Context) error {
 	if records, err := dnsRecords.LoadRecords(); err == nil {
 		resolved := false
 		for _, record := range records {
-			if record["Name"] == "dnsdock" {
+			if record.Name == "dnsdock" {
 				resolved = true
-				cmd.out.Info("DNS and routing services are working. DNSDock resolves to %s", record["IPs"])
+				cmd.out.Info("DNS and routing services are working. DNSDock resolves to %s", record.IPs)
 				break
 			}
 		}

util/filesystem.go

@@ -1,9 +1,12 @@
 package util
 
 import (
+	"bufio"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/kardianos/osext"
 )
@@ -91,3 +94,26 @@ func TouchFile(pathToFile string, workingDir string) error {
 	f.Close()
 	return nil
 }
+
+// LoadFile loads a file into an array, without the newlines
+func LoadFile(filename string) []string {
+	lines := make([]string, 0, 1000)
+	f, err := os.Open(filename)
+	if err != nil {
+		return lines
+	}
+	defer f.Close()
+	r := bufio.NewReader(f)
+	for {
+		switch ln, err := r.ReadString('\n'); err {
+		case nil:
+			ln = strings.Replace(ln, "\r", "", -1)
+			ln = strings.Replace(ln, "\n", "", -1)
+			lines = append(lines, ln)
+		case io.EOF:
+			return lines
+		default:
+			fmt.Println(err)
+		}
+	}
+}