diff --git a/cmd/all/all.go b/cmd/all/all.go index 3efe5ae..1e06686 100644 --- a/cmd/all/all.go +++ b/cmd/all/all.go @@ -2,12 +2,12 @@ package all import ( "net" - "os" command "github.com/esonhugh/k8spider/cmd" "github.com/esonhugh/k8spider/define" "github.com/esonhugh/k8spider/pkg" "github.com/esonhugh/k8spider/pkg/mutli" + "github.com/esonhugh/k8spider/pkg/printer" "github.com/esonhugh/k8spider/pkg/scanner" "github.com/miekg/dns" log "github.com/sirupsen/logrus" @@ -27,60 +27,48 @@ var AllCmd = &cobra.Command{ log.Warn("cidr is required") return } + // Wildcard + records := scanner.DumpWildCard(command.Opts.Zone) + if records != nil { + printer.PrintResult(records, command.Opts.OutputFile) + } + // AXFR Dumping records, err := scanner.DumpAXFR(dns.Fqdn(command.Opts.Zone), "ns.dns."+command.Opts.Zone+":53") if err == nil { - printResult(records) - return - } - log.Errorf("Transfer failed: %v", err) - records = scanner.DumpWildCard(command.Opts.Zone) - if records != nil && len(records) > 0 { - printResult(records) - return + printer.PrintResult(records, command.Opts.OutputFile) + } else { + log.Errorf("Transfer failed: %v", err) } - log.Errorf("WildCard dns dump failed: %v", err) - + // Service Discovery ipNets, err := pkg.ParseStringToIPNet(command.Opts.Cidr) if err != nil { log.Warnf("ParseStringToIPNet failed: %v", err) return } - if command.Opts.BatchMode { - RunBatch(ipNets) + var finalRecord define.Records + if command.Opts.MultiThreadingMode { + finalRecord = RunMultiThread(ipNets, command.Opts.ThreadingNum) } else { - Run(ipNets) + finalRecord = Run(ipNets) } + printer.PrintResult(finalRecord, command.Opts.OutputFile) }, } -func Run(net *net.IPNet) { +func Run(net *net.IPNet) (finalRecord define.Records) { var records define.Records = scanner.ScanSubnet(net) if records == nil || len(records) == 0 { log.Warnf("ScanSubnet Found Nothing") return } records = scanner.ScanSvcForPorts(records) - printResult(records) + return } -func RunBatch(net *net.IPNet) { - scan := mutli.ScanAll(net) - var finalRecord []define.Record +func RunMultiThread(net *net.IPNet, count int) (finalRecord define.Records) { + scan := mutli.ScanAll(net, count) for r := range scan { finalRecord = append(finalRecord, r...) } - printResult(finalRecord) -} - -func printResult(records define.Records) { - if command.Opts.OutputFile != "" { - f, err := os.OpenFile(command.Opts.OutputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - log.Warnf("OpenFile failed: %v", err) - } - defer f.Close() - records.Print(log.StandardLogger().Writer(), f) - } else { - records.Print(log.StandardLogger().Writer()) - } + return } diff --git a/cmd/axfr/axfr.go b/cmd/axfr/axfr.go index 5ba568b..a9d12a7 100644 --- a/cmd/axfr/axfr.go +++ b/cmd/axfr/axfr.go @@ -1,11 +1,11 @@ package axfr import ( - "os" "strings" command "github.com/esonhugh/k8spider/cmd" "github.com/esonhugh/k8spider/define" + "github.com/esonhugh/k8spider/pkg/printer" "github.com/esonhugh/k8spider/pkg/scanner" "github.com/miekg/dns" log "github.com/sirupsen/logrus" @@ -41,16 +41,6 @@ var AxfrCmd = &cobra.Command{ log.Errorf("Transfer failed: %v", err) return } - if command.Opts.OutputFile != "" { - f, err := os.OpenFile(command.Opts.OutputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - log.Warnf("OpenFile failed: %v", err) - } - defer f.Close() - records.Print(log.StandardLogger().Writer(), f) - } else { - records.Print(log.StandardLogger().Writer()) - } - + printer.PrintResult(records, command.Opts.OutputFile) }, } diff --git a/cmd/root.go b/cmd/root.go index 3de8515..2f4fef6 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -19,7 +19,8 @@ var Opts = struct { OutputFile string Verbose string - BatchMode bool + MultiThreadingMode bool + ThreadingNum int }{} func init() { @@ -29,7 +30,8 @@ func init() { RootCmd.PersistentFlags().StringVarP(&Opts.Zone, "zone", "z", "cluster.local", "zone") RootCmd.PersistentFlags().StringVarP(&Opts.OutputFile, "output-file", "o", "", "output file") RootCmd.PersistentFlags().StringVarP(&Opts.Verbose, "verbose", "v", "info", "log level (debug,info,trace,warn,error,fatal,panic)") - RootCmd.PersistentFlags().BoolVarP(&Opts.BatchMode, "batch-mode", "b", false, "batch mode") + RootCmd.PersistentFlags().BoolVarP(&Opts.MultiThreadingMode, "thread", "t", false, "multi threading mode, work pair with -n") + RootCmd.PersistentFlags().IntVarP(&Opts.ThreadingNum, "thread-num", "n", 16, "threading num, default 16") } var RootCmd = &cobra.Command{ diff --git a/cmd/service/service.go b/cmd/service/service.go index cc40416..80e2381 100644 --- a/cmd/service/service.go +++ b/cmd/service/service.go @@ -2,10 +2,10 @@ package service import ( "fmt" - "os" command "github.com/esonhugh/k8spider/cmd" "github.com/esonhugh/k8spider/define" + "github.com/esonhugh/k8spider/pkg/printer" "github.com/esonhugh/k8spider/pkg/scanner" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" @@ -17,7 +17,7 @@ func init() { var ServiceCmd = &cobra.Command{ Use: "service", - Short: "service is a tool to discover k8s services ports", + Short: "service is a tool to discover k8s services", Run: func(cmd *cobra.Command, args []string) { if command.Opts.Zone == "" || command.Opts.SvcDomains == nil || len(command.Opts.SvcDomains) == 0 { log.Warn("zone can't empty and svc-domains can't empty") @@ -28,15 +28,6 @@ var ServiceCmd = &cobra.Command{ records = append(records, define.Record{SvcDomain: fmt.Sprintf("%s.svc.%s", domain, command.Opts.Zone)}) } records = scanner.ScanSvcForPorts(records) - if command.Opts.OutputFile != "" { - f, err := os.OpenFile(command.Opts.OutputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - log.Warnf("OpenFile failed: %v", err) - } - defer f.Close() - records.Print(log.StandardLogger().Writer(), f) - } else { - records.Print(log.StandardLogger().Writer()) - } + printer.PrintResult(records, command.Opts.OutputFile) }, } diff --git a/cmd/subnet/subnet.go b/cmd/subnet/subnet.go index 16afd29..2804889 100644 --- a/cmd/subnet/subnet.go +++ b/cmd/subnet/subnet.go @@ -2,12 +2,12 @@ package subnet import ( "net" - "os" command "github.com/esonhugh/k8spider/cmd" "github.com/esonhugh/k8spider/define" "github.com/esonhugh/k8spider/pkg" "github.com/esonhugh/k8spider/pkg/mutli" + "github.com/esonhugh/k8spider/pkg/printer" "github.com/esonhugh/k8spider/pkg/scanner" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" @@ -19,7 +19,7 @@ func init() { var SubNetCmd = &cobra.Command{ Use: "subnet", - Short: "subnet is a tool to discover k8s available service ip in subnet", + Short: "subnet is a tool to discover k8s available ip in subnet", Run: func(cmd *cobra.Command, args []string) { if command.Opts.Cidr == "" { log.Warn("cidr is required") @@ -30,26 +30,27 @@ var SubNetCmd = &cobra.Command{ log.Warnf("ParseStringToIPNet failed: %v", err) return } - if command.Opts.BatchMode { - BatchRun(ipNets) + var finalRecord define.Records + if command.Opts.MultiThreadingMode { + finalRecord = RunMultiThread(ipNets, command.Opts.ThreadingNum) } else { - Run(ipNets) + finalRecord = Run(ipNets) } + printer.PrintResult(finalRecord, command.Opts.OutputFile) }, } -func Run(net *net.IPNet) { - var records define.Records = scanner.ScanSubnet(net) +func Run(net *net.IPNet) (records define.Records) { + records = scanner.ScanSubnet(net) if records == nil || len(records) == 0 { log.Warnf("ScanSubnet Found Nothing") return } - printResult(records) + return } -func BatchRun(net *net.IPNet) { - scan := mutli.NewSubnetScanner() - var finalRecord []define.Record +func RunMultiThread(net *net.IPNet, num int) (finalRecord define.Records) { + scan := mutli.NewSubnetScanner(num) for r := range scan.ScanSubnet(net) { finalRecord = append(finalRecord, r...) } @@ -57,18 +58,5 @@ func BatchRun(net *net.IPNet) { log.Warn("ScanSubnet Found Nothing") return } - printResult(finalRecord) -} - -func printResult(records define.Records) { - if command.Opts.OutputFile != "" { - f, err := os.OpenFile(command.Opts.OutputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - log.Warnf("OpenFile failed: %v", err) - } - defer f.Close() - records.Print(log.StandardLogger().Writer(), f) - } else { - records.Print(log.StandardLogger().Writer()) - } + return } diff --git a/cmd/wildcard/wildcard.go b/cmd/wildcard/wildcard.go index 8ba4884..94d0e73 100644 --- a/cmd/wildcard/wildcard.go +++ b/cmd/wildcard/wildcard.go @@ -1,10 +1,8 @@ package wildcard import ( - "os" - command "github.com/esonhugh/k8spider/cmd" - "github.com/esonhugh/k8spider/define" + "github.com/esonhugh/k8spider/pkg/printer" "github.com/esonhugh/k8spider/pkg/scanner" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" @@ -22,19 +20,6 @@ var WildCardCmd = &cobra.Command{ log.Warn("zone can't empty") return } - printResult(scanner.DumpWildCard(command.Opts.Zone)) + printer.PrintResult(scanner.DumpWildCard(command.Opts.Zone), command.Opts.OutputFile) }, } - -func printResult(records define.Records) { - if command.Opts.OutputFile != "" { - f, err := os.OpenFile(command.Opts.OutputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - log.Warnf("OpenFile failed: %v", err) - } - defer f.Close() - records.Print(log.StandardLogger().Writer(), f) - } else { - records.Print(log.StandardLogger().Writer()) - } -} diff --git a/define/record.go b/define/record.go index 560ef81..57d6033 100644 --- a/define/record.go +++ b/define/record.go @@ -34,7 +34,8 @@ func (r *Record) Print(writer ...io.Writer) { if len(writer) == 0 { W = os.Stdout } else { - W = io.MultiWriter(writer...) + w := io.MultiWriter(writer...) + W = io.MultiWriter(os.Stdout, w) } data, err := json.Marshal(r) if err != nil { diff --git a/pkg/mutli/executor.go b/pkg/mutli/executor.go index 082faf9..5ddb91f 100644 --- a/pkg/mutli/executor.go +++ b/pkg/mutli/executor.go @@ -6,8 +6,8 @@ import ( "github.com/esonhugh/k8spider/define" ) -func ScanAll(subnet *net.IPNet) (result <-chan []define.Record) { - subs := NewSubnetScanner() +func ScanAll(subnet *net.IPNet, num int) (result <-chan []define.Record) { + subs := NewSubnetScanner(num) result = ScanServiceWithChan(subs.ScanSubnet(subnet)) return result } diff --git a/pkg/mutli/subnet.go b/pkg/mutli/subnet.go index 3e3cc0a..247502e 100644 --- a/pkg/mutli/subnet.go +++ b/pkg/mutli/subnet.go @@ -12,12 +12,20 @@ import ( ) type SubnetScanner struct { - wg *sync.WaitGroup + wg *sync.WaitGroup + count int } -func NewSubnetScanner() *SubnetScanner { - return &SubnetScanner{ - wg: new(sync.WaitGroup), +func NewSubnetScanner(threading ...int) *SubnetScanner { + if len(threading) == 0 { + return &SubnetScanner{ + wg: new(sync.WaitGroup), + } + } else { + return &SubnetScanner{ + wg: new(sync.WaitGroup), + count: threading[0], + } } } @@ -28,10 +36,12 @@ func (s *SubnetScanner) ScanSubnet(subnet *net.IPNet) <-chan []define.Record { } out := make(chan []define.Record, 100) go func() { - log.Debugf("splitting subnet into 16 pices") - if subnets, err := pkg.SubnetShift(subnet, 4); err != nil { + // if subnets, err := pkg.SubnetShift(subnet, 4); err != nil { + if subnets, err := pkg.SubnetInto(subnet, s.count); err != nil { + log.Errorf("Subnet split into %v failed, fallback to single mode, reason: %v", s.count, err) go s.scan(subnet, out) } else { + log.Debugf("Subnet split into %v success", len(subnets)) for _, sn := range subnets { go s.scan(sn, out) } diff --git a/pkg/printer/p_utils.go b/pkg/printer/p_utils.go new file mode 100644 index 0000000..35acba4 --- /dev/null +++ b/pkg/printer/p_utils.go @@ -0,0 +1,21 @@ +package printer + +import ( + "os" + + "github.com/esonhugh/k8spider/define" + log "github.com/sirupsen/logrus" +) + +func PrintResult(records define.Records, OutputFile string) { + if OutputFile != "" { + f, err := os.OpenFile(OutputFile, os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + log.Warnf("OpenFile failed: %v", err) + } + defer f.Close() + records.Print(f) + } else { + records.Print() + } +}