From 5097260ec049fe40b4a4f2572fede900527b1472 Mon Sep 17 00:00:00 2001 From: Esonhugh Date: Mon, 15 Apr 2024 17:21:48 +0800 Subject: [PATCH] update: scan record print --- cmd/all/all.go | 37 ++++++++++++++----------------------- cmd/axfr/axfr.go | 14 ++------------ cmd/service/service.go | 13 ++----------- cmd/subnet/subnet.go | 32 ++++++++++---------------------- define/record.go | 3 ++- pkg/printer/p_utils.go | 21 +++++++++++++++++++++ 6 files changed, 51 insertions(+), 69 deletions(-) create mode 100644 pkg/printer/p_utils.go diff --git a/cmd/all/all.go b/cmd/all/all.go index b166ca2..1b2c045 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,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 index 7539a36..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_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 index b5b6196..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" @@ -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 index 07be242..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" @@ -30,26 +30,27 @@ 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...) } @@ -57,18 +58,5 @@ func RunMultiThread(net *net.IPNet, num int) { 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 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/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() + } +}