package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"

	"github.com/pkg/errors"
	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/disk"
	"github.com/projectdiscovery/nvd"
	"github.com/projectdiscovery/sliceutil"
	"github.com/projectdiscovery/stringsutil"
	"gopkg.in/yaml.v3"
)

const (
	yamlIndentSpaces = 2
)

var cisaKnownExploitedVulnerabilities map[string]struct{}

func init() {
	if err := fetchCISAKnownExploitedVulnerabilities(); err != nil {
		panic(err)
	}
}

var (
	input       = flag.String("i", "", "Templates to annotate")
	templateDir = flag.String("d", "", "Custom template directory for update")
)

func main() {
	flag.Parse()

	if *input == "" || *templateDir == "" {
		log.Fatalf("invalid input, see -h\n")
	}

	if err := process(); err != nil {
		log.Fatalf("could not process: %s\n", err)
	}
}

func process() error {
	tempDir, err := os.MkdirTemp("", "nuclei-nvd-%s")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	client, err := nvd.NewClient(tempDir)
	if err != nil {
		return err
	}
	catalog := disk.NewCatalog(*templateDir)

	paths, err := catalog.GetTemplatePath(*input)
	if err != nil {
		return err
	}
	for _, path := range paths {
		data, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		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
}

var (
	idRegex       = regexp.MustCompile("id: ([C|c][V|v][E|e]-[0-9]+-[0-9]+)")
	severityRegex = regexp.MustCompile(`severity: ([a-z]+)`)
)

const maxReferenceCount = 5

// dead sites to skip for references
var badRefs = []string{
	"osvdb.org/",
	"securityfocus.com/",
	"archives.neohapsis.com/",
	"iss.net/",
	"ntelbras.com/",
	"andmp.com/",
	"blacklanternsecurity.com/",
	"pwnwiki.org/",
	"0dayhack.net/",
	"correkt.horse/",
	"poc.wgpsec.org/",
	"ctf-writeup.revers3c.com/",
	"secunia.com/",
}

func getCVEData(client *nvd.Client, filePath, data string) {
	matches := idRegex.FindAllStringSubmatch(data, 1)
	if len(matches) == 0 {
		return
	}
	cveName := matches[0][1]

	// Perform CISA Known-exploited-vulnerabilities tag annotation
	// if we discover it has been exploited.
	var err error
	if cisaKnownExploitedVulnerabilities != nil {
		_, ok := cisaKnownExploitedVulnerabilities[strings.ToLower(cveName)]
		if ok {
			data, err = parseAndAddCISAKevTagTemplate(filePath, data)
		}
	}
	if err != nil {
		log.Printf("Could not parse cisa data %s: %s\n", cveName, err)
		return
	}

	severityMatches := severityRegex.FindAllStringSubmatch(data, 1)
	if len(severityMatches) == 0 {
		return
	}
	severityValue := severityMatches[0][1]

	cveItem, err := client.FetchCVE(cveName)
	if err != nil {
		log.Printf("Could not fetch cve %s: %s\n", cveName, err)
		return
	}
	var cweID []string
	for _, problemData := range cveItem.CVE.Problemtype.ProblemtypeData {
		for _, description := range problemData.Description {
			cweID = append(cweID, description.Value)
		}
	}
	cvssScore := cveItem.Impact.BaseMetricV3.CvssV3.BaseScore
	cvssMetrics := cveItem.Impact.BaseMetricV3.CvssV3.VectorString

	// Perform some hacky string replacement to place the metadata in templates
	infoBlockIndexData := data[strings.Index(data, "info:"):]
	requestsIndex := strings.Index(infoBlockIndexData, "requests:")
	networkIndex := strings.Index(infoBlockIndexData, "network:")
	variablesIndex := strings.Index(infoBlockIndexData, "variables:")
	if requestsIndex == -1 && networkIndex == -1 && variablesIndex == -1 {
		return
	}
	if networkIndex != -1 {
		requestsIndex = networkIndex
	}
	if variablesIndex != -1 {
		requestsIndex = variablesIndex
	}
	infoBlockData := infoBlockIndexData[:requestsIndex]
	infoBlockClean := strings.TrimRight(infoBlockData, "\n")

	infoBlock := InfoBlock{}
	err = yaml.Unmarshal([]byte(data), &infoBlock)
	if err != nil {
		log.Printf("Could not unmarshal info block: %s\n", err)
	}

	var changed bool
	if newSeverity := isSeverityMatchingCvssScore(severityValue, cvssScore); newSeverity != "" {
		changed = true
		infoBlock.Info.Severity = newSeverity
		fmt.Printf("Adjusting severity for %s from %s=>%s (%.2f)\n", filePath, severityValue, newSeverity, cvssScore)
	}
	isCvssEmpty := cvssScore == 0 || cvssMetrics == ""
	hasCvssChanged := infoBlock.Info.Classification.CvssScore != cvssScore || cvssMetrics != infoBlock.Info.Classification.CvssMetrics
	if !isCvssEmpty && hasCvssChanged {
		changed = true
		infoBlock.Info.Classification.CvssMetrics = cvssMetrics
		infoBlock.Info.Classification.CvssScore = cvssScore
		infoBlock.Info.Classification.CveId = cveName
		if len(cweID) > 0 && (cweID[0] != "NVD-CWE-Other" && cweID[0] != "NVD-CWE-noinfo") {
			infoBlock.Info.Classification.CweId = strings.Join(cweID, ",")
		}
	}
	// If there is no description field, fill the description from CVE information
	hasDescriptionData := len(cveItem.CVE.Description.DescriptionData) > 0
	isDescriptionEmpty := infoBlock.Info.Description == ""
	if isDescriptionEmpty && hasDescriptionData {
		changed = true
		// removes all new lines
		description := stringsutil.ReplaceAny(cveItem.CVE.Description.DescriptionData[0].Value, "", "\n", "\\", "'", "\t")
		description += "\n"
		infoBlock.Info.Description = description
	}

	// we are unmarshaling info block to have valid data
	var referenceDataURLs []string

	// skip sites that are no longer alive
	for _, reference := range cveItem.CVE.References.ReferenceData {
		if stringsutil.ContainsAny(reference.URL, badRefs...) {
			continue
		}
		referenceDataURLs = append(referenceDataURLs, reference.URL)
	}
	hasReferenceData := len(cveItem.CVE.References.ReferenceData) > 0
	areCveReferencesContained := sliceutil.ContainsItems(infoBlock.Info.Reference, referenceDataURLs)
	referencesCount := len(infoBlock.Info.Reference)
	if hasReferenceData && !areCveReferencesContained {
		changed = true
		for _, ref := range referenceDataURLs {
			referencesCount++
			if referencesCount >= maxReferenceCount {
				break
			}
			infoBlock.Info.Reference = append(infoBlock.Info.Reference, ref)
		}
		infoBlock.Info.Reference = sliceutil.PruneEmptyStrings(sliceutil.Dedupe(infoBlock.Info.Reference))
	}

	var newInfoBlock bytes.Buffer
	yamlEncoder := yaml.NewEncoder(&newInfoBlock)
	yamlEncoder.SetIndent(yamlIndentSpaces)
	err = yamlEncoder.Encode(infoBlock)
	if err != nil {
		log.Printf("Could not marshal info block: %s\n", err)
		return
	}
	newInfoBlockData := strings.TrimSuffix(newInfoBlock.String(), "\n")

	newTemplate := strings.ReplaceAll(data, infoBlockClean, newInfoBlockData)
	if changed {
		_ = os.WriteFile(filePath, []byte(newTemplate), 0644)
		fmt.Printf("Wrote updated template to %s\n", filePath)
	}
}

func isSeverityMatchingCvssScore(severity string, score float64) string {
	if score == 0.0 {
		return ""
	}
	var expected string

	if score >= 0.1 && score <= 3.9 {
		expected = "low"
	} else if score >= 4.0 && score <= 6.9 {
		expected = "medium"
	} else if score >= 7.0 && score <= 8.9 {
		expected = "high"
	} else if score >= 9.0 && score <= 10.0 {
		expected = "critical"
	}
	if expected != "" && expected != severity {
		return expected
	}
	return ""
}

type cisaKEVData struct {
	Vulnerabilities []struct {
		CVEID string `json:"cveID"`
	}
}

// fetchCISAKnownExploitedVulnerabilities fetches CISA known exploited
// vulnerabilities catalog for template tag enrichment
func fetchCISAKnownExploitedVulnerabilities() error {
	data := &cisaKEVData{}

	resp, err := http.Get("https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json")
	if err != nil {
		return errors.Wrap(err, "could not get cisa kev catalog")
	}
	defer resp.Body.Close()

	if err := json.NewDecoder(resp.Body).Decode(data); err != nil {
		return errors.Wrap(err, "could not decode cisa kev catalog json data")
	}
	cisaKnownExploitedVulnerabilities = make(map[string]struct{})
	for _, vuln := range data.Vulnerabilities {
		cisaKnownExploitedVulnerabilities[strings.ToLower(vuln.CVEID)] = struct{}{}
	}
	return nil
}

// parseAndAddCISAKevTagTemplate parses and adds `kev` tag to CISA KEV templates.
// also removes cisa tag if it exists
func parseAndAddCISAKevTagTemplate(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
	}

	var cisaIndex = -1
	for i, tag := range splitted {
		// If we already have tag, return
		if tag == "kev" {
			return data, nil
		}
		if tag == "cisa" {
			cisaIndex = i
		}
	}
	// Remove CISA index tag element
	if cisaIndex >= 0 {
		splitted = append(splitted[:cisaIndex], splitted[cisaIndex+1:]...)
	}
	splitted = append(splitted, "kev")
	replaced := strings.ReplaceAll(data, block.Info.Tags, strings.Join(splitted, ","))
	return replaced, os.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, os.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"`
}

type TemplateClassification struct {
	CvssMetrics string  `yaml:"cvss-metrics,omitempty"`
	CvssScore   float64 `yaml:"cvss-score,omitempty"`
	CveId       string  `yaml:"cve-id,omitempty"`
	CweId       string  `yaml:"cwe-id,omitempty"`
}

type TemplateInfo struct {
	Name           string                 `yaml:"name"`
	Author         string                 `yaml:"author"`
	Severity       string                 `yaml:"severity"`
	Description    string                 `yaml:"description,omitempty"`
	Reference      []string               `yaml:"reference,omitempty"`
	Remediation    string                 `yaml:"remediation,omitempty"`
	Classification TemplateClassification `yaml:"classification,omitempty"`
	Metadata       map[string]string      `yaml:"metadata,omitempty"`
	Tags           string                 `yaml:"tags,omitempty"`
}