Merge pull request #3 from pacodes/master

Adds new command for adding domains to hosts file and some tests

cmd/add.go

@@ -20,7 +20,12 @@ If the profile already exists it will be added to it.`,
 
 		h, _ := cmd.Flags().GetString("host-file")
 
-		err := host.AddFromFile(from, h, profile, false)
+		err := host.AddFromFile(&host.AddFromFileOptions{
+			From:    from,
+			Dst:     h,
+			Profile: profile,
+			Reset:   false,
+		})
 		if err != nil {
 			return err
 		}
@@ -41,8 +46,51 @@ If the profile already exists it will be added to it.`,
 	},
 }
 
+// addDomainsCmd represents the fromFile command
+var addDomainsCmd = &cobra.Command{
+	Use:   "domains",
+	Short: "Add content in your hosts file.",
+	Long: `
+Set content in your hosts file.
+If the profile already exists it will be added to it.`,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		src, _ := cmd.Flags().GetString("host-file")
+		ip, _ := cmd.Flags().GetString("ip")
+		profile, _ := cmd.Flags().GetString("profile")
+		if profile == "" {
+			profile = "other"
+		}
+
+		h, _ := cmd.Flags().GetString("host-file")
+
+		err := host.AddFromArgs(&host.AddFromArgsOptions{
+			Domains: args,
+			IP:      ip,
+			Dst:     h,
+			Profile: profile,
+			Reset:   false,
+		})
+		if err != nil {
+			return err
+		}
+
+		return host.ListProfiles(src, &host.ListOptions{
+			Profile: profile,
+		})
+	},
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		profile, _ := cmd.Flags().GetString("profile")
+
+		return host.ValidProfile(profile)
+	},
+}
+
 func init() {
 	rootCmd.AddCommand(addFromFileCmd)
 
 	addFromFileCmd.Flags().StringP("from", "f", "", "file to read")
+
+	addFromFileCmd.AddCommand(addDomainsCmd)
+
+	addDomainsCmd.Flags().String("ip", "127.0.0.1", "domains ip")
 }

cmd/backup.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"errors"
+	"github.com/guumaster/tablewriter"
 	"log"
 	"os"
 
@@ -30,7 +31,13 @@ as extension.
 		src, _ := cmd.Flags().GetString("host-file")
 		dst, _ := cmd.Flags().GetString("path")
 
-		return host.BackupFile(src, dst)
+		backupFile, err := host.BackupFile(src, dst)
+
+		table := tablewriter.NewWriter(os.Stdout)
+		table.Append([]string{backupFile})
+		table.Render()
+
+		return err
 	},
 }
 

cmd/set.go

@@ -21,7 +21,12 @@ If the profile already exists it will be overwritten.
 
 		h, _ := cmd.Flags().GetString("host-file")
 
-		err := host.AddFromFile(from, h, profile, true)
+		err := host.AddFromFile(&host.AddFromFileOptions{
+			From:    from,
+			Dst:     h,
+			Profile: profile,
+			Reset:   true,
+		})
 		if err != nil {
 			return err
 		}
@@ -42,8 +47,51 @@ If the profile already exists it will be overwritten.
 	},
 }
 
+// setDomainsCmd represents the fromFile command
+var setDomainsCmd = &cobra.Command{
+	Use:   "domains",
+	Short: "Add content in your hosts file.",
+	Long: `
+Set content in your hosts file.
+If the profile already exists it will be added to it.`,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		src, _ := cmd.Flags().GetString("host-file")
+		ip, _ := cmd.Flags().GetString("ip")
+		profile, _ := cmd.Flags().GetString("profile")
+		if profile == "" {
+			profile = "other"
+		}
+
+		h, _ := cmd.Flags().GetString("host-file")
+
+		err := host.AddFromArgs(&host.AddFromArgsOptions{
+			Domains: args,
+			IP:      ip,
+			Dst:     h,
+			Profile: profile,
+			Reset:   true,
+		})
+		if err != nil {
+			return err
+		}
+
+		return host.ListProfiles(src, &host.ListOptions{
+			Profile: profile,
+		})
+	},
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		profile, _ := cmd.Flags().GetString("profile")
+
+		return host.ValidProfile(profile)
+	},
+}
+
 func init() {
 	rootCmd.AddCommand(setFromFileCmd)
 
 	setFromFileCmd.Flags().StringP("from", "f", "", "file to read")
+
+	setFromFileCmd.AddCommand(setDomainsCmd)
+
+	setDomainsCmd.Flags().String("ip", "127.0.0.1", "domains ip")
 }

pkg/host/add.go

@@ -5,31 +5,77 @@ import (
 	"os"
 )
 
+// commonAddOptions contains common options for adding.
+type commonAddOptions struct {
+	Dst     string
+	Profile string
+	Reset   bool
+}
+
+// AddFromFileOptions contains available options for adding from file.
+type AddFromFileOptions struct {
+	Dst     string
+	Profile string
+	Reset   bool
+	From string
+}
+
+// AddFromArgsOptions contains available options for adding from arguments.
+type AddFromArgsOptions struct {
+	Dst     string
+	Profile string
+	Reset   bool
+	Domains []string
+	IP      string
+}
+
 // AddFromFile reads content from a file and adds it as a profile into your hosts file.
 // If you pass reset=true it will delete all previous content of the profile.
-func AddFromFile(from, dst, profile string, reset bool) error {
-	if from == "" {
+func AddFromFile(opts *AddFromFileOptions) error {
+	if opts.From == "" {
 		return errors.New("missing source file")
 	}
-	if dst == "" {
+	newData, _ := ReadHostFileStrict(opts.From)
+
+	return add(newData, &commonAddOptions{
+		opts.Dst,
+		opts.Profile,
+		opts.Reset,
+	})
+}
+
+func AddFromArgs(opts *AddFromArgsOptions) error {
+	if len(opts.Domains) == 0 {
+		return errors.New("missing domains")
+	}
+	newData := ReadFromArgs(opts.Domains, opts.IP)
+
+	return add(newData, &commonAddOptions{
+		opts.Dst,
+		opts.Profile,
+		opts.Reset,
+	})
+}
+
+func add(n *hostFile, opts *commonAddOptions) error {
+	if opts.Dst == "" {
 		return errors.New("missing destination file")
 	}
-	if profile == "" {
-		profile = "default"
+	if opts.Profile == "" {
+		opts.Profile = "default"
 	}
 
-	currData, err := ReadHostFile(dst)
+	currData, err := ReadHostFile(opts.Dst)
 	if err != nil {
 		return err
 	}
-	newData, _ := ReadHostFileStrict(from)
 
-	if reset {
-		currData.profiles[profile] = hostLines{}
+	if opts.Reset {
+		currData.profiles[opts.Profile] = hostLines{}
 	}
-	currData.profiles[profile] = append(currData.profiles[profile], newData.profiles["default"]...)
+	currData.profiles[opts.Profile] = append(currData.profiles[opts.Profile], n.profiles["default"]...)
 
-	dstFile, err := os.OpenFile(dst, os.O_APPEND|os.O_CREATE|os.O_RDWR, 0644)
+	dstFile, err := os.OpenFile(opts.Dst, os.O_APPEND|os.O_CREATE|os.O_RDWR, 0644)
 	if err != nil {
 		return err
 	}

pkg/host/backup.go

@@ -9,21 +9,27 @@ import (
 )
 
 // BackupFile creates a copy of your hosts file to a new location with the date as extension
-func BackupFile(src, dstPath string) error {
+func BackupFile(src, dstPath string) (string, error) {
 	srcFile, err := os.Open(src)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer srcFile.Close()
 
-	bkpFilename := fmt.Sprintf("%s.%s", srcFile.Name(), time.Now().UTC().Format("20060102"))
-	bkpFilename = path.Join(dstPath, path.Base(bkpFilename))
+	bkpFilename := getBackupFilename(srcFile.Name(), dstPath, time.Now())
 	dst, err := os.Create(bkpFilename)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer dst.Close()
 
 	_, err = io.Copy(dst, srcFile)
-	return err
+	return bkpFilename, err
+}
+
+func getBackupFilename(srcFilename, dstPath string, t time.Time) string {
+	bkpFilename := fmt.Sprintf("%s.%s", srcFilename, t.UTC().Format("20060102"))
+	bkpFilename = path.Join(dstPath, path.Base(bkpFilename))
+
+	return bkpFilename
 }

pkg/host/backup_test.go

@@ -0,0 +1,28 @@
+package host
+
+import (
+	"testing"
+	"time"
+)
+
+func TestGetBackupFilename(t *testing.T) {
+	t.Run("Test backup filename", func(t *testing.T) {
+		nowTime := time.Now()
+		filename := getBackupFilename("/etc/hosts", "/tmp", nowTime)
+		if len(filename) == 0 {
+			t.Fatalf("filename in empty")
+		}
+		checkValue := "/tmp/hosts." + nowTime.UTC().Format("20060102")
+		if filename != checkValue {
+			t.Fatalf("unexpected filename; got %s; want %s", filename, checkValue)
+		}
+	})
+	t.Run("Test backup filename with last time", func(t *testing.T) {
+		nowTime := time.Now().Add(-24 * time.Hour)
+		filename := getBackupFilename("/tmp/path/fake_file", "/tmp/fake_path", nowTime)
+		checkValue := "/tmp/fake_path/fake_file." + nowTime.UTC().Format("20060102")
+		if filename != checkValue {
+			t.Fatalf("unexpected filename; got %s; want %s", filename, checkValue)
+		}
+	})
+}

pkg/host/host_file.go

@@ -85,6 +85,20 @@ func Read(r io.Reader, strict bool) (*hostFile, error) {
 	return h, nil
 }
 
+// ReadFromArgs read arguments into a hostFile struct
+func ReadFromArgs(domains []string, ip string) *hostFile {
+	dom := make([]string, len(domains))
+	for k, d := range domains {
+		dom[k] = fmt.Sprintf("%s\t%s", ip, d)
+	}
+	newData := &hostFile{
+		profiles: profileMap{
+			"default": dom,
+		},
+	}
+	return newData
+}
+
 // IsHostLine checks if a line is a host line or a comment line.
 func IsHostLine(line string) bool {
 	p := strings.Split(cleanLine(line), " ")

pkg/host/host_file_test.go

@@ -0,0 +1,27 @@
+package host
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestReadFromArgs(t *testing.T) {
+	f := func(domains []string, ip string) {
+		t.Helper()
+		data := ReadFromArgs(domains, ip)
+		if len(data.profiles["default"]) != len(domains) {
+			t.Fatalf("bad count of records; got %d; want %d", len(data.profiles["default"]), len(domains))
+		}
+		for k, d := range domains {
+			got := data.profiles["default"][k]
+			want := fmt.Sprintf("%s\t%s", ip, d)
+			if got != want {
+				t.Fatalf("unexpected record; got %s; want %s", got, want)
+			}
+		}
+	}
+	f([]string{}, "")
+	f([]string{"dom1.local"}, "")
+	f([]string{"dom1.local"}, "127.0.0.1")
+	f([]string{"dom3.local", "dom4.local"}, "localhost")
+}