From 6ebc1ae46c8878ddf0416583c0503d67a1fda751 Mon Sep 17 00:00:00 2001
From: Esonhugh <esonhughoutside@gmail.com>
Date: Thu, 21 Nov 2024 22:30:52 +0800
Subject: [PATCH] update: dns utils

---
 Makefile            |  8 ++++++--
 cmd/all/all.go      |  5 +++++
 cmd/dnsutils/dns.go | 46 +++++++++++++++++++++++++++++++++++++++++++++
 main.go             |  1 +
 pkg/query_utils.go  | 30 ++++++++++++++++++++++++++++-
 5 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 cmd/dnsutils/dns.go

diff --git a/Makefile b/Makefile
index 909a418..4bc8a53 100644
--- a/Makefile
+++ b/Makefile
@@ -3,7 +3,7 @@ BUILD_DIR = bin
 
 MAIN_PROGRAM_NAME = k8spider
 
-default: build build-static
+default: build build-static check-size
 
 # build
 build:
@@ -11,6 +11,10 @@ build:
 
 build-static:
 	GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)-linux-static main.go
-	upx $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)-linux-static
+	upx -9 $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)-linux-static
+
+check-size:
+	ls -alh $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)*
+
 clean:
 	rm -rf $(BUILD_DIR)
\ No newline at end of file
diff --git a/cmd/all/all.go b/cmd/all/all.go
index 4ff6097..ab3335e 100644
--- a/cmd/all/all.go
+++ b/cmd/all/all.go
@@ -48,6 +48,7 @@ var AllCmd = &cobra.Command{
 			log.Warnf("ParseStringToIPNet failed: %v", err)
 			return
 		}
+
 		var finalRecord define.Records
 		if command.Opts.MultiThreadingMode {
 			finalRecord = RunMultiThread(ipNets, command.Opts.ThreadingNum)
@@ -55,6 +56,7 @@ var AllCmd = &cobra.Command{
 			finalRecord = Run(ipNets)
 		}
 		printer.PrintResult(finalRecord, command.Opts.OutputFile)
+
 		PostRun(finalRecord)
 	},
 }
@@ -78,6 +80,9 @@ func RunMultiThread(net *net.IPNet, count int) (finalRecord define.Records) {
 }
 
 func PostRun(finalRecord define.Records) {
+	if finalRecord == nil || len(finalRecord) == 0 {
+		return
+	}
 	log.Info("Extract Namespaces: ")
 	list := post.RecordsDumpNameSpace(finalRecord, command.Opts.Zone)
 	for _, ns := range list {
diff --git a/cmd/dnsutils/dns.go b/cmd/dnsutils/dns.go
new file mode 100644
index 0000000..c95b46a
--- /dev/null
+++ b/cmd/dnsutils/dns.go
@@ -0,0 +1,46 @@
+package dnsutils
+
+import (
+	"strings"
+
+	command "github.com/esonhugh/k8spider/cmd"
+	"github.com/esonhugh/k8spider/pkg"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+var queryType string
+
+func init() {
+	DNSCmd.PersistentFlags().StringVarP(&queryType, "type", "t", "A", "query type")
+	command.RootCmd.AddCommand(DNSCmd)
+}
+
+var DNSCmd = &cobra.Command{
+	Use:     "dns",
+	Aliases: []string{"dig"},
+	Short:   "dns is a command to query dns server",
+	Run: func(cmd *cobra.Command, args []string) {
+		var querier pkg.DnsQuery
+		switch strings.ToLower(queryType) {
+		case "a", "aaaa":
+			querier = pkg.QueryA
+		case "ptr":
+			querier = pkg.QueryPTR
+		case "srv":
+			querier = pkg.QuerySRV
+		case "txt":
+			querier = pkg.QueryTXT
+		default:
+			querier = pkg.QueryA
+		}
+		for _, query := range args {
+			res, err := querier(query)
+			if err != nil {
+				log.Warnf("Query %s failed: %v", query, err)
+				continue
+			}
+			log.Infof("Query [%d] %s: %v", queryType, query, res)
+		}
+	},
+}
diff --git a/main.go b/main.go
index 2ee8f92..ca4739d 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"github.com/esonhugh/k8spider/cmd"
 	_ "github.com/esonhugh/k8spider/cmd/all"
 	_ "github.com/esonhugh/k8spider/cmd/axfr"
+	_ "github.com/esonhugh/k8spider/cmd/dnsutils"
 	_ "github.com/esonhugh/k8spider/cmd/neighbor"
 	_ "github.com/esonhugh/k8spider/cmd/service"
 	_ "github.com/esonhugh/k8spider/cmd/subnet"
diff --git a/pkg/query_utils.go b/pkg/query_utils.go
index 2cb09ea..361080f 100644
--- a/pkg/query_utils.go
+++ b/pkg/query_utils.go
@@ -2,6 +2,7 @@ package pkg
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"regexp"
 	"strings"
@@ -89,7 +90,9 @@ func WarpDnsServer(dnsServer string) *SpiderResolver {
 				return d.DialContext(ctx, network, dnsServer)
 			},
 		},
-		ctx: ctx,
+		ctx:      ctx,
+		filter:   []*regexp.Regexp{},
+		contains: []string{},
 	}
 }
 
@@ -146,3 +149,28 @@ func (s *SpiderResolver) TXTRecord(domain string) ([]string, error) {
 func TXTRecord(domain string) (txts []string, err error) {
 	return NetResolver.TXTRecord(domain)
 }
+
+type DnsQuery func(domain string) ([]string, error)
+
+var (
+	QueryPTR DnsQuery = func(domain string) ([]string, error) {
+		return PTRRecord(net.ParseIP(domain)), nil
+	}
+	QueryA DnsQuery = func(domain string) ([]string, error) {
+		res, err := ARecord(domain)
+		var ret []string
+		for _, r := range res {
+			ret = append(ret, r.String())
+		}
+		return ret, err
+	}
+	QueryTXT DnsQuery = TXTRecord
+	QuerySRV DnsQuery = func(domain string) ([]string, error) {
+		_, res, err := SRVRecord(domain)
+		var ret []string
+		for _, r := range res {
+			ret = append(ret, fmt.Sprintf("%s:%d", r.Target, r.Port))
+		}
+		return ret, err
+	}
+)