package main

import (
	"encoding/csv"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"sort"
	"strconv"
	"strings"
)

// Example:
// map["rancher/v2.9-head"][0]["image", "release", "package_name", ...  "state", "notes"]
// map["rancher/v2.9-head"][1]["image", "release", "package_name", ...  "state", "notes"]
// ...
// map["harvester/master"][0]["image", "release", "package_name", ...  "state", "notes"]
// ...
type ReportsPerRelease map[string][][]string

// Example:
// map["rancher/v2.9-head"]map["rancher/shell"]["CRITICAL"]3
// map["rancher/v2.9-head"]map["rancher/shell"]["HIGH"]5
// map["rancher/v2.9-head"]map["rancher/shell"]["total"]8
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.20"]["CRITICAL"]2
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.20"]["HIGH"]3
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.20"]["total"]5
// ...
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.25"]["CRITICAL"]1
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.25"]["HIGH"]2
// map["rancher/v2.9-head"]map["rancher/shell:v0.1.25"]["total"]3
// ...
type ReportsImageCVECount map[string]map[string]map[string]int

// Example:
// map["CVE-2024-12345"][0]["CVE-2024-12345", "rancher/shell:v0.1.25", "rancher/v2.9-head", "/bin/sh", "pending-review"]
// map["CVE-2024-12345"][1]["CVE-2024-12345", "rancher/rke-tools:v0.1.100", "rancher/v2.9-head", "/bin/sh", "pending-review"]
// ...
type ReportsPerCVE map[string][][]string

type Reports struct {
	ReportsPerRelease              ReportsPerRelease
	ReportsImageCVECount           ReportsImageCVECount
	ReportsImageAggregatedCVECount ReportsImageCVECount
	ReportsPerCVE                  ReportsPerCVE
	ReportsPerCVESUSEDB            ReportsPerCVE
}

// Headers for the CSV reports.
var CSVHeaderReportPerRelease = []string{"image", "release", "package_name", "package_version", "type", "vulnerability_id", "severity", "url", "target", "patched_version", "mirrored", "status", "justification"}
var CSVHeaderReportPerStats = []string{"image", "critical", "high", "total"}
var CSVHeaderPerCVE = []string{"vulnerability_id", "image", "release", "target", "status"}

func formatReleaseTitle(release string) string {
	str := strings.ReplaceAll(release, "/", " ")
	if strings.HasPrefix(release, "rke2") {
		// If the release is "rke2 vX.Y", then we uppercase it to
		// "RKE2 vX.Y"
		return strings.ToUpper(string(str[0:3])) + str[3:]
	}
	return strings.ToUpper(string(str[0])) + str[1:]
}

func isTier1Image(image string) bool {
	return slices.Contains(CONFIG.Images.Tier1, strings.Split(image, ":")[0])
}

func generateReportsPerCVE(reportsPerCVE ReportsPerCVE, dir string) {
	failed := false

	// Saving the reports per unique CVE.
	for k, v := range reportsPerCVE {
		// The baseDir is used to create a year based directory
		// hierarchy for the CVEs.
		// Example: CVE-2024-12345 or SUSE-SU-2024:12345-1 results in
		// baseDir = 2024.
		baseDir := regexp.MustCompile(`\d{4}`).FindString(k)
		if strings.Contains(k, "GHSA") {
			// GHSA advisories have no year in the advisory ID, so
			// we assign all of them to the `ghsa` baseDir.
			baseDir = "ghsa"
		} else if len(baseDir) == 0 {
			fmt.Printf("Skipping unknown vulnerability identifier - %s\n", k)
			failed = true
			continue
		}
		report := fmt.Sprintf("%s.csv", strings.ToLower(strings.ReplaceAll(k, ":", "-")))
		reportSave(report, filepath.Join(CONFIG.ReportDir, dir, baseDir), v, CSVHeaderPerCVE)
	}

	// We fail and exit in case there is a new CVE identifier, besides
	// `CVE-`, `SUSE-SU-` and `GHSA-`, so we can investigate and update the
	// code if needed.
	if failed {
		os.Exit(1)
	}
}

// Function that generates the reports that are sent to SUSE's CVE database page
// (https://www.suse.com/security/cve/). It's one report per CVE in the format
// of:
//
// CVE-2023-45288,rancher/shell:v0.1.23,Harvester master,"usr/local/bin/helm,usr/local/bin/k9s,usr/local/bin/kubectl,usr/local/bin/kustomize",Affected
// CVE-2023-45288,rancher/shell:v0.1.23,Rancher v2.8.4,"usr/local/bin/helm,usr/local/bin/k9s,usr/local/bin/kubectl,usr/local/bin/kustomize",Affected
// CVE-2023-45288,rancher/shell:v0.1.24,Rancher v2.8.4,"usr/local/bin/k9s,usr/local/bin/kubectl,usr/local/bin/kustomize",Affected
//
// When the data is gathered from the scanning results, the map key is
// calculated as `CVE-image-release` to facilitate the extraction of the
// affected targets per image and release combination. The targets are joined in
// the same line. Note that the title is transformed to the format that we want
// to display in the CVE db page.
// The directory where the reports are generated is `reports/suse-cve-db`.
// Nothing should be changed here without first evaluating possible impacts in
// the CVE db page and without prior alignment with SUSE Product Security team,
// as they develop the script that consumes our reports.
// At the present moment we only generate reports for tier 1 images.
func generateReportsPerCVESUSEDB(reportsPerCVESUSEDB ReportsPerCVE) {
	reportsPerCVE := make(ReportsPerCVE)
	keys := make([]string, 0, len(reportsPerCVESUSEDB))

	for k := range reportsPerCVESUSEDB {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		split := strings.Split(k, ",")
		if len(split) < 3 {
			slog.Debug("unexpected key value", "key", k)
			continue
		}
		image := split[1]
		if !isTier1Image(image) {
			continue
		}

		cve := split[0]
		release := formatReleaseTitle(split[2])

		var affected, notApplicable []string
		for _, target := range reportsPerCVESUSEDB[k] {
			// All CVEs that are not false-positive are considered
			// as applicable (affected).
			if target[1] == CONFIG.CVES.Status.NotAffected {
				notApplicable = append(notApplicable, target[0])
			} else {
				affected = append(affected, target[0])
			}
		}

		if len(affected) > 0 {
			reportsPerCVE[cve] = append(reportsPerCVE[cve],
				[]string{cve, image, release, strings.Join(affected, ","), "Affected"})
		}
		if len(notApplicable) > 0 {
			reportsPerCVE[cve] = append(reportsPerCVE[cve],
				[]string{cve, image, release, strings.Join(notApplicable, ","), "Not applicable"})
		}

	}

	generateReportsPerCVE(reportsPerCVE, "suse-cve-db")
}

func generateReportsPerRelease(images Images, reportsPerRelease ReportsPerRelease, reportsImageCVECount ReportsImageCVECount) {
	for k, v := range reportsPerRelease {
		// Saving the full CVE reports per release.
		report := fmt.Sprintf("report-%s-cves.csv", strings.ReplaceAll(k, "/", "-"))
		reportSave(report, CONFIG.ReportDir, v, CSVHeaderReportPerRelease)

		// Saving the stats reports per release.

		// Adding the images that have no CVE and are not EOL.
		for _, i := range images {
			if !imageIsEOL(i.Image) && slices.Contains(i.ReleaseLabels, "release/"+k) {
				if _, ok := reportsImageCVECount[k][i.Image]; !ok {
					reportsImageCVECount[k][i.Image] = make(map[string]int)
				}
			}
		}

		// Sorting the map's keys (image:tag) so we can write the report
		// based on the alphabetical order of the images.
		sk := make([]string, 0, len(reportsImageCVECount[k]))
		for key := range reportsImageCVECount[k] {
			sk = append(sk, key)
		}
		sort.Strings(sk)

		var records [][]string
		for _, key := range sk {
			vulnCount := reportsImageCVECount[k][key]
			var rec []string
			rec = append(rec, key)
			for _, sev := range strings.Split(CONFIG.Scanner.Severity, ",") {
				rec = append(rec, strconv.Itoa(vulnCount[sev]))
			}
			rec = append(rec, strconv.Itoa(vulnCount["total"]))
			records = append(records, rec)
		}
		report = fmt.Sprintf("report-%s-stats.csv", strings.ReplaceAll(k, "/", "-"))
		reportSave(report, CONFIG.ReportDir, records, CSVHeaderReportPerStats)
	}
}

// all the reports, instead of returning individual reports.
func reportExtractData(cveRecords CvesRecords) Reports {
	var reports Reports
	reportsPerRelease := make(ReportsPerRelease)
	reportsImageCVECount := make(ReportsImageCVECount)
	reportsImageAggregatedCVECount := make(ReportsImageCVECount)
	reportsPerCVE := make(ReportsPerCVE)
	reportsPerCVESUSEDB := make(ReportsPerCVE)

	for _, rec := range cveRecords {
		for _, r := range strings.Split(rec[COL_RELEASE], ",") {
			// Retrieve data for the CVE reports.
			reportsPerRelease[r] = append(reportsPerRelease[r],
				[]string{rec[COL_IMAGE], r, rec[COL_PACKAGE_NAME], rec[COL_PACKAGE_VERSION], rec[COL_TYPE],
					rec[COL_VULNERABILITY_ID], rec[COL_SEVERITY], rec[COL_URL], rec[COL_TARGET],
					rec[COL_PATCHED_VERSION], rec[COL_MIRRORED], rec[COL_STATUS], rec[COL_JUSTIFICATION]})

			// Retrieve data about a CVE and images affected by it.
			target := rec[COL_TARGET]
			if strings.Contains(target, rec[COL_IMAGE]) {
				target = rec[COL_PACKAGE_NAME]
			}
			reportsPerCVE[rec[COL_VULNERABILITY_ID]] =
				append(reportsPerCVE[rec[COL_VULNERABILITY_ID]],
					[]string{rec[COL_VULNERABILITY_ID], rec[COL_IMAGE], r, target, rec[COL_STATUS]})

			// Retrieve data about the combination of
			// CVE-image-release (as a key for the map) and the
			// affected targets/binaries and their states.
			// Any change done to the key format must be reflected
			// in generateReportsPerCVESUSEDB().
			key := fmt.Sprintf("%s,%s,%s", rec[COL_VULNERABILITY_ID], rec[COL_IMAGE], r)
			reportsPerCVESUSEDB[key] = append(reportsPerCVESUSEDB[key], []string{target, rec[COL_STATUS]})

			// From this line below, we skip not affected
			// vulnerabilities from the couting.
			if rec[COL_STATUS] == CONFIG.CVES.Status.NotAffected {
				continue
			}

			// Retrieve data about the vulnerability count for each
			// image:tag combination.
			if _, ok := reportsImageCVECount[r]; !ok {
				reportsImageCVECount[r] = make(map[string]map[string]int)
			}
			if _, ok := reportsImageCVECount[r][rec[COL_IMAGE]]; !ok {
				reportsImageCVECount[r][rec[COL_IMAGE]] = make(map[string]int)
			}
			reportsImageCVECount[r][rec[COL_IMAGE]][rec[COL_SEVERITY]]++
			reportsImageCVECount[r][rec[COL_IMAGE]]["total"]++

			// Aggregate data about the vulnerability count
			// affecting all tags of the same image.
			image := strings.Split(rec[COL_IMAGE], ":")[0]
			if _, ok := reportsImageAggregatedCVECount[r]; !ok {
				reportsImageAggregatedCVECount[r] = make(map[string]map[string]int)
			}
			if _, ok := reportsImageAggregatedCVECount[r][image]; !ok {
				reportsImageAggregatedCVECount[r][image] = make(map[string]int)
			}
			reportsImageAggregatedCVECount[r][image][rec[COL_SEVERITY]]++
			reportsImageAggregatedCVECount[r][image]["total"]++
		}
	}

	reports.ReportsPerRelease = reportsPerRelease
	reports.ReportsImageCVECount = reportsImageCVECount
	reports.ReportsImageAggregatedCVECount = reportsImageAggregatedCVECount
	reports.ReportsPerCVE = reportsPerCVE
	reports.ReportsPerCVESUSEDB = reportsPerCVESUSEDB

	return reports
}

func reportSave(report string, dir string, data [][]string, header []string) {
	var csvFile [][]string

	fmt.Printf("Saving report %s\n", report)

	// Reminder to keep a save default permission.
	err := os.MkdirAll(dir, 0o700)
	checkErrAndExit(err)

	// Reminder to keep a save default permission.
	out, err := os.OpenFile(filepath.Join(dir, report), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	checkErrAndExit(err)
	defer out.Close()

	reportCsv := csv.NewWriter(out)

	csvFile = append(csvFile, header)
	csvFile = append(csvFile, data...)

	err = reportCsv.WriteAll(csvFile)
	checkErrAndExit(err)
}

func mainGenerateReports(cveRecords CvesRecords, images Images) {
	fmt.Println("Generating reports")

	reports := reportExtractData(cveRecords)

	generateReportsPerRelease(images, reports.ReportsPerRelease, reports.ReportsImageCVECount)
	generateReportsPerCVE(reports.ReportsPerCVE, "cves")
	generateReportsPerCVESUSEDB(reports.ReportsPerCVESUSEDB)
}