Added support for tracker links in paths, and filtering responses by header

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 <millis> Set the HTTP timeout (default: 10000)\n"
  h += " -v, --verbose Verbose mode\n"
  h += " -X, --method <method> HTTP method (default: GET)\n\n"
+ h += " -tracker <tracker link> Replaces {tracker} in paths with the specified tracker link\n"
+ h += " -sh --saveheader <header> Save only responses containing a specific header\n"
 
  h += "Defaults:\n"
  h += " pathsFile: ./paths\n"

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

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
+}