From e6958d7aead465d254f62d24d7f3e15c280bea2e Mon Sep 17 00:00:00 2001
From: Ice3man543 <nizamulrana@gmail.com>
Date: Mon, 21 Dec 2020 12:04:33 +0530
Subject: [PATCH] Moved colorizer stuff to pkg/output

---
 v2/pkg/colorizer/colorizer.go | 42 -----------------------------------
 v2/pkg/output/output.go       |  7 ++++++
 2 files changed, 7 insertions(+), 42 deletions(-)
 delete mode 100644 v2/pkg/colorizer/colorizer.go

diff --git a/v2/pkg/colorizer/colorizer.go b/v2/pkg/colorizer/colorizer.go
deleted file mode 100644
index 61937bbf..00000000
--- a/v2/pkg/colorizer/colorizer.go
+++ /dev/null
@@ -1,42 +0,0 @@
-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
-}
diff --git a/v2/pkg/output/output.go b/v2/pkg/output/output.go
index 42edf4bd..efaf410d 100644
--- a/v2/pkg/output/output.go
+++ b/v2/pkg/output/output.go
@@ -13,6 +13,8 @@ import (
 type Writer interface {
 	// Close closes the output writer interface
 	Close()
+	// Colorizer returns the colorizer instance for writer
+	Colorizer() aurora.Aurora
 	// Write writes the event to file and/or screen.
 	Write(Event) error
 	// Request writes a log the requests trace log
@@ -135,6 +137,11 @@ func (w *StandardWriter) Request(templateID, url, requestType string, err error)
 	w.traceMutex.Unlock()
 }
 
+// Colorizer returns the colorizer instance for writer
+func (w *StandardWriter) Colorizer() aurora.Aurora {
+	return w.aurora
+}
+
 // Close closes the output writing interface
 func (w *StandardWriter) Close() {
 	if w.outputFile != nil {