From 2240cad41951195981d4699aecc5739637d5ed9b Mon Sep 17 00:00:00 2001 From: Tom Hudson Date: Mon, 1 Jan 2018 15:08:14 +0000 Subject: [PATCH] Moves argument processing into its own function; fixes #16 --- args.go | 134 ++++++++++++++++++++++++++++++++++++++++++++++++++ main.go | 136 ++++++++++----------------------------------------- ratelimit.go | 2 +- 3 files changed, 160 insertions(+), 112 deletions(-) create mode 100644 args.go diff --git a/args.go b/args.go new file mode 100644 index 0000000..b2d5a14 --- /dev/null +++ b/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 Set the concurrency level (defaut: 20)\n" + h += " -d, --delay Milliseconds between requests to the same host (defaut: 5000)\n" + h += " -H, --header
Send a custom HTTP header\n" + h += " -s, --savestatus Save only responses with specific status code\n" + h += " -v, --verbose Verbose mode\n" + h += " -X, --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) + } +} diff --git a/main.go b/main.go index b40b755..9c8f24d 100644 --- a/main.go +++ b/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 Set the concurrency level (defaut: 20)\n" - h += " -d, --delay Milliseconds between requests to the same host (defaut: 5000)\n" - h += " -H, --header
Send a custom HTTP header\n" - h += " -s, --savestatus Save only responses with specific status code\n" - h += " -v, --verbose Verbose mode\n" - h += " -X, --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() +} diff --git a/ratelimit.go b/ratelimit.go index 8a5428a..1bdbcdf 100644 --- a/ratelimit.go +++ b/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