Added output directory flag

Author: rtmor

README.md

@@ -11,14 +11,16 @@ Program will resolve each domain and report execution time statistics.
 ### usage
 
 ```
-Usage: ./main [options] -ns {nameserver}
+Usage: ./dns-subnet-client [options] -ns {nameserver}
 
   -client string
         set edns client-subnet option
   -d string
         dns query wordlists file
   -ns string
         set preferred nameserver (default "8.8.8.8")
+  -o string
+        output directory for data graph (default "data")
   -t int
         number of threads (default 100)
   -v    enable verbose output of dns queries (debug)

graph/graph.go

@@ -10,7 +10,8 @@ import (
 )
 
 // BuildGraph initializes new 2-axis graph
-func BuildGraph(nameserver string, clientStatus bool, t, c []float64, threads, dmnCount int) {
+func BuildGraph(nameserver string, clientStatus bool, t, c []float64,
+	threads, dmnCount int, output string) {
 	mainSeries := chart.ContinuousSeries{
 		Name:    "Rate",
 		XValues: t,
@@ -25,7 +26,7 @@ func BuildGraph(nameserver string, clientStatus bool, t, c []float64, threads, d
 	} // we can optionally set the `WindowSize` property which alters how the moving average is calculated.
 
 	graph := chart.Chart{
-		Title: fmt.Sprintf("%v - [subnet_client: %v, thread_count:%v, domain_list:%v",
+		Title: fmt.Sprintf("%v - +subnet_client: %v, +thread_count:%v, +domain_list:%v",
 			nameserver, clientStatus, threads, dmnCount),
 		TitleStyle: chart.Style{
 			FontSize: 12.0,
@@ -75,9 +76,12 @@ func BuildGraph(nameserver string, clientStatus bool, t, c []float64, threads, d
 		chart.Legend(&graph),
 	}
 
-	outputDir := "data"
+	if _, err := os.Stat(output); os.IsNotExist(err) {
+		os.Mkdir(output, os.ModePerm)
+	}
+
 	f, err := os.Create(fmt.Sprintf("%v/ns-%v_client-%v_%4v.png",
-		outputDir, nameserver, clientStatus, time.Now().Unix()))
+		output, nameserver, clientStatus, time.Now().Unix()))
 	if err != nil {
 		log.Printf("Error writing to file\n%v", err)
 	}

main.go

@@ -19,6 +19,7 @@ var (
 	nameserver = flag.String("ns", "8.8.8.8", "set preferred nameserver")
 	threads    = flag.Int("t", 100, "number of threads")
 	verbose    = flag.Bool("v", false, "enable verbose output of dns queries (debug)")
+	output     = flag.String("o", "data", "output directory for data graph")
 )
 
 type statistics struct {
@@ -125,10 +126,13 @@ func updateStats(done chan bool) {
 
 func finalStats() {
 	graph.BuildGraph(*nameserver, len(*client) != 0, timeValues,
-		rateValues, *threads, domainCount)
+		rateValues, *threads, domainCount, *output)
 
-	fmt.Printf("\n\nFinal Statistics\nAttempts: %v\nSuccess: %v\nFailed: %v\n\n"+
-		"Avg Rate: %v queries/sec",
+	fmt.Printf("\n\nFinal Statistics\n"+
+		"Attempts:   %v\n"+
+		"Success:    %v\n"+
+		"Failed:     %v\n\n"+
+		"Avg Rate:   %v queries/s",
 		Stats.attempts, Stats.success, Stats.fail, getStatAvg())
 }