diff --git a/args.go b/args.go index 90cca52..24916f8 100644 --- a/args.go +++ b/args.go @@ -41,6 +41,8 @@ func (s saveStatusArgs) Includes(search int) bool { } type config struct { + trackerLink string + saveHeader string body string concurrency int delay int @@ -61,6 +63,10 @@ type config struct { func processArgs() config { + // tracker link params + trackerLink := "" + flag.StringVar(&trackerLink, "tracker", "", "") + // body param body := "" flag.StringVar(&body, "body", "", "") @@ -96,6 +102,11 @@ func processArgs() config { flag.Var(&saveStatus, "savestatus", "") flag.Var(&saveStatus, "s", "") + // saveheader parms + saveHeader := "" + flag.StringVar(&saveHeader, "saveheader", "", "") + flag.StringVar(&saveHeader, "sh", "", "") + // timeout param timeout := 10000 flag.IntVar(&timeout, "timeout", 10000, "") @@ -142,6 +153,8 @@ func processArgs() config { } return config{ + trackerLink: trackerLink, + saveHeader: saveHeader, body: body, concurrency: concurrency, delay: delay, @@ -177,6 +190,8 @@ func init() { h += " -t, --timeout Set the HTTP timeout (default: 10000)\n" h += " -v, --verbose Verbose mode\n" h += " -X, --method HTTP method (default: GET)\n\n" + h += " -tracker Replaces {tracker} in paths with the specified tracker link\n" + h += " -sh --saveheader
Save only responses containing a specific header\n" h += "Defaults:\n" h += " pathsFile: ./paths\n" diff --git a/main.go b/main.go index f2a73a2..1188f0f 100644 --- a/main.go +++ b/main.go @@ -8,6 +8,7 @@ import ( "path/filepath" "sync" "time" + "strings" ) const ( @@ -88,6 +89,10 @@ func main() { continue } + if c.saveHeader != "" && !res.hasHeader(c.saveHeader) { + continue + } + if res.err != nil { fmt.Fprintf(os.Stderr, "request failed: %s\n", res.err) continue @@ -115,10 +120,16 @@ func main() { // so we should strip that off and add it to // the beginning of the path. u, err := url.Parse(host) + if err != nil { fmt.Fprintf(os.Stderr, "failed to parse host: %s\n", err) continue } + + if len(c.trackerLink) > 0 { + path = replaceTracker(path, u.Host, c.trackerLink) + } + prefixedPath := u.Path + path u.Path = "" @@ -136,7 +147,7 @@ func main() { timeout: time.Duration(c.timeout * 1000000), } } - } + } // once all of the requests have been sent we can // close the requests channel @@ -169,6 +180,19 @@ func readLines(filename string) ([]string, error) { return lines, sc.Err() } +func replaceTracker(path string, host string, trackerLink string) string { + now := time.Now() + seconds := now.Unix() + + host = strings.Split(host, ":")[0] + + tracker := fmt.Sprintf("%d.%s.%s", seconds, host, trackerLink) + + newPath := strings.Replace(path, "{tracker}", tracker, -1) + + return newPath +} + // 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 diff --git a/response.go b/response.go index 1ca0d0a..e3f94cd 100644 --- a/response.go +++ b/response.go @@ -7,6 +7,7 @@ import ( "io/ioutil" "os" "path" + "strings" ) // a response is a wrapper around an HTTP response; @@ -89,3 +90,12 @@ func (r response) save(pathPrefix string, noHeaders bool) (string, error) { return p, nil } + +func (r response) hasHeader(header string) bool { + for _, item := range r.headers { + if strings.TrimSpace(item) == header { + return true + } + } + return false +}