Refactoring, using structure, passing pointer objects, optimizing jud…

cmd/dismap/dismap.go

@@ -16,5 +16,5 @@ import (
 )
 
 func main() {
-	internal.DisMap()
+	internal.Execute()
 }

configs/rule.go

@@ -22,7 +22,7 @@ type RuleLab struct {
 	Http ReqHttp
 }
 
-var RuleData = []RuleLab{
+var RuleData = []*RuleLab{
 	{1, "EnterCRM", "body", "", InStr{"(Ent.base.js)", "", ""}, ReqHttp{"", "", nil, ""}},
 	{1, "MeterSphere", "body", "", InStr{"(<title>MeterSphere</title>)", "", ""}, ReqHttp{"", "", nil, ""}},
 	{3, "Apache Druid", "body", "", InStr{"(<title>Apache Druid</title>|content=\"Apache Druid console\")", "", ""}, ReqHttp{"", "", nil, ""}},

go.mod

@@ -4,6 +4,8 @@ go 1.16
 
 require (
 	github.com/gookit/color v1.4.2
+	github.com/ivanpirog/coloredcobra v1.0.1
+	github.com/spf13/cobra v1.7.0
 	github.com/spf13/pflag v1.0.5
 	github.com/txthinking/socks5 v0.0.0-20220212043548-414499347d4a
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd

go.sum

@@ -1,11 +1,29 @@
+github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
+github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/gookit/color v1.4.2 h1:tXy44JFSFkKnELV6WaMo/lLfu/meqITX3iAV52do7lk=
 github.com/gookit/color v1.4.2/go.mod h1:fqRyamkC1W8uxl+lxCQxOT09l/vYfZ+QeiX3rKQHCoQ=
+github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/ivanpirog/coloredcobra v1.0.1 h1:aURSdEmlR90/tSiWS0dMjdwOvCVUeYLfltLfbgNxrN4=
+github.com/ivanpirog/coloredcobra v1.0.1/go.mod h1:iho4nEKcnwZFiniGSdcgdvRgZNjxm+h20acv8vqmN6Q=
+github.com/mattn/go-colorable v0.1.9 h1:sqDoxXbdeALODt0DAeJCVp38ps9ZogZEAXjus69YV3U=
+github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
+github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
+github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
+github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/spf13/cobra v1.4.0/go.mod h1:Wo4iy3BUC+X2Fybo0PDqwJIv3dNRiZLHQymsfxlB84g=
+github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
+github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -21,14 +39,19 @@ github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 h1:QldyIu/L63oPpyvQmHg
 github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778/go.mod h1:2MuV+tbUrU1zIOPMxZ5EncGwgmMJsa+9ucAQZXxsObs=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd h1:O7DYs+zxREGLKzKoMQrtrEacpb0ZVXA5rIwylE2Xchk=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
+golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e h1:fLOSk5Q00efkSvAm+4xcoXD+RRmLmmulPn5I3Y9F2EM=
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/checkRule.go

@@ -0,0 +1,49 @@
+package internal
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/zhzyker/dismap/configs"
+	"github.com/zhzyker/dismap/pkg/logger"
+)
+
+var checkCmd = &cobra.Command{
+	Use:   "check",
+	Short: "check rules",
+	Run: func(cmd *cobra.Command, args []string) {
+		for _, rule := range configs.RuleData {
+			modes := strings.Split(rule.Mode, "|")
+			types := strings.Split(rule.Type, "|")
+			// Check the number of matches
+			if rule.Mode == "" {
+				if len(strings.Split(rule.Type, "|")) != 1 {
+					logger.Error(fmt.Sprintf("Abnormal match pattern and quantity name: %-30v type: %-20v mode: %v", rule.Name, rule.Type, rule.Mode))
+				}
+			} else {
+				if len(modes)+1 != len(types) {
+					logger.Error(fmt.Sprintf("Abnormal match pattern and quantity name: %-30v type: %-20v mode: %v", rule.Name, rule.Type, rule.Mode))
+				}
+
+			}
+			// check keyword
+			for _, item := range types {
+				if !(item == "body" || item == "header" || item == "ico") {
+					logger.Error(fmt.Sprintf("Abnormal keyword, name: %-30v type: %-20v mode: %v", rule.Name, rule.Type, rule.Mode))
+					break
+				}
+			}
+			for _, item2 := range modes {
+				if !(item2 == "" || item2 == "and" || item2 == "or") {
+					logger.Error(fmt.Sprintf("Abnormal mode, name: %-30v type: %-20v mode: %v", rule.Name, rule.Type, rule.Mode))
+					break
+				}
+			}
+		}
+	},
+}
+
+func init() {
+	RootCmd.AddCommand(checkCmd)
+}

internal/dismap.go

@@ -1,49 +1,102 @@
 package internal
 
 import (
+	"net/http"
+	"os"
+	"sync"
+
+	"github.com/ivanpirog/coloredcobra"
+	"github.com/spf13/cobra"
 	"github.com/zhzyker/dismap/configs"
 	"github.com/zhzyker/dismap/internal/flag"
 	"github.com/zhzyker/dismap/internal/operate"
 	"github.com/zhzyker/dismap/internal/output"
-	"sync"
 
 	"github.com/zhzyker/dismap/pkg/logger"
 )
 
+func which(wg *sync.WaitGroup, lock *sync.Mutex) {
+	op := output.Open()
 
-func which(Args map[string]interface{}, wg *sync.WaitGroup, lock *sync.Mutex) {
-	op := output.Open(Args)
-
-	address := Args["FlagNetwork"].(string)
+	address := flag.NetWork
 	if address != "" {
-		operate.FlagNetwork(op, wg, lock, address, Args)
+		operate.FlagNetwork(op, wg, lock, address)
 		output.Close(op)
 		return
 	}
 
-	uri := Args["FlagUrl"].(string)
+	uri := flag.InUrl
 	if uri != "" {
-		operate.FlagUrl(op, uri, Args)
+		operate.FlagUrl(op, uri)
 		output.Close(op)
 		return
 	}
 
-	file := Args["FlagFile"].(string)
+	file := flag.File
 	if file != "" {
-		operate.FlagFile(op, wg, lock, file, Args)
+		operate.FlagFile(op, wg, lock, file)
 		output.Close(op)
 		return
 	}
+}
 
+func init() {
+	RootCmd.Flags().StringVarP(&flag.InUrl, "uri", "u", "", "Specify a target URI [e.g. -u https://example.com]")
+	RootCmd.Flags().StringVarP(&flag.NetWork, "ip", "i", "", "Network segment [e.g. -i 192.168.1.0/24 or -i 192.168.1.1-10]")
+	RootCmd.Flags().StringVarP(&flag.Mode, "mode", "m", "", "Specify the protocol [e.g. -m mysql/-m http]")
+	RootCmd.Flags().StringVar(&flag.Type, "type", "", "Specify the type [e.g. --type tcp/--type udp]")
+	RootCmd.Flags().IntVar(&flag.Timeout, "timeout", 5, "Response timeout time, the default is 5 seconds")
+	RootCmd.Flags().IntVarP(&flag.Thread, "thread", "t", 500, "Number of concurrent threads")
+	RootCmd.Flags().StringVarP(&flag.Port, "port", "p", "", "Custom scan ports [e.g. -p 80,443 or -p 1-65535]")
+	RootCmd.Flags().StringVarP(&flag.Output, "output", "o", "output.txt", "Save the scan results to the specified file")
+	RootCmd.Flags().StringVarP(&flag.File, "file", "f", "", "Parse the target from the specified file for batch recognition")
+	RootCmd.Flags().BoolVar(&flag.NoIcmp, "np", false, "Not use ICMP/PING to detect surviving hosts")
+	RootCmd.Flags().StringVarP(&flag.OutJson, "json", "j", "", "Scan result in json format [e.g. -j r.json]")
+	RootCmd.Flags().BoolVar(&flag.NoColor, "nc", false, "Do not print character colors")
+	RootCmd.Flags().IntVarP(&flag.Level, "level", "l", 3, "Specify log level (0:Fatal 1:Error 2:Info 3:Warning 4:Debug 5:Verbose)")
+	RootCmd.Flags().StringVarP(&flag.Proxy, "proxy", "", "", "Use proxy scan, support http/socks5 protocol [e.g. --proxy socks5://127.0.0.1:1080]")
+	RootCmd.Flags().BoolVarP(&flag.Pprof, "pprof", "d", false, "use pprof debug, on http://localhost:6060/debug/pprof/")
 }
 
-func DisMap() {
-	configs.Banner()
-	Args := flag.Flags()
-	wg := &sync.WaitGroup{}
-	lock := &sync.Mutex{}
+func Execute() {
+	coloredcobra.Init(&coloredcobra.Config{
+		RootCmd:         RootCmd,
+		Headings:        coloredcobra.HiGreen + coloredcobra.Underline,
+		Commands:        coloredcobra.Cyan + coloredcobra.Bold,
+		Example:         coloredcobra.Italic,
+		ExecName:        coloredcobra.Bold,
+		Flags:           coloredcobra.Cyan + coloredcobra.Bold,
+		NoExtraNewlines: true,
+	})
+	err := RootCmd.Execute()
+	if err != nil {
+		os.Exit(1)
+	}
+}
+
+var RootCmd = &cobra.Command{
+	Use: "dismap",
+	Run: func(cmd *cobra.Command, args []string) {
+		if flag.NetWork == "" && flag.File == "" && flag.InUrl == "" {
+			configs.Banner()
+			cmd.Help()
+			return
+		}
+		_wg := &sync.WaitGroup{}
+		if flag.Pprof {
+			_wg.Add(1)
+			go func() {
+				logger.Info(http.ListenAndServe("localhost:6060", nil).Error())
+			}()
+		}
+
+		configs.Banner()
+		wg := &sync.WaitGroup{}
+		lock := &sync.Mutex{}
+		which(wg, lock)
+
+		logger.Info("Identification completed and ended")
 
-	information()
-	which(Args, wg, lock)
-	logger.Info("Identification completed and ended")
+		_wg.Wait()
+	},
 }

internal/flag/flag.go

@@ -1,66 +1,61 @@
 package flag
 
-import (
-	flag "github.com/spf13/pflag"
-	"os"
-)
-
 var (
-    NetWork	string
-    InUrl	string
-    Timeout	int
-    Thread	int
-    Port	string
-    Output	string
-    File	string
-    NoIcmp	bool
-    NoColor	bool
-    Mode	string
-    Type	string
-    Help	bool
-    Level	int
-    Proxy	string
-	OutJson	string
+	NetWork string
+	InUrl   string
+	Timeout int
+	Thread  int
+	Port    string
+	Output  string
+	File    string
+	NoIcmp  bool
+	NoColor bool
+	Mode    string
+	Type    string
+	Help    bool
+	Level   int
+	Proxy   string
+	OutJson string
+	Pprof   bool
 )
 
+// func init() {
+// 	flag.StringVarP(&InUrl, "uri", "u", "", "Specify a target URI [e.g. -u https://example.com]")
+// 	flag.StringVarP(&NetWork,"ip", "i", "", "Network segment [e.g. -i 192.168.1.0/24 or -i 192.168.1.1-10]")
+// 	flag.StringVarP(&Mode,"mode", "m", "", "Specify the protocol [e.g. -m mysql/-m http]")
+// 	flag.StringVar(&Type,"type", "", "Specify the type [e.g. --type tcp/--type udp]")
+// 	flag.IntVar(&Timeout, "timeout", 5, "Response timeout time, the default is 5 seconds")
+// 	flag.IntVarP(&Thread, "thread", "t", 500, "Number of concurrent threads")
+// 	flag.StringVarP(&Port, "port", "p", "", "Custom scan ports [e.g. -p 80,443 or -p 1-65535]")
+// 	flag.StringVarP(&Output, "output", "o", "output.txt", "Save the scan results to the specified file")
+// 	flag.StringVarP(&File, "file", "f", "", "Parse the target from the specified file for batch recognition")
+// 	flag.BoolVar(&NoIcmp, "np",false, "Not use ICMP/PING to detect surviving hosts")
+// 	flag.StringVarP(&OutJson, "json", "j", "", "Scan result in json format [e.g. -j r.json]")
+// 	flag.BoolVar(&NoColor, "nc", false, "Do not print character colors")
+// 	flag.IntVarP(&Level, "level", "l", 3, "Specify log level (0:Fatal 1:Error 2:Info 3:Warning 4:Debug 5:Verbose)")
+// 	flag.StringVarP(&Proxy, "proxy", "", "", "Use proxy scan, support http/socks5 protocol [e.g. --proxy socks5://127.0.0.1:1080]")
+// 	flag.BoolVarP(&Help, "help", "h",false, "Show help")
+// }
 
-func init() {
-	flag.StringVarP(&InUrl, "uri", "u", "", "Specify a target URI [e.g. -u https://example.com]")
-	flag.StringVarP(&NetWork,"ip", "i", "", "Network segment [e.g. -i 192.168.1.0/24 or -i 192.168.1.1-10]")
-	flag.StringVarP(&Mode,"mode", "m", "", "Specify the protocol [e.g. -m mysql/-m http]")
-	flag.StringVar(&Type,"type", "", "Specify the type [e.g. --type tcp/--type udp]")
-	flag.IntVar(&Timeout, "timeout", 5, "Response timeout time, the default is 5 seconds")
-	flag.IntVarP(&Thread, "thread", "t", 500, "Number of concurrent threads")
-	flag.StringVarP(&Port, "port", "p", "", "Custom scan ports [e.g. -p 80,443 or -p 1-65535]")
-	flag.StringVarP(&Output, "output", "o", "output.txt", "Save the scan results to the specified file")
-	flag.StringVarP(&File, "file", "f", "", "Parse the target from the specified file for batch recognition")
-	flag.BoolVar(&NoIcmp, "np",false, "Not use ICMP/PING to detect surviving hosts")
-	flag.StringVarP(&OutJson, "json", "j", "", "Scan result in json format [e.g. -j r.json]")
-	flag.BoolVar(&NoColor, "nc", false, "Do not print character colors")
-	flag.IntVarP(&Level, "level", "l", 3, "Specify log level (0:Fatal 1:Error 2:Info 3:Warning 4:Debug 5:Verbose)")
-	flag.StringVarP(&Proxy, "proxy", "", "", "Use proxy scan, support http/socks5 protocol [e.g. --proxy socks5://127.0.0.1:1080]")
-	flag.BoolVarP(&Help, "help", "h",false, "Show help")
-}
-
-func Flags() map[string]interface{} {
-	flag.Parse()
-	if Help {
-		flag.PrintDefaults()
-		os.Exit(0)
-	}
-	flags := map[string]interface{}{
-		"FlagUrl":       InUrl,
-		"FlagNetwork":   NetWork,
-		"FlagMode":      Mode,
-		"FlagType":      Type,
-		"FlagTimeout":   Timeout,
-		"FlagThread":    Thread,
-		"FlagPort":      Port,
-		"FlagOutput":    Output,
-		"FlagFile":      File,
-		"FlagNoIcmp":    NoIcmp,
-		"FlagProxy":	 Proxy,
-		"FlagOutJson":	 OutJson,
-	}
-	return flags
-}
+// func Flags() map[string]interface{} {
+// 	flag.Parse()
+// 	if Help {
+// 		flag.PrintDefaults()
+// 		os.Exit(0)
+// 	}
+// 	flags := map[string]interface{}{
+// 		"FlagUrl":       InUrl,
+// 		"FlagNetwork":   NetWork,
+// 		"FlagMode":      Mode,
+// 		"FlagType":      Type,
+// 		"FlagTimeout":   Timeout,
+// 		"FlagThread":    Thread,
+// 		"FlagPort":      Port,
+// 		"FlagOutput":    Output,
+// 		"FlagFile":      File,
+// 		"FlagNoIcmp":    NoIcmp,
+// 		"FlagProxy":	 Proxy,
+// 		"FlagOutJson":	 OutJson,
+// 	}
+// 	return flags
+// }