mirror of
https://github.com/Vyntral/god-eye.git
synced 2026-02-12 16:52:45 +00:00
🔧 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:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user