added function that allows to save output into file

indigo-sadland · dwisiswant0 · commit da722db5682d · 2022-01-15T20:52:18.000+07:00
diff --git a/internal/executor/executor.go b/internal/executor/executor.go
@@ -1,24 +1,51 @@
 package executor
 
 import (
+	"bufio"
 	"fmt"
+	"github.com/projectdiscovery/gologger"
+	"os"
+	"sync"
 
 	"github.com/logrusorgru/aurora"
-	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/retryabledns"
 	"github.com/pwnesia/dnstake/internal/errors"
 	"github.com/pwnesia/dnstake/internal/option"
 	"github.com/pwnesia/dnstake/pkg/dnstake"
 	"github.com/pwnesia/dnstake/pkg/fingerprint"
 )
 
+var mu sync.Mutex
+
+// WriteToFile writes output data into specified file.
+func WriteToFile(data, outputFile string) {
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	file, err := os.OpenFile(outputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		gologger.Error().Msg(err.Error())
+	}
+
+	wrt := bufio.NewWriter(file)
+
+	_, err = wrt.WriteString(data + "\n")
+	if err != nil {
+		gologger.Error().Msg(err.Error())
+	}
+
+	wrt.Flush()
+	file.Close()
+}
+
 // New to execute target hostname
 func New(opt *option.Options, hostname string) {
 	var out = ""
 
 	vuln, DNS, err := exec(hostname)
 	if err != nil {
-		gologger.Error().Msgf("%s: %s", hostname, err.Error())
+		out += fmt.Sprintf("[%s] %s: %s", aurora.Red("ERR"), hostname, err.Error())
 	}
 
 	if vuln {
@@ -43,8 +70,14 @@ func New(opt *option.Options, hostname string) {
 			}
 		}
 
-		fmt.Println(out)
 	}
+
+	if opt.Output != "" {
+		WriteToFile(out, opt.Output)
+	}
+
+	fmt.Println(out)
+
 }
 
 func exec(hostname string) (bool, fingerprint.DNS, error) {
diff --git a/internal/option/const.go b/internal/option/const.go
@@ -18,10 +18,12 @@ const (
   -t, --target <HOST/FILE>    Define single target host/list to check
   -c, --concurrent <i>        Set the concurrency level (default: 25)
   -s, --silent                Suppress errors and/or clean output
+  -o, --output                Define path where you want to store the tool's output
   -h, --help                  Display its help`
   examples = `
   dnstake -t (sub.)domain.tld
   dnstake -t hosts.txt
+  dnstake -t hosts.txt -o ./dnstake.out
   cat hosts.txt | dnstake
   subfinder -silent -d domain.tld | dnstake
   `
diff --git a/internal/option/options.go b/internal/option/options.go
@@ -2,8 +2,9 @@ package option
 
 // Options define available flag/options
 type Options struct {
-	Target      string
-	Concurrency int
-	Silent      bool
-	List        []string
+	Target      	string
+	Concurrency 	int
+	Silent      	bool
+	List        	[]string
+	Output			string
 }
diff --git a/internal/option/parse.go b/internal/option/parse.go
@@ -20,6 +20,9 @@ func Parse() *Options {
 	flag.BoolVar(&opt.Silent, "s", false, "")
 	flag.BoolVar(&opt.Silent, "silent", false, "")
 
+	flag.StringVar(&opt.Output, "o", "", "")
+	flag.StringVar(&opt.Output, "output", "", "")
+
 	flag.Usage = func() {
 		showBanner()
 		h := []string{
