🔧 Fix CVE detection with direct NVD API integration

- 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.
This commit is contained in:
Vyntral
2025-11-20 13:51:32 +01:00
parent 123d6123c4
commit 163b51f141
2 changed files with 39 additions and 17 deletions

View File

@@ -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)

View File

@@ -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)
}