the project structure modified and new features added

4 files changed, 53 insertions, 23 deletions

diff --git a/cmd/cache.go b/cmd/cache.go
index c72af16..5e17fdd 100644
--- a/cmd/cache.go
+++ b/cmd/cache.go
@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"dtool/prober"
-	"dtool/scheduler"
 	"dtool/utils"
 
 	"github.com/spf13/cobra"
@@ -11,6 +10,8 @@ import (
 var query_cnt int
 var inputfile string
 var outputfile string
+var goroutine_num int
+var controlled_domain string
 var cacheCmd = &cobra.Command{
 	Use:   "cache",
 	Short: "cache related test",
@@ -24,14 +25,16 @@ func cache_test(cmd *cobra.Command, args []string) {
 			prober.RecursiveCacheTest(args[0], query_cnt)
 		}
 	} else {
-		scheduler.CreateTask(prober.RecursiveCacheProbe, inputfile, outputfile, 10)
+		prober.CreateTask(prober.RecursiveCacheProbe, controlled_domain, inputfile, outputfile, goroutine_num)
 	}
 }
 
 func init() {
+	cacheCmd.Flags().StringVarP(&controlled_domain, "domain", "d", "echodns.xyz", "controlled domain")
 	cacheCmd.Flags().StringVarP(&inputfile, "input", "i", "", "input file(optional)")
 	cacheCmd.Flags().StringVarP(&outputfile, "output", "o", "", "output file(optional)")
 	cacheCmd.MarkFlagsRequiredTogether("input", "output")
 	cacheCmd.Flags().IntVarP(&query_cnt, "num", "n", 20, "number of queries in one test")
+	cacheCmd.Flags().IntVarP(&goroutine_num, "concurrency", "t", 150, "number of goroutine")
 	rootCmd.AddCommand(cacheCmd)
 }
diff --git a/cmd/record.go b/cmd/record.go
new file mode 100644
index 0000000..9ede8d7
--- /dev/null
+++ b/cmd/record.go
@@ -0,0 +1,47 @@
+package cmd
+
+import (
+	"dtool/prober"
+	"dtool/utils"
+
+	"fmt"
+
+	"github.com/spf13/cobra"
+)
+
+var record_input string
+var record_output string
+var record_type string
+var record_domain string
+var recordCmd = &cobra.Command{
+	Use:   "record",
+	Short: "get specific record response",
+	Long:  "get specific record response",
+	Run:   record_probe,
+}
+
+func record_probe(cmd *cobra.Command, args []string) {
+	if len(args) == 1 {
+		if utils.IsValidIP(args[0]) {
+			result, err := prober.SVCBProbeOnce(args[0], record_domain)
+			if err == nil {
+				if output_str, err := prober.OutputHandler(result); err == nil {
+					fmt.Println(output_str)
+				}
+			} else {
+				fmt.Println(err)
+			}
+		}
+	} else {
+		prober.CreateTask(prober.SVCBProbe, record_domain, record_input, record_output, 500)
+	}
+}
+
+func init() {
+	recordCmd.Flags().StringVarP(&record_input, "input", "i", "", "input file(optional)")
+	recordCmd.Flags().StringVarP(&record_output, "output", "o", "", "output file(optional)")
+	recordCmd.MarkFlagsRequiredTogether("input", "output")
+	recordCmd.Flags().StringVarP(&record_type, "type", "t", "A", "request record type")
+	recordCmd.Flags().StringVarP(&record_domain, "domain", "d", "example.com", "requested domain")
+	rootCmd.AddCommand(recordCmd)
+}
diff --git a/cmd/upstream.go b/cmd/upstream.go
index c40528b..a9f40ef 100644
--- a/cmd/upstream.go
+++ b/cmd/upstream.go
@@ -27,7 +27,7 @@ input target can be added as an argument or as a file
 
 func upstream(cmd *cobra.Command, args []string) {
 	if len(args) > 1 {
-		panic(errors.New("too many arguments!"))
+		panic(errors.New("too many arguments"))
 	} else if len(args) == 1 {
 		if utils.IsValidIP(args[0]) {
 			prober.Get_upstream_ip(args[0])
diff --git a/cmd/version.go b/cmd/version.go
deleted file mode 100644
index 3c6712b..0000000
--- a/cmd/version.go
+++ /dev/null
@@ -1,20 +0,0 @@
-package cmd
-
-import (
-	"github.com/spf13/cobra"
-)
-
-var versionCmd = &cobra.Command{
-	Use:   "version",
-	Short: "get server version with version.bind",
-	Long:  "get server version with version.bind chaos txt request",
-	Run:   version,
-}
-
-func version(cmd *cobra.Command, args []string) {
-
-}
-
-func init() {
-	rootCmd.AddCommand(versionCmd)
-}