STRLCPY/afrog

■ ■ ■ ■ ■ ■

cmd/afrog/main.go

		skipped 10 lines
11	11		"github.com/zan8in/afrog/pkg/html"
12	12		"github.com/zan8in/afrog/pkg/log"
13	13		"github.com/zan8in/afrog/pkg/poc"
	14	+	"github.com/zan8in/afrog/pkg/upgrade"
14	15		)
15	16
16	17		var options = &config.Options{}
		skipped 15 lines
32	33
33	34		app.Action = func(c *cli.Context) error {
34	35
35		-	title := log.LogColor.Title("一个挖洞工具 - afrog V" + c.App.Version)
36		-	defconfig := log.LogColor.Info("默认配置 " + options.Config.GetConfigPath())
37		-	defpocdir := log.LogColor.Info("默认脚本 " + poc.GetPocPath())
38		-	fmt.Println(title + "\r\n" + defconfig + "\r\n" + defpocdir)
	36	+	title := log.LogColor.Vulner("A tool for finding vulnerabilities - afrog V" + config.Version)
	37	+
	38	+	upgrade := upgrade.New()
	39	+	upgrade.UpgradeAfrogPocs()
	40	+
	41	+	defconfig := log.LogColor.Low("Default Conf " + options.Config.GetConfigPath())
	42	+	defpocdir := log.LogColor.Low("Default Pocs " + poc.GetPocPath())
	43	+
	44	+	fmt.Println(title + "\r\n" + defconfig + "\r\n" + defpocdir + " v" + upgrade.LastestVersion + "")
39	45
40	46		htemplate.Filename = options.Output
41	47		if err := htemplate.New(); err != nil {
		skipped 28 lines
70	76
71	77		err := app.Run(os.Args)
72	78		if err != nil {
73		-	fmt.Println(log.LogColor.High("启动 afrog 出错，", err.Error()))
	79	+	fmt.Println(log.LogColor.High("Failed to start afrog，", err.Error()))
74	80		}
75	81		}
76	82
		skipped 6 lines

■ ■ ■ ■ ■ ■

cmd/download/main.go

		skipped 1 lines
2	2
3	3		import (
4	4		"fmt"
	5	+	"os"
5	6		"time"
	7	+
	8	+	"github.com/cavaliergopher/grab/v3"
	9	+	"github.com/zan8in/afrog/pkg/upgrade"
6	10		)
7	11
8		-	// Bar ...
9		-	type Bar struct {
10		-	percent int64 // progress percentage
11		-	cur int64 // current progress
12		-	total int64 // total value for progress
13		-	rate string // the actual progress bar to be printed
14		-	graph string // the fill value for progress bar
	12	+	func VersionOrdinal(version string) string {
	13	+	// ISO/IEC 14651:2011
	14	+	const maxByte = 1<<8 - 1
	15	+	vo := make([]byte, 0, len(version)+8)
	16	+	j := -1
	17	+	for i := 0; i < len(version); i++ {
	18	+	b := version[i]
	19	+	if '0' > b \|\| b > '9' {
	20	+	vo = append(vo, b)
	21	+	j = -1
	22	+	continue
	23	+	}
	24	+	if j == -1 {
	25	+	vo = append(vo, 0x00)
	26	+	j = len(vo) - 1
	27	+	}
	28	+	if vo[j] == 1 && vo[j+1] == '0' {
	29	+	vo[j+1] = b
	30	+	continue
	31	+	}
	32	+	if vo[j]+1 > maxByte {
	33	+	panic("VersionOrdinal: invalid version")
	34	+	}
	35	+	vo = append(vo, b)
	36	+	vo[j]++
	37	+	}
	38	+	return string(vo)
15	39		}
16	40
17		-	func (bar *Bar) NewOption(start, total int64) {
18		-	bar.cur = start
19		-	bar.total = total
20		-	if bar.graph == "" {
21		-	bar.graph = "█"
22		-	}
23		-	bar.percent = bar.getPercent()
24		-	for i := 0; i < int(bar.percent); i += 2 {
25		-	bar.rate += bar.graph // initial progress position
26		-	}
27		-	}
	41	+	func main() {
	42	+	fmt.Println("1.2" > "1.1.1")
	43	+
	44	+	return
	45	+	upgrade := upgrade.New()
	46	+	upgrade.UpgradeAfrogPocs()
	47	+	return
	48	+	// uz := utils.NewUnzip()
	49	+
	50	+	// files, err := uz.Extract("./afrog-pocs.zip", "./cmd/download/")
	51	+	// if err != nil {
	52	+	// fmt.Println(err)
	53	+	// }
	54	+
	55	+	// fmt.Printf("extracted files count: %d", len(files))
	56	+	// fmt.Printf("files list: %v", files)
	57	+	// create client
	58	+	client := grab.NewClient()
	59	+	req, _ := grab.NewRequest(".", "http://binbin.run/afrog-pocs.zip")
	60	+
	61	+	// start download
	62	+	fmt.Printf("Downloading %v...\n", req.URL())
	63	+	resp := client.Do(req)
	64	+	fmt.Printf(" %v\n", resp.HTTPResponse.Status)
	65	+
	66	+	// start UI loop
	67	+	t := time.NewTicker(500 * time.Millisecond)
	68	+	defer t.Stop()
28	69
29		-	func (bar *Bar) getPercent() int64 {
30		-	return int64((float32(bar.cur) / float32(bar.total)) * 100)
31		-	}
	70	+	Loop:
	71	+	for {
	72	+	select {
	73	+	case <-t.C:
	74	+	fmt.Printf("\rtransferred %v / %v bytes (%.2f%%)\n",
	75	+	resp.BytesComplete(),
	76	+	resp.Size,
	77	+	100*resp.Progress())
32	78
33		-	func (bar *Bar) Play(cur int64) {
34		-	bar.cur = cur
35		-	last := bar.percent
36		-	bar.percent = bar.getPercent()
37		-	if bar.percent != last && bar.percent%2 == 0 {
38		-	bar.rate += bar.graph
	79	+	case <-resp.Done:
	80	+	// download is complete
	81	+	break Loop
	82	+	}
39	83		}
40		-	fmt.Printf("\r%-50s%3d%% %8d/%d", bar.rate, bar.percent, bar.cur, bar.total)
41		-	}
42	84
43		-	func (bar *Bar) Finish() {
44		-	fmt.Println()
45		-	}
	85	+	// check for errors
	86	+	if err := resp.Err(); err != nil {
	87	+	fmt.Fprintf(os.Stderr, "Download failed: %v\n", err)
	88	+	os.Exit(1)
	89	+	}
	90	+
	91	+	fmt.Printf("Download saved to %v \n", resp.Filename)
46	92
47		-	func main() {
48		-	var bar Bar
49		-	bar.NewOption(0, 100)
50		-	for i := 0; i <= 100; i++ {
51		-	time.Sleep(100 * time.Millisecond)
52		-	bar.Play(int64(i))
53		-	}
54		-	bar.Finish()
	93	+	// Output:
	94	+	// Downloading http://www.golang-book.com/public/pdf/gobook.pdf...
	95	+	// 200 OK
	96	+	// transferred 42970 / 2893557 bytes (1.49%)
	97	+	// transferred 1207474 / 2893557 bytes (41.73%)
	98	+	// transferred 2758210 / 2893557 bytes (95.32%)
	99	+	// Download saved to ./gobook.pdf
55	100		}
56	101

■ ■ ■ ■ ■ ■

go.mod

		skipped 4 lines
5	5		require (
6	6		github.com/andybalholm/brotli v1.0.4 // indirect
7	7		github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20210826220005-b48c857c3a0e // indirect
	8	+	github.com/cavaliergopher/grab/v3 v3.0.1 // indirect
8	9		github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d // indirect
9	10		github.com/dlclark/regexp2 v1.4.0 // indirect
10	11		github.com/fatih/color v1.13.0 // indirect
		skipped 37 lines

■ ■ ■ ■ ■ ■

go.sum

		skipped 7 lines
8	8		github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20210826220005-b48c857c3a0e h1:GCzyKMDDjSGnlpl3clrdAK7I1AaVoaiKDOYkUzChZzg=
9	9		github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20210826220005-b48c857c3a0e/go.mod h1:F7bn7fEU90QkQ3tnmaTx3LTKLEDqnwWODIYppRQ5hnY=
10	10		github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
	11	+	github.com/cavaliergopher/grab/v3 v3.0.1 h1:4z7TkBfmPjmLAAmkkAZNX/6QJ1nNFdv3SdIHXju0Fr4=
	12	+	github.com/cavaliergopher/grab/v3 v3.0.1/go.mod h1:1U/KNnD+Ft6JJiYoYBAimKH2XrYptb8Kl3DFGmsjpq4=
11	13		github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
12	14		github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
13	15		github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
		skipped 238 lines

■ ■ ■ ■ ■ ■

internal/runner/runner.go

		skipped 2 lines
3	3		import (
4	4		"errors"
5	5		"fmt"
	6	+	"os"
6	7
7	8		"github.com/zan8in/afrog/pkg/catalog"
8	9		"github.com/zan8in/afrog/pkg/config"
		skipped 29 lines
38	39		options.Config = config
39	40
40	41		if len(options.Config.Reverse.Ceye.Domain) == 0 \|\| len(options.Config.Reverse.Ceye.ApiKey) == 0 {
41		-	return errors.New("rerverse CeyeApiKey or CeyeDomain is Empty in your `/home/[yourname]/.config/afrog/afrog-config.yaml`")
	42	+	homeDir, _ := os.UserHomeDir()
	43	+	return errors.New("please edit `api-key` and `domain` in `" + homeDir + "/.config/afrog/afrog-config.yaml`")
42	44		}
43	45
44	46		// init targets
		skipped 44 lines

■ ■ ■ ■ ■ ■

pkg/config/config.go

		skipped 10 lines
11	11
12	12		// Config is a afrog-config.yaml catalog helper implementation
13	13		type Config struct {
14		-	ConfigVersion string `yaml:"version"`
15	14		PocSizeWaitGroup int32 `yaml:"poc_sizewaitgroup"`
16	15		TargetSizeWaitGroup int32 `yaml:"target_sizewaitgroup"`
17	16		ConfigHttp ConfigHttp `yaml:"http"`
		skipped 28 lines
46	45		func New() (*Config, error) {
47	46		if isExistConfigFile() != nil {
48	47		c := Config{}
49		-	c.ConfigVersion = Version
50	48		c.PocSizeWaitGroup = 25
51	49		c.TargetSizeWaitGroup = 6
52	50		configHttp := c.ConfigHttp
		skipped 104 lines

■ ■ ■ ■ ■ ■

pkg/poc/poc.go

		skipped 2 lines
3	3		import (
4	4		"os"
5	5		"path/filepath"
	6	+	"strings"
6	7
7	8		"github.com/zan8in/afrog/pkg/utils"
8	9		"gopkg.in/yaml.v2"
		skipped 115 lines
124	125		return ""
125	126		}
126	127		return configFile
	128	+	}
	129	+
	130	+	func GetPocVersionNumber() (string, error) {
	131	+	version := GetPocPath() + "/version"
	132	+	v, err := utils.ReadFromFile(version)
	133	+	if err != nil {
	134	+	return "0", nil
	135	+	}
	136	+	return strings.TrimSpace(string(v)), nil
127	137		}
128	138
129	139		// Read a poc yaml file from disk.
		skipped 62 lines

■ ■ ■ ■ ■ ■

pkg/upgrade/upgrade.go

1	+	package upgrade
2	+
3	+	import (
4	+	"errors"
5	+	"fmt"
6	+	"io/ioutil"
7	+	"net/http"
8	+	"os"
9	+	"strings"
10	+
11	+	"github.com/cavaliergopher/grab/v3"
12	+	"github.com/zan8in/afrog/pkg/log"
13	+	"github.com/zan8in/afrog/pkg/poc"
14	+	"github.com/zan8in/afrog/pkg/utils"
15	+	)
16	+
17	+	type Upgrade struct {
18	+	HomeDir string
19	+	CurrVersion string
20	+	RemoteVersion string
21	+	LastestVersion string
22	+	}
23	+
24	+	const (
25	+	upHost = "http://binbin.run/afrog-release"
26	+	upPathName = "/afrog-pocs"
27	+	upPath = "/afrog-pocs.zip"
28	+	upRemoteVersion = "/version"
29	+	)
30	+
31	+	func New() *Upgrade {
32	+	homeDir, _ := os.UserHomeDir()
33	+	return &Upgrade{HomeDir: homeDir}
34	+	}
35	+
36	+	func (u *Upgrade) CheckUpgrade() (bool, error) {
37	+	curVersion, err := poc.GetPocVersionNumber()
38	+	if err != nil {
39	+	return false, errors.New("failed to get local version number")
40	+	}
41	+
42	+	resp, err := http.Get(upHost + upRemoteVersion)
43	+	if err != nil {
44	+	return false, errors.New("failed to get remote version number")
45	+	}
46	+	defer resp.Body.Close()
47	+
48	+	remoteVersion, err := ioutil.ReadAll(resp.Body)
49	+	if err != nil {
50	+	return false, errors.New("failed to get remote version number")
51	+	}
52	+
53	+	u.CurrVersion = curVersion
54	+	u.RemoteVersion = strings.TrimSpace(string(remoteVersion))
55	+
56	+	return strings.TrimSpace(string(remoteVersion)) > curVersion, nil
57	+	}
58	+
59	+	func (u *Upgrade) UpgradeAfrogPocs() {
60	+	isUp, err := u.CheckUpgrade()
61	+	if err != nil {
62	+	u.LastestVersion = u.CurrVersion
63	+	return
64	+	}
65	+	if !isUp {
66	+	u.LastestVersion = u.CurrVersion
67	+	return
68	+	}
69	+	if isUp {
70	+	fmt.Println(log.LogColor.Info("Downloading latest release.."))
71	+	u.LastestVersion = u.RemoteVersion
72	+	u.Download()
73	+	}
74	+	}
75	+
76	+	func (u *Upgrade) Download() {
77	+	resp, err := grab.Get(u.HomeDir, upHost+upPath)
78	+	if err != nil {
79	+	fmt.Println(log.LogColor.Low(err.Error()))
80	+	return
81	+	}
82	+
83	+	utils.RandSleep(1000)
84	+
85	+	u.Unzip(resp.Filename)
86	+
87	+	utils.RandSleep(1000)
88	+
89	+	os.Remove(resp.Filename)
90	+	}
91	+
92	+	func (u *Upgrade) Unzip(src string) {
93	+	uz := utils.NewUnzip()
94	+
95	+	_, err := uz.Extract(src, u.HomeDir)
96	+	if err != nil {
97	+	fmt.Println(log.LogColor.Low("Failed updated afrog-pocs ", err))
98	+	}
99	+
100	+	fmt.Println(log.LogColor.Info("Successfully updated afrog-pocs to ", u.HomeDir+upPathName))
101	+	}
102	+

■ ■ ■ ■ ■ ■

pkg/utils/unzip.go

1	+	package utils
2	+
3	+	import (
4	+	"archive/zip"
5	+	"fmt"
6	+	"io"
7	+	"os"
8	+	"path/filepath"
9	+	"strings"
10	+	)
11	+
12	+	// reference: https://raw.githubusercontent.com/artdarek/go-unzip/master/pkg/unzip/unzip.go
13	+
14	+	type Unzip struct {
15	+	}
16	+
17	+	func NewUnzip() *Unzip {
18	+	return &Unzip{}
19	+	}
20	+
21	+	func (uz Unzip) Extract(source, destination string) ([]string, error) {
22	+	r, err := zip.OpenReader(source)
23	+	if err != nil {
24	+	return nil, err
25	+	}
26	+
27	+	defer func() {
28	+	if err := r.Close(); err != nil {
29	+	panic(err)
30	+	}
31	+	}()
32	+
33	+	err = os.MkdirAll(destination, 0755)
34	+	if err != nil {
35	+	return nil, err
36	+	}
37	+
38	+	var extractedFiles []string
39	+	for _, f := range r.File {
40	+	err := uz.extractAndWriteFile(destination, f)
41	+	if err != nil {
42	+	return nil, err
43	+	}
44	+
45	+	extractedFiles = append(extractedFiles, f.Name)
46	+	}
47	+
48	+	return extractedFiles, nil
49	+	}
50	+
51	+	func (Unzip) extractAndWriteFile(destination string, f *zip.File) error {
52	+	rc, err := f.Open()
53	+	if err != nil {
54	+	return err
55	+	}
56	+	defer func() {
57	+	if err := rc.Close(); err != nil {
58	+	panic(err)
59	+	}
60	+	}()
61	+
62	+	path := filepath.Join(destination, f.Name)
63	+	if !strings.HasPrefix(path, filepath.Clean(destination)+string(os.PathSeparator)) {
64	+	return fmt.Errorf("%s: illegal file path", path)
65	+	}
66	+
67	+	if f.FileInfo().IsDir() {
68	+	err = os.MkdirAll(path, f.Mode())
69	+	if err != nil {
70	+	return err
71	+	}
72	+	} else {
73	+	err = os.MkdirAll(filepath.Dir(path), f.Mode())
74	+	if err != nil {
75	+	return err
76	+	}
77	+
78	+	f, err := os.OpenFile(path, os.O_WRONLY\|os.O_CREATE\|os.O_TRUNC, f.Mode())
79	+	if err != nil {
80	+	return err
81	+	}
82	+	defer func() {
83	+	if err := f.Close(); err != nil {
84	+	panic(err)
85	+	}
86	+	}()
87	+
88	+	_, err = io.Copy(f, rc)
89	+	if err != nil {
90	+	return err
91	+	}
92	+	}
93	+
94	+	return nil
95	+	}
96	+

add check upgrade afrog-pocs