update: scan record print

Esonhugh · Apr 15, 2024 · 5097260 · 5097260
1 parent f601f2f
commit 5097260
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 69 deletions.
diff --git 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,52 +27,43 @@ var AllCmd = &cobra.Command{
 			log.Warn("cidr is required")
 			return
 		}
+		// AXFR Dumping
 		records, err := scanner.DumpAXFR(dns.Fqdn(command.Opts.Zone), "ns.dns."+command.Opts.Zone+":53")
 		if err == nil {
-			printResult(records)
+			printer.PrintResult(records, command.Opts.OutputFile)
+		} else {
+			log.Errorf("Transfer failed: %v", err)
 		}
-		log.Errorf("Transfer failed: %v", err)
+		// Service Discovery
 		ipNets, err := pkg.ParseStringToIPNet(command.Opts.Cidr)
 		if err != nil {
 			log.Warnf("ParseStringToIPNet failed: %v", err)
 			return
 		}
+		var finalRecord define.Records
 		if command.Opts.MultiThreadingMode {
-			RunMultiThread(ipNets, command.Opts.ThreadingNum)
+			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 RunMultiThread(net *net.IPNet, count int) {
+func RunMultiThread(net *net.IPNet, count int) (finalRecord define.Records) {
 	scan := mutli.ScanAll(net, count)
-	var finalRecord []define.Record
 	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_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
@@ -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_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/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"
@@ -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_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
@@ -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"
@@ -30,45 +30,33 @@ var SubNetCmd = &cobra.Command{
 			log.Warnf("ParseStringToIPNet failed: %v", err)
 			return
 		}
+		var finalRecord define.Records
 		if command.Opts.MultiThreadingMode {
-			RunMultiThread(ipNets, command.Opts.ThreadingNum)
+			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 RunMultiThread(net *net.IPNet, num int) {
+func RunMultiThread(net *net.IPNet, num int) (finalRecord define.Records) {
 	scan := mutli.NewSubnetScanner(num)
-	var finalRecord []define.Record
 	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_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/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/printer/p_utils.go 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()
+	}
+}