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