Optimize cve/generate to use last release (#18269)

cmd/cve/generate.go

@@ -7,12 +7,14 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/fleetdm/fleet/v4/pkg/fleethttp"
 	"github.com/fleetdm/fleet/v4/server/vulnerabilities/nvd"
 	nvdsync "github.com/fleetdm/fleet/v4/server/vulnerabilities/nvd/sync"
 	"github.com/go-kit/log"
@@ -44,11 +46,34 @@ func main() {
 		panic(err)
 	}
 
+	// Download the last released NVD feed
+	logger.Log("msg", "Downloading latest release")
+	maxRetries := 3
+	for i := 0; i < maxRetries; i++ {
+		err := downloadLatestRelease(*dbDir, *debug, logger)
+		if err == nil {
+			break
+		}
+
+		if i == maxRetries-1 {
+			logger.Log("msg", "Failed to download latest release. Continuing with full NVD Sync", "err", err)
+			break
+		}
+
+		logger.Log("msg", "Failed to download latest release. Retrying in 30 seconds", "err", err)
+		time.Sleep(30 * time.Second)
+	}
+
 	// Sync the CVE files
 	if err := nvd.GenerateCVEFeeds(*dbDir, *debug, logger); err != nil {
 		panic(err)
 	}
 
+	// Remove Vulncheck archive
+	if err := os.RemoveAll(filepath.Join(*dbDir, "vulncheck.zip")); err != nil {
+		logger.Log("msg", "Failed to remove vulncheck.zip", "err", err)
+	}
+
 	// Read in every cpe file and create a corresponding metadata file
 	// nvd data feeds start in 2002
 	logger.Log("msg", "Generating metadata files ...")
@@ -76,6 +101,67 @@ func main() {
 	createEmptyFiles(*dbDir, "recent")
 }
 
+func downloadLatestRelease(dbDir string, debug bool, logger log.Logger) error {
+	// Download the latest release
+	err := nvd.DownloadCVEFeed(dbDir, "", debug, logger)
+	if err != nil {
+		return fmt.Errorf("download cve feed: %w", err)
+	}
+
+	// gunzip json files
+	files, err := filepath.Glob(filepath.Join(dbDir, "nvdcve-1.1-*.json.gz"))
+	if err != nil {
+		return fmt.Errorf("glob json files: %w", err)
+	}
+	for _, file := range files {
+		err = gunzipFileToDisk(file, dbDir)
+		if err != nil {
+			return fmt.Errorf("gunzip file %s to disk: %w", file, err)
+		}
+	}
+
+	// Download the last mod start date
+	err = downloadLatestGitHubAsset(dbDir, "last_mod_start_date.txt")
+	if err != nil {
+		return fmt.Errorf("downloading last_mod_start_date asset: %w", err)
+	}
+
+	return nil
+}
+
+// downloadAsset downloads the asset from the latest release and writes it to a file
+func downloadLatestGitHubAsset(dbDir, fileName string) error {
+	assetPath, err := nvd.GetGitHubCVEAssetPath()
+	if err != nil {
+		return fmt.Errorf("get github cve asset path: %w", err)
+	}
+
+	client := fleethttp.NewClient()
+	resp, err := client.Get(assetPath + fileName)
+	if err != nil {
+		return fmt.Errorf("get last mod start date: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("get last mod start date: %w", fmt.Errorf("unexpected status code %d", resp.StatusCode))
+	}
+
+	lastModStartDate, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return fmt.Errorf("read last mod start date: %w", err)
+	}
+
+	// Write the last mod start date to a file
+	lastModStartDateFile := filepath.Join(dbDir, fileName)
+	err = os.WriteFile(lastModStartDateFile, lastModStartDate, 0o644)
+	if err != nil {
+		return fmt.Errorf("write last mod start date: %w", err)
+	}
+
+	return nil
+}
+
 func createMetadata(fileName string, metaName string) {
 	fileInfo, err := os.Stat(fileName)
 	if err != nil {
@@ -152,3 +238,35 @@ func gunzipFileAndComputeSHA256(filename string) (string, error) {
 	defer f.Close()
 	return gunzipAndComputeSHA256(f)
 }
+
+func gunzipFileToDisk(filename, dbpath string) error {
+	f, err := os.Open(filename)
+	if err != nil {
+		return fmt.Errorf("open file: %w", err)
+	}
+	defer f.Close()
+
+	gz, err := gzip.NewReader(f)
+	if err != nil {
+		return fmt.Errorf("new gzip reader: %w", err)
+	}
+	defer gz.Close()
+
+	filepath := filepath.Join(dbpath, strings.TrimSuffix(filepath.Base(filename), ".gz"))
+
+	out, err := os.Create(filepath)
+	if err != nil {
+		return fmt.Errorf("create file: %w", err)
+	}
+	defer out.Close()
+
+	// Using a maxBytes limit to prevent decompression bombs: gosec G110
+	maxBytes := 200 * 1024 * 1024 // 200MB
+	_, err = io.CopyN(out, gz, int64(maxBytes))
+	if err != nil && err != io.EOF {
+		msg := fmt.Sprintf("error copying file %s: %v", f.Name(), err)
+		panic(msg)
+	}
+
+	return nil
+}

server/vulnerabilities/nvd/cve.go

@@ -69,7 +69,7 @@ func DownloadCVEFeed(vulnPath, cveFeedPrefixURL string, debug bool, logger log.L
 	var err error
 
 	if cveFeedPrefixURL == "" {
-		cveFeedPrefixURL, err = getGitHubCVEAssetPath()
+		cveFeedPrefixURL, err = GetGitHubCVEAssetPath()
 		if err != nil {
 			return fmt.Errorf("get cve asset path: %w", err)
 		}
@@ -83,12 +83,17 @@ func DownloadCVEFeed(vulnPath, cveFeedPrefixURL string, debug bool, logger log.L
 	return nil
 }
 
-func getGitHubCVEAssetPath() (string, error) {
+func GetGitHubCVEAssetPath() (string, error) {
+	vulnOwner := os.Getenv("TEST_VULN_GITHUB_OWNER")
+	if vulnOwner == "" {
+		vulnOwner = owner
+	}
+
 	ghClient := github.NewClient(fleethttp.NewGithubClient())
 
 	releases, _, err := ghClient.Repositories.ListReleases(
 		context.Background(),
-		owner,
+		vulnOwner,
 		vulnRepo,
 		&github.ListOptions{Page: 0, PerPage: 10},
 	)
@@ -115,7 +120,7 @@ func getGitHubCVEAssetPath() (string, error) {
 		return "", errors.New("no CVE feed found")
 	}
 
-	return fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/", owner, vulnRepo, found), nil
+	return fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/", vulnOwner, vulnRepo, found), nil
 }
 
 func downloadNVDCVELegacy(vulnPath string, cveFeedPrefixURL string) error {