embed poc

Author: zan8in

internal/runner/runner.go

@@ -61,25 +61,25 @@ func New(options *config.Options, htemplate *html.HtmlTemplate, acb config.ApiCa
 	}
 
 	// init pocs
-	allPocsYamlSlice := []string{}
+	allPocsEmbedYamlSlice := []string{}
 	if len(options.PocsFilePath) > 0 {
 		options.PocsDirectory.Set(options.PocsFilePath)
 		// console print
 		fmt.Println("   " + options.PocsFilePath)
 	} else {
 		// init default afrog-pocs
 		if allDefaultPocsYamlSlice, err := pocs.GetPocs(); err == nil {
-			allPocsYamlSlice = append(allPocsYamlSlice, allDefaultPocsYamlSlice...)
+			allPocsEmbedYamlSlice = append(allPocsEmbedYamlSlice, allDefaultPocsYamlSlice...)
 		}
 		// init ~/afrog-pocs
 		pocsDir, _ := poc.InitPocHomeDirectory()
 		if len(pocsDir) > 0 {
 			options.PocsDirectory.Set(pocsDir)
 		}
 	}
-	allPocsYamlSlice = append(allPocsYamlSlice, runner.catalog.GetPocsPath(options.PocsDirectory)...)
+	allPocsYamlSlice := runner.catalog.GetPocsPath(options.PocsDirectory)
 
-	if len(allPocsYamlSlice) == 0 {
+	if len(allPocsYamlSlice) == 0 && len(allPocsEmbedYamlSlice) == 0 {
 		return errors.New("未找到可执行脚本(POC)，请检查`默认脚本`或指定新の脚本(POC)")
 	}
 
@@ -89,7 +89,7 @@ func New(options *config.Options, htemplate *html.HtmlTemplate, acb config.ApiCa
 	}
 
 	// init scan sum
-	options.Count = len(options.Targets) * len(allPocsYamlSlice)
+	options.Count = len(options.Targets) * (len(allPocsYamlSlice) + len(allPocsEmbedYamlSlice))
 
 	// fmt.Println(ShowUsage())
 
@@ -109,7 +109,7 @@ func New(options *config.Options, htemplate *html.HtmlTemplate, acb config.ApiCa
 
 	//
 	e := core.New(options)
-	e.Execute(allPocsYamlSlice)
+	e.Execute(allPocsYamlSlice, allPocsEmbedYamlSlice)
 
 	return nil
 }

pkg/core/excute.go

@@ -4,19 +4,18 @@ import (
 	"github.com/zan8in/afrog/pkg/log"
 	"github.com/zan8in/afrog/pkg/poc"
 	"github.com/zan8in/afrog/pkg/utils"
+	"github.com/zan8in/afrog/pocs"
 )
 
 var (
 	ReverseCeyeApiKey string
 	ReverseCeyeDomain string
 )
 
-func (e *Engine) Execute(allPocsYamlSlice utils.StringSlice) {
+func (e *Engine) Execute(allPocsYamlSlice, allPocsEmbedYamlSlice utils.StringSlice) {
 	ReverseCeyeApiKey = e.options.Config.Reverse.Ceye.ApiKey
 	ReverseCeyeDomain = e.options.Config.Reverse.Ceye.Domain
 
-	//http2.Init(e.options)
-
 	var pocSlice []poc.Poc
 
 	for _, pocYaml := range allPocsYamlSlice {
@@ -28,6 +27,15 @@ func (e *Engine) Execute(allPocsYamlSlice utils.StringSlice) {
 		pocSlice = append(pocSlice, p)
 	}
 
+	for _, pocEmbedYaml := range allPocsEmbedYamlSlice {
+		p, err := pocs.ReadPocs(pocEmbedYaml)
+		if err != nil {
+			log.Log().Error(err.Error())
+			continue
+		}
+		pocSlice = append(pocSlice, p)
+	}
+
 	swg := e.workPool.PocSwg
 	for _, p := range pocSlice {
 		swg.Add()

pocs/poc.go

@@ -2,50 +2,44 @@ package pocs
 
 import (
 	"embed"
-	"fmt"
 	"io/fs"
-	"os"
-	"path/filepath"
 	"strings"
+
+	"github.com/zan8in/afrog/pkg/poc"
+	"gopkg.in/yaml.v2"
 )
 
 //go:embed afrog-pocs/*
 var f embed.FS
 
 func GetPocs() ([]string, error) {
 	allPocs := []string{}
+
 	err := fs.WalkDir(f, ".", func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
 		// fmt.Printf("path=%q, isDir=%v\n", path, d.IsDir())
 		if !d.IsDir() && strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") {
-			file := filepath.Base(path)
-			absPath, err := resolvePath(filepath.Dir(path))
-			if err != nil {
-				return err
-			}
-			allPocs = append(allPocs, filepath.Join(absPath, file))
+			allPocs = append(allPocs, path)
 		}
 		return nil
 	})
+
 	return allPocs, err
 }
 
-func resolvePath(pocName string) (string, error) {
-	if filepath.IsAbs(pocName) {
-		return pocName, nil
-	}
+func ReadPocs(path string) (poc.Poc, error) {
+	var poc = poc.Poc{}
 
-	curDirectory, err := os.Getwd()
+	file, err := f.Open(path)
 	if err != nil {
-		return "", err
+		return poc, err
 	}
+	defer file.Close()
 
-	pocPath := filepath.Join(curDirectory, "pocs", pocName)
-	if len(pocPath) > 0 {
-		return pocPath, nil
+	if err := yaml.NewDecoder(file).Decode(&poc); err != nil {
+		return poc, err
 	}
-
-	return "", fmt.Errorf("no such path found: %s", pocName)
+	return poc, nil
 }