use flag

Author: zyazhb

src/main.go

@@ -1,26 +1,32 @@
 package main
 
 import (
-	"os"
+	"flag"
 	"sync"
 	"time"
 
 	"github.com/sirupsen/logrus"
 )
 
 func main() {
+	domain := flag.String("d", "", "domain")
+	flag.Parse()
+	if *domain == "" {
+		logrus.Error("[-]Please specify a domain")
+		flag.Usage()
+		return
+	}
 	res := make(map[string]time.Duration)
 	reslock := sync.RWMutex{}
-	domain := os.Args[1]
-	logrus.Infof("[+]Testing domain %s...", domain)
-	iplist := RunLocalCore(domain, "")
+	logrus.Infof("[+]Testing domain %s...", *domain)
+	iplist := RunLocalCore(*domain, "")
 	wg := sync.WaitGroup{}
 	logrus.Info("[+]Testing delay...")
 	for _, ip := range iplist {
 		ip0 := ip
 		wg.Add(1)
 		go func() {
-			delay := Delay(domain, ip0)
+			delay := Delay(*domain, ip0)
 			logrus.Infof("[+]Delay: %s - %s", ip0, delay)
 			if delay != time.Duration(-1) {
 				reslock.Lock()