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