From 163b51f141cb5f426b285cb82eef88be143307bc Mon Sep 17 00:00:00 2001
From: Vyntral <luca.lorenzi.1996@gmail.com>
Date: Thu, 20 Nov 2025 13:51:32 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20Fix=20CVE=20detection=20with=20d?=
 =?UTF-8?q?irect=20NVD=20API=20integration?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Replaced AI function calling with direct SearchCVE() calls for reliability
- Added rate limiting (6s between requests) to prevent NVD API 429 errors
- CVE detection now shows real vulnerabilities with CVE IDs and CVSS scores
- Improved prompt context for AI analysis of CVE results

The function calling approach wasn't working reliably with Ollama models.
Now CVE data is fetched directly from NVD and passed to AI for analysis.
---
 internal/ai/cve.go    | 13 +++++++++++++
 internal/ai/ollama.go | 43 ++++++++++++++++++++++++++-----------------
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/internal/ai/cve.go b/internal/ai/cve.go
index 766c48e..b414b17 100644
--- a/internal/ai/cve.go
+++ b/internal/ai/cve.go
@@ -59,6 +59,10 @@ var (
 		Timeout: 10 * time.Second,
 	}
 	nvdBaseURL = "https://services.nvd.nist.gov/rest/json/cves/2.0"
+
+	// Rate limiting: max 5 requests per 30 seconds (NVD allows 10 req/60s without API key)
+	lastNVDRequest time.Time
+	nvdRateLimit   = 6 * time.Second // Wait 6 seconds between requests
 )
 
 // SearchCVE searches for CVE vulnerabilities using NVD API
@@ -121,6 +125,15 @@ func SearchCVE(technology string, version string) (string, error) {
 
 // queryNVD queries the NVD API for CVE information
 func queryNVD(keyword string) ([]CVEInfo, error) {
+	// Rate limiting: wait if necessary
+	if !lastNVDRequest.IsZero() {
+		elapsed := time.Since(lastNVDRequest)
+		if elapsed < nvdRateLimit {
+			time.Sleep(nvdRateLimit - elapsed)
+		}
+	}
+	lastNVDRequest = time.Now()
+
 	// Build URL with query parameters
 	params := url.Values{}
 	params.Add("keywordSearch", keyword)
diff --git a/internal/ai/ollama.go b/internal/ai/ollama.go
index 4cc8aec..a00fdc4 100644
--- a/internal/ai/ollama.go
+++ b/internal/ai/ollama.go
@@ -254,27 +254,34 @@ Be concise and professional.`,
 	return response, nil
 }
 
-// CVEMatch checks for known vulnerabilities in detected technologies using function calling
+// CVEMatch checks for known vulnerabilities in detected technologies
 func (c *OllamaClient) CVEMatch(technology, version string) (string, error) {
-	prompt := fmt.Sprintf(`Check if %s version %s has known CVE vulnerabilities. Use the search_cve tool to look up real CVE data from the NVD database.
-
-After getting CVE results, analyze them and provide:
-1. Summary of findings
-2. Severity assessment
-3. Specific recommendations
-
-If version is unknown, still search using just the technology name.`, technology, version)
-
-	// Use function calling with tools
-	response, err := c.queryWithTools(c.DeepModel, prompt, 30*time.Second)
+	// Call SearchCVE directly instead of using function calling (more reliable)
+	cveData, err := SearchCVE(technology, version)
 	if err != nil {
 		return "", err
 	}
 
-	if strings.Contains(strings.ToLower(response), "no known cve") {
+	// If no CVEs found, return empty
+	if strings.Contains(cveData, "No known CVE vulnerabilities found") {
 		return "", nil
 	}
 
+	// Ask AI to analyze the CVE data
+	prompt := fmt.Sprintf(`Analyze these CVE vulnerabilities for %s and provide a concise security assessment:
+
+%s
+
+Provide a 2-3 sentence summary focusing on:
+- Most critical vulnerabilities
+- Key recommendations`, technology, cveData)
+
+	response, err := c.query(c.FastModel, prompt, 20*time.Second)
+	if err != nil {
+		// If AI analysis fails, return the raw CVE data
+		return cveData, nil
+	}
+
 	return response, nil
 }
 
@@ -424,13 +431,15 @@ func (c *OllamaClient) queryWithTools(model, prompt string, timeout time.Duratio
 			}
 		}
 
-		// Send tool results back to AI for final analysis
-		followUpPrompt := fmt.Sprintf(`%s
+		// Build a clear follow-up prompt with context
+		followUpPrompt := fmt.Sprintf(`You previously asked to use tools to answer this question:
+"%s"
+
+Here are the results from the tools you requested:
 
-Tool Results:
 %s
 
-Based on these results, provide your analysis.`, prompt, formatToolResults(toolResults))
+Now, based on these SPECIFIC results above, provide a detailed security analysis. Use the actual CVE data provided, including CVE IDs, severity scores, and descriptions. Do NOT say you don't have information - the data is RIGHT ABOVE.`, prompt, formatToolResults(toolResults))
 
 		return c.query(model, followUpPrompt, timeout)
 	}