Added reference based tag addition to templates (#2464)

* Added reference based tag addition to templates

* reference mapping list update

* Misc changes as per review

Co-authored-by: sandeep <sandeep@projectdiscovery.io>

v2/cmd/cve-annotate/main.go

@@ -8,6 +8,7 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
@@ -71,7 +72,16 @@ func process() error {
 		if err != nil {
 			return err
 		}
-		getCVEData(client, path, string(data))
+		dataString := string(data)
+
+		// First try to resolve references to tags
+		dataString, err = parseAndAddReferenceBasedTags(path, dataString)
+		if err != nil {
+			log.Printf("Could not parse reference tags %s: %s\n", path, err)
+			continue
+		}
+		// Next try and fill CVE data
+		getCVEData(client, path, dataString)
 	}
 	return nil
 }
@@ -290,12 +300,15 @@ func parseAndAddCISAKevTagTemplate(path string, data string) (string, error) {
 		return "", errors.Wrap(err, "could not decode template yaml")
 	}
 	splitted := strings.Split(block.Info.Tags, ",")
+	if len(splitted) == 0 {
+		return data, nil
+	}
 
 	var cisaIndex = -1
 	for i, tag := range splitted {
 		// If we already have tag, return
 		if tag == "kev" {
-			return "", nil
+			return data, nil
 		}
 		if tag == "cisa" {
 			cisaIndex = i
@@ -306,11 +319,69 @@ func parseAndAddCISAKevTagTemplate(path string, data string) (string, error) {
 		splitted = append(splitted[:cisaIndex], splitted[cisaIndex+1:]...)
 	}
 	splitted = append(splitted, "kev")
-	final := strings.Join(splitted, ",")
-	replaced := strings.Replace(data, block.Info.Tags, final, -1)
+	replaced := strings.ReplaceAll(data, block.Info.Tags, strings.Join(splitted, ","))
 	return replaced, ioutil.WriteFile(path, []byte(replaced), os.ModePerm)
 }
 
+// parseAndAddReferenceBasedTags parses and adds reference based tags to templates
+func parseAndAddReferenceBasedTags(path string, data string) (string, error) {
+	block := &InfoBlock{}
+	if err := yaml.NewDecoder(strings.NewReader(data)).Decode(block); err != nil {
+		return "", errors.Wrap(err, "could not decode template yaml")
+	}
+	splitted := strings.Split(block.Info.Tags, ",")
+	if len(splitted) == 0 {
+		return data, nil
+	}
+	tagsCurrent := fmt.Sprintf("tags: %s", block.Info.Tags)
+	newTags := suggestTagsBasedOnReference(block.Info.Reference, splitted)
+
+	if len(newTags) == len(splitted) {
+		return data, nil
+	}
+	replaced := strings.ReplaceAll(data, tagsCurrent, fmt.Sprintf("tags: %s", strings.Join(newTags, ",")))
+	return replaced, ioutil.WriteFile(path, []byte(replaced), os.ModePerm)
+}
+
+var referenceMapping = map[string]string{
+	"huntr.dev":               "huntr",
+	"hackerone.com":           "hackerone",
+	"tenable.com":             "tenable",
+	"packetstormsecurity.org": "packetstorm",
+	"seclists.org":            "seclists",
+	"wpscan.com":              "wpscan",
+	"packetstormsecurity.com": "packetstorm",
+	"exploit-db.com":          "edb",
+	"https://github.com/rapid7/metasploit-framework/": "msf",
+	"https://github.com/vulhub/vulhub/":               "vulhub",
+}
+
+func suggestTagsBasedOnReference(references, currentTags []string) []string {
+	uniqueTags := make(map[string]struct{})
+	for _, value := range currentTags {
+		uniqueTags[value] = struct{}{}
+	}
+
+	for _, reference := range references {
+		parsed, err := url.Parse(reference)
+		if err != nil {
+			continue
+		}
+		hostname := parsed.Hostname()
+
+		for value, tag := range referenceMapping {
+			if strings.HasSuffix(hostname, value) || strings.HasPrefix(reference, value) {
+				uniqueTags[tag] = struct{}{}
+			}
+		}
+	}
+	newTags := make([]string, 0, len(uniqueTags))
+	for tag := range uniqueTags {
+		newTags = append(newTags, tag)
+	}
+	return newTags
+}
+
 // Cloning struct from nuclei as we don't want any validation
 type InfoBlock struct {
 	Info TemplateInfo `yaml:"info"`