From 38965d4836a748733869ffe9f1af7ef81bb9e3c2 Mon Sep 17 00:00:00 2001
From: esonhugh <esonhughoutside@gmail.com>
Date: Fri, 18 Oct 2024 17:29:05 +0800
Subject: [PATCH] update: filter results in k8spider

---
 cmd/root.go        | 13 ++++++++-
 pkg/query_utils.go | 70 +++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index e57cbbd..c325bb4 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -22,6 +22,9 @@ var Opts = struct {
 	ThreadingNum       int
 
 	SkipKubeDNSCheck bool
+
+	FilterRules   []string
+	FilterStrings []string
 }{}
 
 func init() {
@@ -43,6 +46,9 @@ func init() {
 	RootCmd.PersistentFlags().IntVarP(&Opts.ThreadingNum, "thread-num", "n", 16, "threading num, default 16")
 
 	RootCmd.PersistentFlags().BoolVarP(&Opts.SkipKubeDNSCheck, "skip-kube-dns-check", "k", false, "skip kube-dns check, force check if current environment is matched kube-dns schema")
+
+	RootCmd.PersistentFlags().StringSliceVarP(&Opts.FilterRules, "filter-rules", "f", []string{}, "filter regexp rules")
+	RootCmd.PersistentFlags().StringSliceVarP(&Opts.FilterStrings, "filter-strings", "s", []string{}, "filter contained strings")
 }
 
 var RootCmd = &cobra.Command{
@@ -62,7 +68,12 @@ var RootCmd = &cobra.Command{
 		if Opts.DnsServer != "" {
 			pkg.NetResolver = pkg.WarpDnsServer(Opts.DnsServer)
 		}
-
+		for _, rules := range Opts.FilterRules {
+			pkg.NetResolver.SetFilter(rules)
+		}
+		for _, rules := range Opts.FilterStrings {
+			pkg.NetResolver.SetFilter(rules)
+		}
 		// Check if current environment is a kubernetes cluster
 		// If the command is whereisdns, which means DNS is not sure , so skip this check!
 		// If SkipKubeDNSCheck is true, skip this check!
diff --git a/pkg/query_utils.go b/pkg/query_utils.go
index 14a9e60..2cb09ea 100644
--- a/pkg/query_utils.go
+++ b/pkg/query_utils.go
@@ -3,6 +3,8 @@ package pkg
 import (
 	"context"
 	"net"
+	"regexp"
+	"strings"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -15,18 +17,65 @@ var (
 )
 
 type SpiderResolver struct {
-	dns string
-	ctx context.Context
-	r   *net.Resolver
+	dns      string
+	ctx      context.Context
+	r        *net.Resolver
+	filter   []*regexp.Regexp
+	contains []string
 }
 
 func DefaultResolver() *SpiderResolver {
 	ctx, _ := context.WithTimeout(context.Background(), time.Duration(DnsTimeout)*time.Second) // I don't think if a inside cluster dns query has more than 2s latency.
 	return &SpiderResolver{
-		dns: "<default-dns>",
-		r:   net.DefaultResolver,
-		ctx: ctx,
+		dns:      "<default-dns>",
+		r:        net.DefaultResolver,
+		ctx:      ctx,
+		filter:   []*regexp.Regexp{},
+		contains: []string{},
+	}
+}
+
+func (r *SpiderResolver) SetFilter(filters ...string) {
+	for _, filter := range filters {
+		r.filter = append(r.filter, regexp.MustCompile(filter))
+	}
+}
+
+func (r *SpiderResolver) SetContainsFilter(name ...string) {
+	r.contains = append(r.contains, name...)
+}
+
+func (r *SpiderResolver) SetSuffixFilter(filter string) {
+	r.SetFilter(filter + "$")
+}
+
+func (r *SpiderResolver) filterString(target string) bool {
+	log.Tracef("filtering %s", target)
+	for _, re := range r.filter {
+		if re.MatchString(target) {
+			log.Tracef("target %s matched regexp rule %s", target, re.String())
+			return true
+		}
+	}
+	for _, re := range r.contains {
+		if strings.Contains(target, re) {
+			log.Tracef("target %s matched contains rule %s", target, re)
+			return true
+		}
 	}
+	return false
+}
+
+func (r *SpiderResolver) filterStringArray(target []string) []string {
+	var filtered []string
+	for _, re := range target {
+		if r.filterString(re) {
+			continue
+		}
+		filtered = append(filtered, re)
+	}
+	log.Tracef("filtering %s \nresult: %s", strings.Join(target, " "), strings.Join(filtered, " "))
+	return filtered
 }
 
 func WarpDnsServer(dnsServer string) *SpiderResolver {
@@ -54,7 +103,7 @@ func (s *SpiderResolver) PTRRecord(ip net.IP) []string {
 		log.Debugf("LookupAddr failed: %v", err)
 		return nil
 	}
-	return names
+	return s.filterStringArray(names)
 }
 
 func PTRRecord(ip net.IP) []string {
@@ -63,6 +112,13 @@ func PTRRecord(ip net.IP) []string {
 
 func (s *SpiderResolver) SRVRecord(svcDomain string) (string, []*net.SRV, error) {
 	cname, srvs, err := s.r.LookupSRV(s.ctx, "", "", svcDomain)
+	var finalsrv []*net.SRV
+	for _, srv := range srvs {
+		if s.filterString(srv.Target) {
+			continue
+		}
+		finalsrv = append(finalsrv, srv)
+	}
 	return cname, srvs, err
 }