cmd 控制台输出美化、report报告美化及体验提升

Author: zan8in

cmd/afrog/main.go

@@ -2,8 +2,8 @@ package main
 
 import (
 	"fmt"
-	"github.com/zan8in/afrog/pkg/utils"
 	"os"
+	"sync"
 
 	"github.com/urfave/cli/v2"
 	"github.com/zan8in/afrog/internal/runner"
@@ -13,10 +13,13 @@ import (
 	"github.com/zan8in/afrog/pkg/log"
 	"github.com/zan8in/afrog/pkg/poc"
 	"github.com/zan8in/afrog/pkg/upgrade"
+	"github.com/zan8in/afrog/pkg/utils"
 )
 
 var options = &config.Options{}
 var htemplate = &html.HtmlTemplate{}
+var lock sync.Mutex
+var number = 0
 
 func main() {
 	app := cli.NewApp()
@@ -37,7 +40,7 @@ func main() {
 		upgrade := upgrade.New()
 		upgrade.UpgradeAfrogPocs()
 
-		showBanner(upgrade.LastestAfrogVersion)
+		runner.ShowBanner2(upgrade.LastestAfrogVersion)
 
 		fmt.Println("PATH:")
 		fmt.Println("   " + options.Config.GetConfigPath())
@@ -51,25 +54,29 @@ func main() {
 		err := runner.New(options, func(result interface{}) {
 			r := result.(*core.Result)
 
-			options.OptLock.Lock()
-			defer options.OptLock.Unlock()
+			lock.Lock()
+			defer lock.Unlock()
 
 			if !options.Silent {
 				options.CurrentCount++
 			}
 
 			if r.IsVul {
-				r.PrintColorResultInfoConsole()
+				number++
 
 				if len(r.Output) > 0 {
 					htemplate.Result = r
+					htemplate.Number = utils.GetNumberText(number)
 					htemplate.Append()
 				}
+
+				r.PrintColorResultInfoConsole(utils.GetNumberText(number))
 			}
 
 			if !options.Silent {
 				fmt.Printf("\r%d/%d | %d%% ", options.CurrentCount, options.Count, options.CurrentCount*100/options.Count)
 			}
+
 		})
 		if err != nil {
 			return err
@@ -84,13 +91,3 @@ func main() {
 		fmt.Println(log.LogColor.High("Failed to start afrog，", err.Error()))
 	}
 }
-
-func showBanner(afrogLatestversion string) {
-	title := "NAME:\n   " + log.LogColor.Banner(runner.ShowBanner()) + " - v" + config.Version
-	old := ""
-	if utils.Compare(afrogLatestversion, ">", config.Version) {
-		old = log.LogColor.High(" (outdated)")
-		old += log.LogColor.Title(" --> https://github.com/zan8in/afrog/releases/tag/v" + afrogLatestversion)
-	}
-	fmt.Println(title + old + "\n")
-}

internal/runner/banner.go

@@ -1,9 +1,27 @@
 package runner
 
+import (
+	"fmt"
+
+	"github.com/zan8in/afrog/pkg/config"
+	"github.com/zan8in/afrog/pkg/log"
+	"github.com/zan8in/afrog/pkg/utils"
+)
+
 func ShowBanner() string {
 	return "afrog"
 }
 
 func ShowUsage() string {
 	return "\nUSAGE:\n   afrog -t example.com -o result.html\n   afrog -T urls.txt -o result.html\n   afrog -T urls.txt -s -o result.html\n   afrog -t example.com -P ./pocs/poc-test.yaml -o result.html\n   afrog -t example.com -P ./pocs/ -o result.html\n"
 }
+
+func ShowBanner2(afrogLatestversion string) {
+	title := "NAME:\n   " + log.LogColor.Banner(ShowBanner()) + " - v" + config.Version
+	old := ""
+	if utils.Compare(afrogLatestversion, ">", config.Version) {
+		old = log.LogColor.High(" (outdated)")
+		old += log.LogColor.Title(" --> https://github.com/zan8in/afrog/releases/tag/v" + afrogLatestversion)
+	}
+	fmt.Println(title + old + "\n")
+}

pkg/core/result.go

@@ -80,8 +80,8 @@ func (r *Result) PrintResultInfo() string {
 	return "[" + utils.GetNowDateTime() + "] [" + r.PocInfo.Id + "] [" + r.PocInfo.Info.Severity + "] " + r.Target
 }
 
-func (r *Result) PrintColorResultInfoConsole() {
-	fmt.Printf("\r" + log.LogColor.Time(utils.GetNowDateTime()) + " " +
+func (r *Result) PrintColorResultInfoConsole(number string) {
+	fmt.Printf("\r" + log.LogColor.Time(number+" "+utils.GetNowDateTime()) + " " +
 		log.LogColor.Vulner(""+r.PocInfo.Id+"") + " " +
 		log.LogColor.GetColor(r.PocInfo.Info.Severity, ""+
 			strings.ToUpper(r.PocInfo.Info.Severity)+"") + " " + r.Target + "\r\n")

pkg/html/html.go

@@ -15,6 +15,7 @@ import (
 type HtmlTemplate struct {
 	Result   *core.Result
 	Filename string
+	Number   string
 }
 
 const outputDirectory = "./reports"
@@ -48,10 +49,10 @@ func (ht *HtmlTemplate) Html() string {
 	}
 	title := fmt.Sprintf(`<table>
 	<thead onclick="$(this).next('tbody').toggle()" style="background:#f5f5f5">
-		<td class="vuln">%s</td>
+		<td class="vuln">%s&nbsp;&nbsp;%s</td>
 		<td class="security %s">%s</td>
 		<td class="url">%s</td>
-	</thead>`, htResult.PocInfo.Id, htResult.PocInfo.Info.Severity, strings.ToUpper(htResult.PocInfo.Info.Severity), htResult.Target)
+	</thead>`, ht.Number, htResult.PocInfo.Id, htResult.PocInfo.Info.Severity, strings.ToUpper(htResult.PocInfo.Info.Severity), htResult.Target)
 
 	info := fmt.Sprintf("<b>name:</b> %s&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<b>author:</b> %s&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<b>security:</b> %s",
 		htResult.PocInfo.Info.Name, htResult.PocInfo.Info.Author, htResult.PocInfo.Info.Severity,
@@ -74,7 +75,30 @@ func (ht *HtmlTemplate) Html() string {
 
 	body := ""
 	for _, v := range htResult.AllPocResult {
-		fullurl := fmt.Sprintf("%s://%s%s", v.ResultRequest.Url.Scheme, v.ResultRequest.Url.Host, v.ResultRequest.Url.Path)
+		if !v.IsVul {
+			continue
+		}
+		schema := ""
+		host := ""
+		path := ""
+		query := ""
+		frament := ""
+		reqraw := []byte{}
+		respraw := []byte{}
+		if v.ResultRequest.Url != nil {
+			schema = v.ResultRequest.Url.Scheme
+			host = v.ResultRequest.Url.Host
+			path = v.ResultRequest.Url.Path
+			if len(v.ResultRequest.Url.Query) > 0 {
+				query = "?" + v.ResultRequest.Url.Query
+			}
+			if len(v.ResultRequest.Url.Fragment) > 0 {
+				frament = "#" + v.ResultRequest.Url.Fragment
+			}
+			reqraw = v.ResultRequest.GetRaw()
+			respraw = v.ResultResponse.GetRaw()
+		}
+		fullurl := fmt.Sprintf("%s://%s%s%s%s", schema, host, path, query, frament)
 		body += fmt.Sprintf(`<tr>
 		<td colspan="3" style="background:#f8f8f8"><a href="%s" target="_blank">%s</a></td>
 	</tr><tr>
@@ -91,7 +115,7 @@ func (ht *HtmlTemplate) Html() string {
 			</div>
 			</td>
 		</tr>
-	`, fullurl, fullurl, v.ResultRequest.GetRaw(), v.ResultResponse.GetRaw())
+	`, fullurl, fullurl, reqraw, respraw)
 	}
 
 	footer := "</tbody></table>"
@@ -102,8 +126,10 @@ func (ht *HtmlTemplate) Html() string {
 func (ht *HtmlTemplate) Append() {
 	r := ht.Html()
 	if len(r) > 0 {
-		utils.BufferWriteAppend(ht.Filename, r)
+		utils.AppendString(ht.Filename, r)
+		// fmt.Println(err)
 	}
+	// fmt.Println(len(r))
 }
 
 func header() string {

pkg/utils/file.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
@@ -81,7 +82,7 @@ func WriteFile(filename string, data []byte) error {
 }
 
 func BufferWriteAppend(filename string, param string) error {
-	fileHandle, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
+	fileHandle, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_APPEND|os.O_SYNC, 0660)
 	if err != nil {
 		return err
 	}
@@ -96,3 +97,26 @@ func BufferWriteAppend(filename string, param string) error {
 	// 将缓冲中的数据写入
 	return buf.Flush()
 }
+
+const (
+	NEW_FILE_PERM = 0666
+)
+
+// AppendString appends the contents of the string to filename.
+func AppendString(filename, content string) error {
+	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, NEW_FILE_PERM)
+	if err != nil {
+		return err
+	}
+	data := []byte(content)
+	n, err := f.Write(data)
+	if err == nil && n < len(data) {
+		err = io.ErrShortWrite
+		// fmt.Println(err)
+	}
+	if err1 := f.Close(); err == nil {
+		err = err1
+		// fmt.Println(err)
+	}
+	return err
+}

pkg/utils/utils.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -151,3 +152,13 @@ func GetNowDateTime() string {
 	now := time.Now()
 	return now.Format("2006-01-02 15:04:05")
 }
+
+func GetNumberText(number int) string {
+	num := strconv.Itoa(number)
+	if len(num) == 1 {
+		num = "00" + num
+	} else if len(num) == 2 {
+		num = "0" + num
+	}
+	return num
+}