From 38be845348f138b70b7a6e6f623bdfbc7a918171 Mon Sep 17 00:00:00 2001 From: MDK Date: Mon, 26 Jun 2023 15:58:04 +0800 Subject: [PATCH] finish upstream subcommand --- .gitignore | 1 + cmd/upstream.go | 30 ++++++++++++++++++++++++------ go.mod | 6 ++++-- go.sum | 1 + prober/rdns_prober.go | 9 ++++----- utils/other_utils.go | 13 +++++++++++++ utils/output_utils.go | 15 ++++++++++++--- 7 files changed, 59 insertions(+), 16 deletions(-) create mode 100644 .gitignore create mode 100644 utils/other_utils.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..32bcdf1 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/test \ No newline at end of file diff --git a/cmd/upstream.go b/cmd/upstream.go index e7d4071..5c315de 100644 --- a/cmd/upstream.go +++ b/cmd/upstream.go @@ -1,13 +1,17 @@ package cmd import ( + "dtool/prober" _ "dtool/prober" "dtool/utils" + _ "dtool/utils" + "errors" "github.com/spf13/cobra" ) var filename string +var output_file string var thread_num int var upstreamCmd = &cobra.Command{ Use: "upstream", @@ -17,16 +21,30 @@ var upstreamCmd = &cobra.Command{ input target can be added as an argument or as a file -f input file with limited ip addresses (limit=50) + -o output file default json type -t number of goroutine`, - Args: cobra.ExactArgs(1), - Run: func(cmd *cobra.Command, args []string) { - //prober.Get_upstream_ip(args[0]) - utils.SendTencentHttpdnsQuery() - }, + //Args: cobra.ExactArgs(1), + Run: upstream, +} + +func upstream(cmd *cobra.Command, args []string) { + if len(args) > 1 { + panic(errors.New("too many arguments!")) + } else if len(args) == 1 { + if utils.IsValidIP(args[0]) { + prober.Get_upstream_ip(args[0]) + } else { + panic(errors.New("invalid ip address")) + } + } else if len(args) == 0 { + prober.Get_upstream_file(filename, output_file, thread_num) + } } func init() { - upstreamCmd.Flags().StringVarP(&filename, "file", "f", "", "input filename") + upstreamCmd.Flags().StringVarP(&filename, "file", "f", "", "input file(optional)") + upstreamCmd.Flags().StringVarP(&output_file, "output", "o", "", "output file(optional)") + upstreamCmd.MarkFlagsRequiredTogether("file", "output") upstreamCmd.Flags().IntVarP(&thread_num, "threads", "t", 10, "number of concurrent threads") rootCmd.AddCommand(upstreamCmd) } diff --git a/go.mod b/go.mod index c08d21e..bdca553 100644 --- a/go.mod +++ b/go.mod @@ -2,11 +2,13 @@ module dtool go 1.20 -require github.com/spf13/cobra v1.7.0 +require ( + github.com/miekg/dns v1.1.54 + github.com/spf13/cobra v1.7.0 +) require ( github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/miekg/dns v1.1.54 // indirect github.com/spf13/pflag v1.0.5 // indirect golang.org/x/mod v0.7.0 // indirect golang.org/x/net v0.2.0 // indirect diff --git a/go.sum b/go.sum index 7fa2eff..937fa3d 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,7 @@ golang.org/x/mod v0.7.0 h1:LapD9S96VoQRhi/GrNTqeBJFrUjs5UHCAtTlgwA5oZA= golang.org/x/mod v0.7.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/net v0.2.0 h1:sZfSu1wtKLGlWI4ZZayP0ck9Y73K1ynO6gqzTdBVdPU= golang.org/x/net v0.2.0/go.mod h1:KqCZLdyyvdV855qA2rE3GC2aiw5xGR5TEjj8smXukLY= +golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o= golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A= golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/tools v0.3.0 h1:SrNbZl6ECOS1qFzgTdQfWXZM9XBkiA6tkFrH9YSTPHM= diff --git a/prober/rdns_prober.go b/prober/rdns_prober.go index 08d500c..8fae0e3 100644 --- a/prober/rdns_prober.go +++ b/prober/rdns_prober.go @@ -24,8 +24,7 @@ func retrieve_ip(pool chan string, filename string) { cnt := 0 f, err := os.Open(filename) if err != nil { - fmt.Printf("cannot open file %s\n", filename) - return + panic(err) } fmt.Println("sending msg ...") reader := bufio.NewReader(f) @@ -104,7 +103,7 @@ func store_data(pool chan Data, wg *sync.WaitGroup) { wg.Done() } -func Get_upstream_file(filename string, prober_num int) { +func Get_upstream_file(filename string, output string, prober_num int) { dataset = map[string][]string{} ip_pool := make(chan string, 500) data_pool := make(chan Data, 200) @@ -117,7 +116,7 @@ func Get_upstream_file(filename string, prober_num int) { probe_tasks.Wait() close(data_pool) store_task.Wait() - utils.OutputJSON(dataset) + utils.OutputJSON(dataset, output) } func Get_upstream_ip(ip string) { @@ -129,5 +128,5 @@ func Get_upstream_ip(ip string) { temp = append(temp, rdns) } dataset[data.target] = temp - utils.OutputJSON(dataset) + utils.OutputJSON(dataset, "-") } diff --git a/utils/other_utils.go b/utils/other_utils.go new file mode 100644 index 0000000..9940a0c --- /dev/null +++ b/utils/other_utils.go @@ -0,0 +1,13 @@ +package utils + +import ( + "net" +) + +func IsValidIP(ip string) bool { + res := net.ParseIP(ip) + if res == nil { + return false + } + return true +} diff --git a/utils/output_utils.go b/utils/output_utils.go index 88a649c..d914494 100644 --- a/utils/output_utils.go +++ b/utils/output_utils.go @@ -3,13 +3,22 @@ package utils import ( "encoding/json" "fmt" + "io/ioutil" ) -func OutputJSON(data map[string][]string) { +func OutputJSON(data interface{}, filename string) error { jsonstr, err := json.MarshalIndent(data, "", " ") if err != nil { fmt.Println("JSON encoding error:", err) - return + return err } - fmt.Println(string(jsonstr)) + if filename == "-" { + fmt.Println(string(jsonstr)) + } else { + err := ioutil.WriteFile(filename, jsonstr, 0666) + if err != nil { + return err + } + } + return nil }