Merge pull request #315 from vzamanillo/serverity-on-match

Added severity to match output message

internal/progress/progress.go

@@ -21,7 +21,7 @@ const (
 	mili        = 1000.
 )
 
-// Encapsulates progress tracking.
+// IProgress encapsulates progress tracking.
 type IProgress interface {
 	InitProgressbar(hostCount int64, templateCount int, requestCount int64)
 	AddToTotal(delta int64)
@@ -37,7 +37,7 @@ type Progress struct {
 	initialTotal int64
 
 	totalMutex *sync.Mutex
-	colorizer  aurora.Aurora
+	colorizer  *aurora.Aurora
 
 	renderChan         chan time.Time
 	captureData        *captureData
@@ -49,8 +49,8 @@ type Progress struct {
 	stdRenderWaitGroup *sync.WaitGroup
 }
 
-// Creates and returns a new progress tracking object.
+// NewProgress creates and returns a new progress tracking object.
-func NewProgress(noColor, active bool) IProgress {
+func NewProgress(colorizer aurora.Aurora, active bool) IProgress {
 	if !active {
 		return &NoOpProgress{}
 	}
@@ -65,7 +65,7 @@ func NewProgress(noColor, active bool) IProgress {
 			mpb.WithManualRefresh(renderChan),
 		),
 		totalMutex: &sync.Mutex{},
-		colorizer:  aurora.NewAurora(!noColor),
+		colorizer:  &colorizer,
 
 		renderChan:         renderChan,
 		stdCaptureMutex:    &sync.Mutex{},
@@ -85,7 +85,7 @@ func (p *Progress) InitProgressbar(hostCount int64, rulesCount int, requestCount
 		panic("A global progressbar is already present.")
 	}
 
-	color := p.colorizer
+	color := *p.colorizer
 
 	barName := color.Sprintf(
 		color.Cyan("%d %s, %d %s"),
@@ -193,7 +193,7 @@ func (p *Progress) renderStdData() {
 
 // Creates and returns a progress bar.
 func (p *Progress) setupProgressbar(name string, total int64, priority int) *mpb.Bar {
-	color := p.colorizer
+	color := *p.colorizer
 
 	p.total = total
 	p.initialTotal = total

internal/runner/processor.go

@@ -64,7 +64,7 @@ func (r *Runner) processTemplateWithList(ctx context.Context, p progress.IProgre
 			JSONRequests:    r.options.JSONRequests,
 			CookieReuse:     value.CookieReuse,
 			ColoredOutput:   !r.options.NoColor,
-			Colorizer:       r.colorizer,
+			Colorizer:       &r.colorizer,
 			Decolorizer:     r.decolorizer,
 		})
 	}
@@ -233,7 +233,7 @@ func (r *Runner) preloadWorkflowTemplates(p progress.IProgress, workflow *workfl
 					JSONRequests:  r.options.JSONRequests,
 					CookieJar:     jar,
 					ColoredOutput: !r.options.NoColor,
-					Colorizer:     r.colorizer,
+					Colorizer:     &r.colorizer,
 					Decolorizer:   r.decolorizer,
 				}
 			} else if len(t.RequestsDNS) > 0 {

internal/runner/runner.go

@@ -16,6 +16,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/internal/bufwriter"
 	"github.com/projectdiscovery/nuclei/v2/internal/progress"
 	"github.com/projectdiscovery/nuclei/v2/pkg/atomicboolean"
+	"github.com/projectdiscovery/nuclei/v2/pkg/colorizer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v2/pkg/workflows"
 )
@@ -38,7 +39,7 @@ type Runner struct {
 	progress progress.IProgress
 
 	// output coloring
-	colorizer   aurora.Aurora
+	colorizer   colorizer.NucleiColorizer
 	decolorizer *regexp.Regexp
 }
 
@@ -52,6 +53,15 @@ func New(options *Options) (*Runner, error) {
 		gologger.Labelf("Could not update templates: %s\n", err)
 	}
 
+	// output coloring
+	useColor := !options.NoColor
+	runner.colorizer = *colorizer.NewNucleiColorizer(aurora.NewAurora(useColor))
+
+	if useColor {
+		// compile a decolorization regex to cleanup file output messages
+		runner.decolorizer = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`)
+	}
+
 	if options.TemplateList {
 		runner.listAvailableTemplates()
 		os.Exit(0)
@@ -65,15 +75,6 @@ func New(options *Options) (*Runner, error) {
 		runner.readNucleiIgnoreFile()
 	}
 
-	// output coloring
-	useColor := !options.NoColor
-	runner.colorizer = aurora.NewAurora(useColor)
-
-	if useColor {
-		// compile a decolorization regex to cleanup file output messages
-		runner.decolorizer = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`)
-	}
-
 	// If we have stdin, write it to a new file
 	if options.Stdin {
 		tempInput, err := ioutil.TempFile("", "stdin-input-*")
@@ -159,7 +160,7 @@ func New(options *Options) (*Runner, error) {
 	}
 
 	// Creates the progress tracking object
-	runner.progress = progress.NewProgress(runner.options.NoColor, options.EnableProgressBar)
+	runner.progress = progress.NewProgress(runner.colorizer.Colorizer, options.EnableProgressBar)
 
 	runner.limiter = make(chan struct{}, options.Threads)
 
@@ -211,9 +212,9 @@ func (r *Runner) RunEnumeration() {
 	}
 
 	gologger.Infof("Using %s rules (%s templates, %s workflows)",
-		r.colorizer.Bold(templateCount).String(),
+		r.colorizer.Colorizer.Bold(templateCount).String(),
-		r.colorizer.Bold(templateCount-workflowCount).String(),
+		r.colorizer.Colorizer.Bold(templateCount-workflowCount).String(),
-		r.colorizer.Bold(workflowCount).String())
+		r.colorizer.Colorizer.Bold(workflowCount).String())
 
 	// precompute total request count
 	var totalRequests int64 = 0

internal/runner/templates.go

@@ -8,22 +8,11 @@ import (
 	"strings"
 
 	"github.com/karrick/godirwalk"
-	"github.com/logrusorgru/aurora"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v2/pkg/workflows"
 )
 
-const fgOrange uint8 = 208
-
-var severityMap = map[string]string{
-	"info":     aurora.Blue("info").String(),
-	"low":      aurora.Green("low").String(),
-	"medium":   aurora.Yellow("medium").String(),
-	"high":     aurora.Index(fgOrange, "high").String(),
-	"critical": aurora.Red("critical").String(),
-}
-
 // getTemplatesFor parses the specified input template definitions and returns a list of unique, absolute template paths.
 func (r *Runner) getTemplatesFor(definitions []string) []string {
 	// keeps track of processed dirs and files
@@ -193,12 +182,12 @@ func (r *Runner) parseTemplateFile(file string) (interface{}, error) {
 func (r *Runner) templateLogMsg(id, name, author, severity string) string {
 	// Display the message for the template
 	message := fmt.Sprintf("[%s] %s (%s)",
-		r.colorizer.BrightBlue(id).String(),
+		r.colorizer.Colorizer.BrightBlue(id).String(),
-		r.colorizer.Bold(name).String(),
+		r.colorizer.Colorizer.Bold(name).String(),
-		r.colorizer.BrightYellow("@"+author).String())
+		r.colorizer.Colorizer.BrightYellow("@"+author).String())
 
 	if severity != "" {
-		message += " [" + severityMap[strings.ToLower(severity)] + "]"
+		message += " [" + r.colorizer.GetColorizedSeverity(severity) + "]"
 	}
 
 	return message
@@ -234,12 +223,11 @@ func (r *Runner) listAvailableTemplates() {
 		r.templatesConfig.CurrentVersion,
 		r.templatesConfig.TemplatesDirectory,
 	)
-	r.colorizer = aurora.NewAurora(true)
 	err := directoryWalker(
 		r.templatesConfig.TemplatesDirectory,
 		func(path string, d *godirwalk.Dirent) error {
 			if d.IsDir() && path != r.templatesConfig.TemplatesDirectory {
-				gologger.Silentf("\n%s:\n\n", r.colorizer.Bold(r.colorizer.BgBrightBlue(d.Name())).String())
+				gologger.Silentf("\n%s:\n\n", r.colorizer.Colorizer.Bold(r.colorizer.Colorizer.BgBrightBlue(d.Name())).String())
 			} else if strings.HasSuffix(path, ".yaml") {
 				r.logAvailableTemplate(path)
 			}

pkg/colorizer/colorizer.go

@@ -0,0 +1,42 @@
+package colorizer
+
+import (
+	"strings"
+
+	"github.com/logrusorgru/aurora"
+)
+
+const (
+	fgOrange  uint8  = 208
+	undefined string = "undefined"
+)
+
+// NucleiColorizer contains the severity color mapping
+type NucleiColorizer struct {
+	Colorizer   aurora.Aurora
+	SeverityMap map[string]string
+}
+
+// NewNucleiColorizer initializes the new nuclei colorizer
+func NewNucleiColorizer(colorizer aurora.Aurora) *NucleiColorizer {
+	return &NucleiColorizer{
+		Colorizer: colorizer,
+		SeverityMap: map[string]string{
+			"info":     colorizer.Blue("info").String(),
+			"low":      colorizer.Green("low").String(),
+			"medium":   colorizer.Yellow("medium").String(),
+			"high":     colorizer.Index(fgOrange, "high").String(),
+			"critical": colorizer.Red("critical").String(),
+		},
+	}
+}
+
+// GetColorizedSeverity returns the colorized severity string
+func (r *NucleiColorizer) GetColorizedSeverity(severity string) string {
+	sev := r.SeverityMap[strings.ToLower(severity)]
+	if sev == "" {
+		return undefined
+	}
+
+	return sev
+}

pkg/executer/executer_dns.go

@@ -5,11 +5,11 @@ import (
 	"os"
 	"regexp"
 
-	"github.com/logrusorgru/aurora"
 	"github.com/pkg/errors"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/internal/bufwriter"
 	"github.com/projectdiscovery/nuclei/v2/internal/progress"
+	"github.com/projectdiscovery/nuclei/v2/pkg/colorizer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/matchers"
 	"github.com/projectdiscovery/nuclei/v2/pkg/requests"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
@@ -29,7 +29,7 @@ type DNSExecuter struct {
 	dnsRequest    *requests.DNSRequest
 	writer        *bufwriter.Writer
 
-	colorizer   aurora.Aurora
+	colorizer   colorizer.NucleiColorizer
 	decolorizer *regexp.Regexp
 }
 
@@ -51,7 +51,7 @@ type DNSOptions struct {
 	DNSRequest    *requests.DNSRequest
 	Writer        *bufwriter.Writer
 
-	Colorizer   aurora.Aurora
+	Colorizer   colorizer.NucleiColorizer
 	Decolorizer *regexp.Regexp
 }
 

pkg/executer/executer_http.go

@@ -16,12 +16,11 @@ import (
 	"strings"
 	"time"
 
-	"github.com/logrusorgru/aurora"
-
 	"github.com/pkg/errors"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/internal/bufwriter"
 	"github.com/projectdiscovery/nuclei/v2/internal/progress"
+	"github.com/projectdiscovery/nuclei/v2/pkg/colorizer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/matchers"
 	"github.com/projectdiscovery/nuclei/v2/pkg/requests"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
@@ -49,7 +48,7 @@ type HTTPExecuter struct {
 	customHeaders   requests.CustomHeaders
 	CookieJar       *cookiejar.Jar
 
-	colorizer   aurora.Aurora
+	colorizer   colorizer.NucleiColorizer
 	decolorizer *regexp.Regexp
 }
 
@@ -69,7 +68,7 @@ type HTTPOptions struct {
 	ProxySocksURL   string
 	CustomHeaders   requests.CustomHeaders
 	CookieJar       *cookiejar.Jar
-	Colorizer       aurora.Aurora
+	Colorizer       *colorizer.NucleiColorizer
 	Decolorizer     *regexp.Regexp
 }
 
@@ -114,7 +113,7 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) {
 		customHeaders:   options.CustomHeaders,
 		CookieJar:       options.CookieJar,
 		coloredOutput:   options.ColoredOutput,
-		colorizer:       options.Colorizer,
+		colorizer:       *options.Colorizer,
 		decolorizer:     options.Decolorizer,
 	}
 

pkg/executer/output_dns.go

@@ -57,16 +57,23 @@ func (e *DNSExecuter) writeOutputDNS(domain string, req, resp *dns.Msg, matcher
 	colorizer := e.colorizer
 
 	builder.WriteRune('[')
-	builder.WriteString(colorizer.BrightGreen(e.template.ID).String())
+	builder.WriteString(colorizer.Colorizer.BrightGreen(e.template.ID).String())
 
 	if matcher != nil && len(matcher.Name) > 0 {
 		builder.WriteString(":")
-		builder.WriteString(colorizer.BrightGreen(matcher.Name).Bold().String())
+		builder.WriteString(colorizer.Colorizer.BrightGreen(matcher.Name).Bold().String())
 	}
 
 	builder.WriteString("] [")
-	builder.WriteString(colorizer.BrightBlue("dns").String())
+	builder.WriteString(colorizer.Colorizer.BrightBlue("dns").String())
 	builder.WriteString("] ")
+
+	if e.template.Info.Severity != "" {
+		builder.WriteString("[")
+		builder.WriteString(colorizer.GetColorizedSeverity(e.template.Info.Severity))
+		builder.WriteString("] ")
+	}
+
 	builder.WriteString(domain)
 
 	// If any extractors, write the results
@@ -74,7 +81,7 @@ func (e *DNSExecuter) writeOutputDNS(domain string, req, resp *dns.Msg, matcher
 		builder.WriteString(" [")
 
 		for i, result := range extractorResults {
-			builder.WriteString(colorizer.BrightCyan(result).String())
+			builder.WriteString(colorizer.Colorizer.BrightCyan(result).String())
 
 			if i != len(extractorResults)-1 {
 				builder.WriteRune(',')

pkg/executer/output_http.go

@@ -73,17 +73,23 @@ func (e *HTTPExecuter) writeOutputHTTP(req *requests.HTTPRequest, resp *http.Res
 	colorizer := e.colorizer
 
 	builder.WriteRune('[')
-	builder.WriteString(colorizer.BrightGreen(e.template.ID).String())
+	builder.WriteString(colorizer.Colorizer.BrightGreen(e.template.ID).String())
 
 	if matcher != nil && len(matcher.Name) > 0 {
 		builder.WriteString(":")
-		builder.WriteString(colorizer.BrightGreen(matcher.Name).Bold().String())
+		builder.WriteString(colorizer.Colorizer.BrightGreen(matcher.Name).Bold().String())
 	}
 
 	builder.WriteString("] [")
-	builder.WriteString(colorizer.BrightBlue("http").String())
+	builder.WriteString(colorizer.Colorizer.BrightBlue("http").String())
 	builder.WriteString("] ")
 
+	if e.template.Info.Severity != "" {
+		builder.WriteString("[")
+		builder.WriteString(colorizer.GetColorizedSeverity(e.template.Info.Severity))
+		builder.WriteString("] ")
+	}
+
 	// Escape the URL by replacing all % with %%
 	escapedURL := strings.ReplaceAll(URL, "%", "%%")
 	builder.WriteString(escapedURL)
@@ -93,7 +99,7 @@ func (e *HTTPExecuter) writeOutputHTTP(req *requests.HTTPRequest, resp *http.Res
 		builder.WriteString(" [")
 
 		for i, result := range extractorResults {
-			builder.WriteString(colorizer.BrightCyan(result).String())
+			builder.WriteString(colorizer.Colorizer.BrightCyan(result).String())
 
 			if i != len(extractorResults)-1 {
 				builder.WriteRune(',')
@@ -110,7 +116,7 @@ func (e *HTTPExecuter) writeOutputHTTP(req *requests.HTTPRequest, resp *http.Res
 		var metas []string
 
 		for name, value := range req.Meta {
-			metas = append(metas, colorizer.BrightYellow(name).Bold().String()+"="+colorizer.BrightYellow(value.(string)).String())
+			metas = append(metas, colorizer.Colorizer.BrightYellow(name).Bold().String()+"="+colorizer.Colorizer.BrightYellow(value.(string)).String())
 		}
 
 		builder.WriteString(strings.Join(metas, ","))

pkg/workflows/var.go

@@ -9,6 +9,7 @@ import (
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/internal/progress"
 	"github.com/projectdiscovery/nuclei/v2/pkg/atomicboolean"
+	"github.com/projectdiscovery/nuclei/v2/pkg/colorizer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/executer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/generators"
 )
@@ -76,7 +77,7 @@ func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) {
 				template.HTTPOptions.BulkHTTPRequest = request
 
 				if template.HTTPOptions.Colorizer == nil {
-					template.HTTPOptions.Colorizer = aurora.NewAurora(true)
+					template.HTTPOptions.Colorizer = colorizer.NewNucleiColorizer(aurora.NewAurora(true))
 				}
 
 				httpExecuter, err := executer.NewHTTPExecuter(template.HTTPOptions)