(#16) Fixed bug and Improve codes

hahwul · Sep 18, 2023 · 661bea8 · 661bea8
1 parent 1361c83
commit 661bea8
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 26 deletions.
diff --git a/cmd/ajaxspider.go b/cmd/ajaxspider.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+
 	zap "github.com/hahwul/mzap/pkg/zap"
 	"github.com/spf13/cobra"
 )
@@ -12,6 +13,10 @@ var ajaxspiderCmd = &cobra.Command{
 	Short: "Add AjaxSpider ZAP",
 	Run: func(cmd *cobra.Command, args []string) {
 		if URLs != "" {
+			options := zap.OptionsZAP{
+				APIKey: APIKey,
+				URLs:   URLs,
+			}
 			zap.AjaxSpider(URLs, apiHosts, options)
 		} else {
 			fmt.Println("Please input --urls flag")

diff --git a/cmd/ascan.go b/cmd/ascan.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+
 	zap "github.com/hahwul/mzap/pkg/zap"
 	"github.com/spf13/cobra"
 )
@@ -12,6 +13,10 @@ var ascanCmd = &cobra.Command{
 	Short: "Add ActiveScan ZAP",
 	Run: func(cmd *cobra.Command, args []string) {
 		if URLs != "" {
+			options := zap.OptionsZAP{
+				APIKey: APIKey,
+				URLs:   URLs,
+			}
 			zap.ActiveScan(URLs, apiHosts, options)
 		} else {
 			fmt.Println("Please input --urls flag")

diff --git a/cmd/root.go b/cmd/root.go
@@ -2,17 +2,16 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/spf13/cobra"
 	"os"
 
+	"github.com/spf13/cobra"
+
 	version "github.com/hahwul/mzap/pkg/version"
-	zap "github.com/hahwul/mzap/pkg/zap"
 	homedir "github.com/mitchellh/go-homedir"
 	"github.com/spf13/viper"
 )
 
 var cfgFile, URLs, apiHosts, APIKey string
-var options zap.OptionsZAP
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
@@ -43,12 +42,6 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&APIKey, "apikey", "", "ZAP API Key / if you disable apikey, not use this option")
 	rootCmd.PersistentFlags().StringVar(&URLs, "urls", "", "URL list file / e.g --urls hosts.txt")
 	rootCmd.PersistentFlags().StringVar(&apiHosts, "apis", "http://localhost:8090", "ZAP API Host(s) address\ne.g --apis http://localhost:8090,http://192.168.0.4:8090")
-
-	options = zap.OptionsZAP{
-		APIKey: APIKey,
-		URLs:   URLs,
-	}
-	_ = options
 }
 
 // initConfig reads in config file and ENV variables if set.

diff --git a/cmd/spider.go b/cmd/spider.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+
 	zap "github.com/hahwul/mzap/pkg/zap"
 	"github.com/spf13/cobra"
 )
@@ -12,6 +13,10 @@ var spiderCmd = &cobra.Command{
 	Short: "Add ZAP spider",
 	Run: func(cmd *cobra.Command, args []string) {
 		if URLs != "" {
+			options := zap.OptionsZAP{
+				APIKey: APIKey,
+				URLs:   URLs,
+			}
 			zap.Spider(URLs, apiHosts, options)
 		} else {
 			fmt.Println("Please input --urls flag")

diff --git a/cmd/stop.go b/cmd/stop.go
@@ -13,6 +13,10 @@ var stopCmd = &cobra.Command{
 	Short: "Stop Scanning",
 	Run: func(cmd *cobra.Command, args []string) {
 		if len(args) >= 1 {
+			options := zap.OptionsZAP{
+				APIKey: APIKey,
+				URLs:   URLs,
+			}
 			if args[0] == "spider" {
 				zap.StopSpider(apiHosts, options)
 			} else if args[0] == "ascan" {

diff --git a/pkg/zap/zap.go b/pkg/zap/zap.go
@@ -20,6 +20,9 @@ type ZapObject struct {
 	subtreeOnly bool
 }
 
+// AccessAPI is api for access url
+const AccessAPI = "/JSON/core/action/accessUrl"
+
 // SpiderAPI is api for add scan
 const SpiderAPI = "/JSON/spider/action/scan/"
 
@@ -29,6 +32,23 @@ const AScanAPI = "/JSON/ascan/action/scan/"
 // AjaxSpiderAPI is api for add scan
 const AjaxSpiderAPI = "/JSON/ajaxSpider/action/scan/"
 
+func callAPI(target, urls, api, prefix string, options OptionsZAP) error {
+	req, err := http.NewRequest("GET", api+prefix, nil)
+	if err != nil {
+		panic(err)
+	}
+	q := req.URL.Query()
+	q.Add("url", target)
+	req.URL.RawQuery = q.Encode()
+	if options.APIKey != "" {
+		req.Header.Add("X-ZAP-API-Key", options.APIKey)
+	}
+
+	client := &http.Client{}
+	_, err = client.Do(req)
+	return err
+}
+
 // Run is running app
 func Run(urls, apis, prefix string, options OptionsZAP) {
 	var scanType string
@@ -67,32 +87,32 @@ func Run(urls, apis, prefix string, options OptionsZAP) {
 
 	for _, target := range arrayUrls {
 		var api = arrayAPIs[count]
-		req, err := http.NewRequest("GET", api+prefix, nil)
+		err = callAPI(target, urls, api, AccessAPI, options)
 		if err != nil {
-			panic(err)
-		}
-		q := req.URL.Query()
-		q.Add("url", target)
-		req.URL.RawQuery = q.Encode()
-		if options.APIKey != "" {
-			req.Header.Add("X-ZAP-API-Key", options.APIKey)
+			log.WithFields(logrus.Fields{
+				"data2": target,
+				"data1": api,
+			}).Warn("error (access)")
 		}
 
-		client := &http.Client{}
-		resp, err := client.Do(req)
-		log.WithFields(logrus.Fields{
-			"data2": target,
-			"data1": api,
-		}).Info("added")
-
+		err = callAPI(target, urls, api, prefix, options)
 		if err != nil {
-			//panic(err)
+			log.WithFields(logrus.Fields{
+				"data2": target,
+				"data1": api,
+			}).Warn("error (scan)")
+		} else {
+			log.WithFields(logrus.Fields{
+				"data2": target,
+				"data1": api,
+			}).Info("added")
 		}
+
 		if len(arrayAPIs)-1 > count {
 			count = count + 1
 		} else {
 			count = 0
 		}
-		defer resp.Body.Close()
+
 	}
 }