STRLCPY/afrog

■ ■ ■ ■ ■ ■

internal/runner/runner.go

		skipped 60 lines
61	61		}
62	62
63	63		// init pocs
64		-	allPocsYamlSlice := []string{}
	64	+	allPocsEmbedYamlSlice := []string{}
65	65		if len(options.PocsFilePath) > 0 {
66	66		options.PocsDirectory.Set(options.PocsFilePath)
67	67		// console print
		skipped 1 lines
69	69		} else {
70	70		// init default afrog-pocs
71	71		if allDefaultPocsYamlSlice, err := pocs.GetPocs(); err == nil {
72		-	allPocsYamlSlice = append(allPocsYamlSlice, allDefaultPocsYamlSlice...)
	72	+	allPocsEmbedYamlSlice = append(allPocsEmbedYamlSlice, allDefaultPocsYamlSlice...)
73	73		}
74	74		// init ~/afrog-pocs
75	75		pocsDir, _ := poc.InitPocHomeDirectory()
		skipped 1 lines
77	77		options.PocsDirectory.Set(pocsDir)
78	78		}
79	79		}
80		-	allPocsYamlSlice = append(allPocsYamlSlice, runner.catalog.GetPocsPath(options.PocsDirectory)...)
	80	+	allPocsYamlSlice := runner.catalog.GetPocsPath(options.PocsDirectory)
81	81
82		-	if len(allPocsYamlSlice) == 0 {
	82	+	if len(allPocsYamlSlice) == 0 && len(allPocsEmbedYamlSlice) == 0 {
83	83		return errors.New("未找到可执行脚本(POC)，请检查`默认脚本`或指定新の脚本(POC)")
84	84		}
85	85
		skipped 3 lines
89	89		}
90	90
91	91		// init scan sum
92		-	options.Count = len(options.Targets) * len(allPocsYamlSlice)
	92	+	options.Count = len(options.Targets) * (len(allPocsYamlSlice) + len(allPocsEmbedYamlSlice))
93	93
94	94		// fmt.Println(ShowUsage())
95	95
		skipped 13 lines
109	109
110	110		//
111	111		e := core.New(options)
112		-	e.Execute(allPocsYamlSlice)
	112	+	e.Execute(allPocsYamlSlice, allPocsEmbedYamlSlice)
113	113
114	114		return nil
115	115		}
		skipped 7 lines

■ ■ ■ ■ ■ ■

pkg/core/excute.go

		skipped 3 lines
4	4		"github.com/zan8in/afrog/pkg/log"
5	5		"github.com/zan8in/afrog/pkg/poc"
6	6		"github.com/zan8in/afrog/pkg/utils"
	7	+	"github.com/zan8in/afrog/pocs"
7	8		)
8	9
9	10		var (
		skipped 1 lines
11	12		ReverseCeyeDomain string
12	13		)
13	14
14		-	func (e *Engine) Execute(allPocsYamlSlice utils.StringSlice) {
	15	+	func (e *Engine) Execute(allPocsYamlSlice, allPocsEmbedYamlSlice utils.StringSlice) {
15	16		ReverseCeyeApiKey = e.options.Config.Reverse.Ceye.ApiKey
16	17		ReverseCeyeDomain = e.options.Config.Reverse.Ceye.Domain
17	18
18		-	//http2.Init(e.options)
19		-
20	19		var pocSlice []poc.Poc
21	20
22	21		for _, pocYaml := range allPocsYamlSlice {
23	22		p, err := poc.ReadPocs(pocYaml)
	23	+	if err != nil {
	24	+	log.Log().Error(err.Error())
	25	+	continue
	26	+	}
	27	+	pocSlice = append(pocSlice, p)
	28	+	}
	29	+
	30	+	for _, pocEmbedYaml := range allPocsEmbedYamlSlice {
	31	+	p, err := pocs.ReadPocs(pocEmbedYaml)
24	32		if err != nil {
25	33		log.Log().Error(err.Error())
26	34		continue
		skipped 54 lines

■ ■ ■ ■ ■ ■

pocs/poc.go

		skipped 1 lines
2	2
3	3		import (
4	4		"embed"
5		-	"fmt"
6	5		"io/fs"
7		-	"os"
8		-	"path/filepath"
9	6		"strings"
	7	+
	8	+	"github.com/zan8in/afrog/pkg/poc"
	9	+	"gopkg.in/yaml.v2"
10	10		)
11	11
12	12		//go:embed afrog-pocs/*
		skipped 1 lines
14	14
15	15		func GetPocs() ([]string, error) {
16	16		allPocs := []string{}
	17	+
17	18		err := fs.WalkDir(f, ".", func(path string, d fs.DirEntry, err error) error {
18	19		if err != nil {
19	20		return err
20	21		}
21	22		// fmt.Printf("path=%q, isDir=%v\n", path, d.IsDir())
22	23		if !d.IsDir() && strings.HasSuffix(path, ".yaml") \|\| strings.HasSuffix(path, ".yml") {
23		-	file := filepath.Base(path)
24		-	absPath, err := resolvePath(filepath.Dir(path))
25		-	if err != nil {
26		-	return err
27		-	}
28		-	allPocs = append(allPocs, filepath.Join(absPath, file))
	24	+	allPocs = append(allPocs, path)
29	25		}
30	26		return nil
31	27		})
	28	+
32	29		return allPocs, err
33	30		}
34	31
35		-	func resolvePath(pocName string) (string, error) {
36		-	if filepath.IsAbs(pocName) {
37		-	return pocName, nil
38		-	}
	32	+	func ReadPocs(path string) (poc.Poc, error) {
	33	+	var poc = poc.Poc{}
39	34
40		-	curDirectory, err := os.Getwd()
	35	+	file, err := f.Open(path)
41	36		if err != nil {
42		-	return "", err
	37	+	return poc, err
43	38		}
	39	+	defer file.Close()
44	40
45		-	pocPath := filepath.Join(curDirectory, "pocs", pocName)
46		-	if len(pocPath) > 0 {
47		-	return pocPath, nil
	41	+	if err := yaml.NewDecoder(file).Decode(&poc); err != nil {
	42	+	return poc, err
48	43		}
49		-
50		-	return "", fmt.Errorf("no such path found: %s", pocName)
	44	+	return poc, nil
51	45		}
52	46