Moves argument processing into its own function; fixes #16

args.go

@@ -0,0 +1,134 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+)
+
+// headerArgs is the type used to store the header arguments
+type headerArgs []string
+
+func (h *headerArgs) Set(val string) error {
+	*h = append(*h, val)
+	return nil
+}
+
+func (h headerArgs) String() string {
+	return "string"
+}
+
+type config struct {
+	concurrency int
+	delay       int
+	headers     headerArgs
+	method      string
+	saveStatus  int
+	verbose     bool
+	suffix      string
+	prefix      string
+	output      string
+}
+
+func processArgs() config {
+
+	// concurrency param
+	concurrency := 20
+	flag.IntVar(&concurrency, "concurrency", 20, "")
+	flag.IntVar(&concurrency, "c", 20, "")
+
+	// delay param
+	delay := 5000
+	flag.IntVar(&delay, "delay", 5000, "")
+	flag.IntVar(&delay, "d", 5000, "")
+
+	// headers param
+	var headers headerArgs
+	flag.Var(&headers, "header", "")
+	flag.Var(&headers, "H", "")
+
+	// method param
+	method := "GET"
+	flag.StringVar(&method, "method", "GET", "")
+	flag.StringVar(&method, "X", "GET", "")
+
+	// savestatus param
+	saveStatus := 0
+	flag.IntVar(&saveStatus, "savestatus", 0, "")
+	flag.IntVar(&saveStatus, "s", 0, "")
+
+	// verbose param
+	verbose := false
+	flag.BoolVar(&verbose, "verbose", false, "")
+	flag.BoolVar(&verbose, "v", false, "")
+
+	flag.Parse()
+
+	// suffixes might be in a file, or it might be a single value
+	suffix := flag.Arg(0)
+	if suffix == "" {
+		suffix = "suffixes"
+	}
+
+	// prefixes are always in a file
+	prefix := flag.Arg(1)
+	if prefix == "" {
+		prefix = "prefixes"
+	}
+
+	// default the output directory to ./out
+	output := flag.Arg(2)
+	if output == "" {
+		output = "./out"
+	}
+
+	return config{
+		concurrency: concurrency,
+		delay:       delay,
+		headers:     headers,
+		method:      method,
+		saveStatus:  saveStatus,
+		verbose:     verbose,
+		suffix:      suffix,
+		prefix:      prefix,
+		output:      output,
+	}
+}
+
+func init() {
+	flag.Usage = func() {
+		h := "Request many paths (suffixes) for many hosts (prefixes)\n\n"
+
+		h += "Usage:\n"
+		h += "  meg [suffix|suffixFile] [prefixFile] [outputDir]\n\n"
+
+		h += "Options:\n"
+		h += "  -c, --concurrency <val>    Set the concurrency level (defaut: 20)\n"
+		h += "  -d, --delay <val>          Milliseconds between requests to the same host (defaut: 5000)\n"
+		h += "  -H, --header <header>      Send a custom HTTP header\n"
+		h += "  -s, --savestatus <status>  Save only responses with specific status code\n"
+		h += "  -v, --verbose              Verbose mode\n"
+		h += "  -X, --method <method>      HTTP method (default: GET)\n\n"
+
+		h += "Defaults:\n"
+		h += "  suffixFile: ./suffixes\n"
+		h += "  prefixFile: ./prefixes\n"
+		h += "  outputDir:  ./out\n\n"
+
+		h += "Suffix file format:\n"
+		h += "  /robots.txt\n"
+		h += "  /package.json\n"
+		h += "  /security.txt\n\n"
+
+		h += "Prefix file format:\n"
+		h += "  http://example.com\n"
+		h += "  https://example.edu\n"
+		h += "  https://example.net\n\n"
+
+		h += "Examples:\n"
+		h += "  meg /robots.txt\n"
+		h += "  meg hosts.txt paths.txt output\n"
+
+		fmt.Fprintf(os.Stderr, h)
+	}
+}

main.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"bufio"
-	"flag"
 	"fmt"
 	"net/url"
 	"os"
@@ -33,94 +32,41 @@ type response struct {
 // a requester is a function that makes HTTP requests
 type requester func(request) response
 
-type headerArgs []string
-
-func (h *headerArgs) Set(val string) error {
-	*h = append(*h, val)
-	return nil
-}
-
-func (h headerArgs) String() string {
-	return "string"
-}
-
 func main() {
 
-	// concurrency param
-	concurrency := 20
-	flag.IntVar(&concurrency, "concurrency", 20, "")
-	flag.IntVar(&concurrency, "c", 20, "")
-
-	// delay param
-	delay := 5000
-	flag.IntVar(&delay, "delay", 5000, "")
-	flag.IntVar(&delay, "d", 5000, "")
-
-	// headers param
-	var headers headerArgs
-	flag.Var(&headers, "header", "")
-	flag.Var(&headers, "H", "")
+	// get the config struct
+	c := processArgs()
 
-	// method param
-	method := "GET"
-	flag.StringVar(&method, "method", "GET", "")
-	flag.StringVar(&method, "X", "GET", "")
-
-	// savestatus param
-	saveStatus := 0
-	flag.IntVar(&saveStatus, "savestatus", 0, "")
-	flag.IntVar(&saveStatus, "s", 0, "")
-
-	// verbose param
-	verbose := false
-	flag.BoolVar(&verbose, "verbose", false, "")
-	flag.BoolVar(&verbose, "v", false, "")
-
-	flag.Parse()
-
-	// suffixes might be in a file, or it might be a single value
-	suffixArg := flag.Arg(0)
-	if suffixArg == "" {
-		suffixArg = "suffixes"
-	}
+	// if the suffix argument is a file, read it; otherwise
+	// treat it as a literal value
 	var suffixes []string
-	if f, err := os.Stat(suffixArg); err == nil && f.Mode().IsRegular() {
-		lines, err := readLines(suffixArg)
+	if isFile(c.suffix) {
+		lines, err := readLines(c.suffix)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "failed to open suffix file: %s\n", err)
 			os.Exit(1)
 		}
 		suffixes = lines
-	} else if suffixArg != "suffixes" {
-		// don't treat the default suffixes filename as a literal value
-		suffixes = []string{suffixArg}
+	} else if c.suffix != "suffixes" {
+		suffixes = []string{c.suffix}
 	}
 
-	// prefixes are always in a file
-	prefixFile := flag.Arg(1)
-	if prefixFile == "" {
-		prefixFile = "prefixes"
-	}
-	prefixes, err := readLines(prefixFile)
+	// read the prefix file
+	prefixes, err := readLines(c.prefix)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to open prefix file: %s\n", err)
 		os.Exit(1)
 	}
 
-	// default the output directory to ./out
-	outputDir := flag.Arg(2)
-	if outputDir == "" {
-		outputDir = "./out"
-	}
-
-	err = os.MkdirAll(outputDir, 0750)
+	// make the output directory
+	err = os.MkdirAll(c.output, 0750)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to create output directory: %s\n", err)
 		os.Exit(1)
 	}
 
 	// open the index file
-	indexFile := filepath.Join(outputDir, "index")
+	indexFile := filepath.Join(c.output, "index")
 	index, err := os.OpenFile(indexFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to open index file for writing: %s\n", err)
@@ -129,7 +75,7 @@ func main() {
 
 	// set up a rate limiter
 	rl := &rateLimiter{
-		delay: time.Duration(delay * 1000000),
+		delay: time.Duration(c.delay * 1000000),
 		reqs:  make(map[string]time.Time),
 	}
 
@@ -140,7 +86,7 @@ func main() {
 
 	// spin up some workers to do the requests
 	var wg sync.WaitGroup
-	for i := 0; i < concurrency; i++ {
+	for i := 0; i < c.concurrency; i++ {
 		wg.Add(1)
 
 		go func() {
@@ -158,18 +104,18 @@ func main() {
 	owg.Add(1)
 	go func() {
 		for res := range responses {
-			if saveStatus != 0 && saveStatus != res.statusCode {
+			if c.saveStatus != 0 && c.saveStatus != res.statusCode {
 				continue
 			}
 
-			path, err := res.save(outputDir)
+			path, err := res.save(c.output)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "failed to save file: %s\n", err)
 			}
 
 			line := fmt.Sprintf("%s %s (%s)\n", path, res.request.url, res.status)
 			fmt.Fprintf(index, "%s", line)
-			if verbose {
+			if c.verbose {
 				fmt.Printf("%s", line)
 			}
 		}
@@ -184,7 +130,7 @@ func main() {
 				fmt.Printf("failed to parse url: %s\n", err)
 				continue
 			}
-			requests <- request{method: method, url: u, headers: headers}
+			requests <- request{method: c.method, url: u, headers: c.headers}
 		}
 	}
 
@@ -201,44 +147,6 @@ func main() {
 
 }
 
-func init() {
-	flag.Usage = func() {
-		h := "Request many paths (suffixes) for many hosts (prefixes)\n\n"
-
-		h += "Usage:\n"
-		h += "  meg [suffix|suffixFile] [prefixFile] [outputDir]\n\n"
-
-		h += "Options:\n"
-		h += "  -c, --concurrency <val>    Set the concurrency level (defaut: 20)\n"
-		h += "  -d, --delay <val>          Milliseconds between requests to the same host (defaut: 5000)\n"
-		h += "  -H, --header <header>      Send a custom HTTP header\n"
-		h += "  -s, --savestatus <status>  Save only responses with specific status code\n"
-		h += "  -v, --verbose              Verbose mode\n"
-		h += "  -X, --method <method>      HTTP method (default: GET)\n\n"
-
-		h += "Defaults:\n"
-		h += "  suffixFile: ./suffixes\n"
-		h += "  prefixFile: ./prefixes\n"
-		h += "  outputDir:  ./out\n\n"
-
-		h += "Suffix file format:\n"
-		h += "  /robots.txt\n"
-		h += "  /package.json\n"
-		h += "  /security.txt\n\n"
-
-		h += "Prefix file format:\n"
-		h += "  http://example.com\n"
-		h += "  https://example.edu\n"
-		h += "  https://example.net\n\n"
-
-		h += "Examples:\n"
-		h += "  meg /robots.txt\n"
-		h += "  meg hosts.txt paths.txt output\n"
-
-		fmt.Fprintf(os.Stderr, h)
-	}
-}
-
 // readLines reads all of the lines from a text file in to
 // a slice of strings, returning the slice and any error
 func readLines(filename string) ([]string, error) {
@@ -256,3 +164,9 @@ func readLines(filename string) ([]string, error) {
 
 	return lines, sc.Err()
 }
+
+// isFile returns true if its argument is a regular file
+func isFile(path string) bool {
+	f, err := os.Stat(path)
+	return err == nil && f.Mode().IsRegular()
+}

ratelimit.go

@@ -20,7 +20,7 @@ func (r *rateLimiter) Block(u *url.URL) {
 
 	// if there's nothing in the map we can
 	// return straight away
-	if _, ok := r.reqs[u.Hostname()]; !ok {
+	if _, ok := r.reqs[key]; !ok {
 		r.reqs[key] = now
 		r.Unlock()
 		return