From 03d21308b4d39c6d22e34c225639cbfe65f71957 Mon Sep 17 00:00:00 2001
From: Tom Hudson <mail@tomnomnom.com>
Date: Sun, 7 Jan 2018 11:58:28 +0000
Subject: [PATCH] Adds support for a single host; fixes #20

---
 args.go | 10 +++-------
 main.go | 45 +++++++++++++++++++++++++++++++--------------
 2 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/args.go b/args.go
index 26e43f8..025fa04 100644
--- a/args.go
+++ b/args.go
@@ -68,28 +68,24 @@ func processArgs() config {
 	flag.BoolVar(&verbose, "verbose", false, "")
 	flag.BoolVar(&verbose, "v", false, "")
 
-	if verbose {
-		fmt.Println("sdgfsd")
-	}
-
 	flag.Parse()
 
 	// paths might be in a file, or it might be a single value
 	paths := flag.Arg(0)
 	if paths == "" {
-		paths = "paths"
+		paths = defaultPathsFile
 	}
 
 	// hosts are always in a file
 	hosts := flag.Arg(1)
 	if hosts == "" {
-		hosts = "hosts"
+		hosts = defaultHostsFile
 	}
 
 	// default the output directory to ./out
 	output := flag.Arg(2)
 	if output == "" {
-		output = "./out"
+		output = defaultOutputDir
 	}
 
 	// set the requester function to use
diff --git a/main.go b/main.go
index 1e6f097..29b74fd 100644
--- a/main.go
+++ b/main.go
@@ -10,7 +10,14 @@ import (
 	"time"
 )
 
-const userAgent = "Mozilla/5.0 (compatible; meg/0.1; +https://github.com/tomnomnom/meg)"
+const (
+	userAgent = "Mozilla/5.0 (compatible; meg/0.2; +https://github.com/tomnomnom/meg)"
+
+	// argument defaults
+	defaultPathsFile = "./paths"
+	defaultHostsFile = "./hosts"
+	defaultOutputDir = "./out"
+)
 
 // a requester is a function that makes HTTP requests
 type requester func(request) response
@@ -20,22 +27,15 @@ func main() {
 	// get the config struct
 	c := processArgs()
 
-	// if the paths argument is a file, read it; otherwise
-	// treat it as a literal value
-	var paths []string
-	if isFile(c.paths) {
-		lines, err := readLines(c.paths)
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "failed to open paths file: %s\n", err)
-			os.Exit(1)
-		}
-		paths = lines
-	} else if c.paths != "paths" {
-		paths = []string{c.paths}
+	// read the paths file
+	paths, err := readLinesOrLiteral(c.paths, defaultPathsFile)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to open paths file: %s\n", err)
+		os.Exit(1)
 	}
 
 	// read the hosts file
-	hosts, err := readLines(c.hosts)
+	hosts, err := readLinesOrLiteral(c.hosts, defaultHostsFile)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to open hosts file: %s\n", err)
 		os.Exit(1)
@@ -166,6 +166,23 @@ func readLines(filename string) ([]string, error) {
 	return lines, sc.Err()
 }
 
+// readLinesOrLiteral tries to read lines from a file, returning
+// the arg in a string slice if the file doesn't exist, unless
+// the arg matches its default value
+func readLinesOrLiteral(arg, argDefault string) ([]string, error) {
+	if isFile(arg) {
+		return readLines(arg)
+	}
+
+	// if the argument isn't a file, but it is the default, don't
+	// treat it as a literal value
+	if arg == argDefault {
+		return []string{}, fmt.Errorf("file %s not found", arg)
+	}
+
+	return []string{arg}, nil
+}
+
 // isFile returns true if its argument is a regular file
 func isFile(path string) bool {
 	f, err := os.Stat(path)