From 8878dbcddcbfc540b30a0ade656df33713269d6f Mon Sep 17 00:00:00 2001
From: kchason <keith.chason@gmail.com>
Date: Wed, 19 Jul 2023 17:02:59 -0400
Subject: [PATCH] Sort markdown exports by host, severity, or template

---
 v2/internal/runner/options.go                 |  7 ++++
 v2/internal/runner/runner.go                  |  2 ++
 .../reporting/exporters/markdown/markdown.go  | 35 +++++++++++++++++--
 v2/pkg/types/types.go                         |  2 ++
 4 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go
index 3da37f45c4..b0c61e4e1d 100644
--- a/v2/internal/runner/options.go
+++ b/v2/internal/runner/options.go
@@ -407,6 +407,13 @@ func readEnvInputVars(options *types.Options) {
 	options.GitLabTemplateDisableDownload = getBoolEnvValue("DISABLE_NUCLEI_TEMPLATES_GITLAB_DOWNLOAD")
 	options.AwsTemplateDisableDownload = getBoolEnvValue("DISABLE_NUCLEI_TEMPLATES_AWS_DOWNLOAD")
 	options.AzureTemplateDisableDownload = getBoolEnvValue("DISABLE_NUCLEI_TEMPLATES_AZURE_DOWNLOAD")
+
+	// Options to modify the behavior of exporters
+	options.MarkdownExportSortMode = strings.ToLower(os.Getenv("MARKDOWN_EXPORT_SORT_MODE"))
+	// If the user has not specified a valid sort mode, use the default
+	if options.MarkdownExportSortMode != "template" && options.MarkdownExportSortMode != "severity" && options.MarkdownExportSortMode != "host" {
+		options.MarkdownExportSortMode = "none"
+	}
 }
 
 func getBoolEnvValue(key string) bool {
diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go
index 711ea769e2..01c49c9e95 100644
--- a/v2/internal/runner/runner.go
+++ b/v2/internal/runner/runner.go
@@ -346,12 +346,14 @@ func createReportingOptions(options *types.Options) (*reporting.Options, error)
 			reportingOptions.MarkdownExporter = &markdown.Options{
 				Directory:         options.MarkdownExportDirectory,
 				IncludeRawPayload: !options.OmitRawRequests,
+				SortMode:          options.MarkdownExportSortMode,
 			}
 		} else {
 			reportingOptions = &reporting.Options{}
 			reportingOptions.MarkdownExporter = &markdown.Options{
 				Directory:         options.MarkdownExportDirectory,
 				IncludeRawPayload: !options.OmitRawRequests,
+				SortMode:          options.MarkdownExportSortMode,
 			}
 		}
 	}
diff --git a/v2/pkg/reporting/exporters/markdown/markdown.go b/v2/pkg/reporting/exporters/markdown/markdown.go
index 362c65718e..fd8070a55f 100644
--- a/v2/pkg/reporting/exporters/markdown/markdown.go
+++ b/v2/pkg/reporting/exporters/markdown/markdown.go
@@ -2,6 +2,7 @@ package markdown
 
 import (
 	"bytes"
+	"github.com/projectdiscovery/gologger"
 	"os"
 	"path/filepath"
 	"strings"
@@ -25,6 +26,7 @@ type Options struct {
 	// Directory is the directory to export found results to
 	Directory         string `yaml:"directory"`
 	IncludeRawPayload bool   `yaml:"include-raw-payload"`
+	SortMode          string `yaml:"sort-mode"`
 }
 
 // New creates a new markdown exporter integration client based on options.
@@ -69,7 +71,36 @@ func (exporter *Exporter) Export(event *output.ResultEvent) error {
 	defer file.Close()
 
 	filename := createFileName(event)
-	host := util.CreateLink(event.Host, filename)
+
+	// If the sort mode is set to severity, host, or template, then we need to get a safe version of the name for a
+	// subdirectory to store the file in.
+	// This will allow us to sort the files into subdirectories based on the sort mode. The subdirectory will need to
+	// be created if it does not exist.
+	fileUrl := filename
+	subdirectory := ""
+	switch exporter.options.SortMode {
+	case "severity":
+		subdirectory = event.Info.SeverityHolder.Severity.String()
+	case "host":
+		subdirectory = event.Host
+	case "template":
+		subdirectory = event.TemplateID
+	}
+	if subdirectory != "" {
+		// Sanitize the subdirectory name to remove any characters that are not allowed in a directory name
+		subdirectory = sanitizeFilename(subdirectory)
+
+		// Prepend the subdirectory name to the filename for the fileUrl
+		fileUrl = filepath.Join(subdirectory, filename)
+
+		// Create the subdirectory if it does not exist
+		err = os.MkdirAll(filepath.Join(exporter.directory, subdirectory), 0755)
+		if err != nil {
+			gologger.Warning().Msgf("Could not create subdirectory for markdown report: %s", err)
+		}
+	}
+
+	host := util.CreateLink(event.Host, fileUrl)
 	finding := event.TemplateID + " " + event.MatcherName
 	severity := event.Info.SeverityHolder.Severity.String()
 
@@ -85,7 +116,7 @@ func (exporter *Exporter) Export(event *output.ResultEvent) error {
 	dataBuilder.WriteString(format.CreateReportDescription(event, util.MarkdownFormatter{}))
 	data := dataBuilder.Bytes()
 
-	return os.WriteFile(filepath.Join(exporter.directory, filename), data, 0644)
+	return os.WriteFile(filepath.Join(exporter.directory, subdirectory, filename), data, 0644)
 }
 
 func createFileName(event *output.ResultEvent) string {
diff --git a/v2/pkg/types/types.go b/v2/pkg/types/types.go
index 10411ea3f0..48347a6008 100644
--- a/v2/pkg/types/types.go
+++ b/v2/pkg/types/types.go
@@ -93,6 +93,8 @@ type Options struct {
 	ReportingConfig string
 	// MarkdownExportDirectory is the directory to export reports in Markdown format
 	MarkdownExportDirectory string
+	// MarkdownExportSortMode is the method to sort the markdown reports (options: severity, template, host, none)
+	MarkdownExportSortMode string
 	// SarifExport is the file to export sarif output format to
 	SarifExport string
 	// CloudURL is the URL for the nuclei cloud endpoint