update: wildcard logic with multi-threading

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
 }

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)
 	},
 }

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{

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)
 	},
 }

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,45 +30,33 @@ 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...)
 	}
 	if len(finalRecord) == 0 {
 		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
 }

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())
-	}
-}

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 {

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
 }

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)
 			}

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()
+	}
+}