From cb69ad89b99f200167cd81436acb7504da55423d Mon Sep 17 00:00:00 2001 From: EdgarPsda Date: Sun, 29 Mar 2026 19:22:03 -0700 Subject: [PATCH] Add AI fix suggestions for HIGH/CRITICAL findings (opt-in) - New cli/ai package: Client supporting Ollama, OpenAI, Anthropic providers - Cache suggestions by sha256(ruleID+message) to avoid duplicate API calls - Only enriches HIGH and CRITICAL findings to keep noise low - Add AISuggestion field to Finding struct (shown in terminal + JSON output) - Add ai: section to security-config.yml (disabled by default) - API keys read from config or OPENAI_API_KEY / ANTHROPIC_API_KEY env vars - Terminal reporter shows fix suggestion inline under each finding - Default provider: ollama (local, privacy-first, no data sent externally) - 5 unit tests: defaults, severity filter, caching, prompt content, cache keys --- cli/ai/cache.go | 37 ++++ cli/ai/suggestions.go | 253 +++++++++++++++++++++++++ cli/ai/suggestions_test.go | 133 +++++++++++++ cli/cmd/scan.go | 38 ++++ cli/config/loader.go | 26 ++- cli/reporters/terminal.go | 4 + cli/scanners/types.go | 17 +- cli/templates/security-config.yml.tmpl | 10 + 8 files changed, 502 insertions(+), 16 deletions(-) create mode 100644 cli/ai/cache.go create mode 100644 cli/ai/suggestions.go create mode 100644 cli/ai/suggestions_test.go diff --git a/cli/ai/cache.go b/cli/ai/cache.go new file mode 100644 index 0000000..f761826 --- /dev/null +++ b/cli/ai/cache.go @@ -0,0 +1,37 @@ +package ai + +import ( + "crypto/sha256" + "fmt" + "sync" +) + +// Cache stores AI suggestions keyed by a hash of rule+message. +// Thread-safe for concurrent enrichment. +type Cache struct { + mu sync.RWMutex + store map[string]string +} + +func NewCache() *Cache { + return &Cache{store: make(map[string]string)} +} + +func (c *Cache) Get(key string) (string, bool) { + c.mu.RLock() + defer c.mu.RUnlock() + v, ok := c.store[key] + return v, ok +} + +func (c *Cache) Set(key, value string) { + c.mu.Lock() + defer c.mu.Unlock() + c.store[key] = value +} + +// cacheKey returns a stable key for a rule ID + message pair +func cacheKey(ruleID, message string) string { + h := sha256.Sum256([]byte(ruleID + "\x00" + message)) + return fmt.Sprintf("%x", h[:8]) +} diff --git a/cli/ai/suggestions.go b/cli/ai/suggestions.go new file mode 100644 index 0000000..e669014 --- /dev/null +++ b/cli/ai/suggestions.go @@ -0,0 +1,253 @@ +package ai + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/edgarpsda/devsecops-kit/cli/scanners" +) + +// Config holds AI provider configuration +type Config struct { + Enabled bool + Provider string // "ollama", "openai", "anthropic" + Model string + Endpoint string // for ollama; defaults to http://localhost:11434 + APIKey string // for openai/anthropic; reads from env if empty +} + +// Client generates fix suggestions for security findings +type Client struct { + cfg Config + cache *Cache + http *http.Client +} + +// NewClient creates an AI client with the given config +func NewClient(cfg Config) *Client { + endpoint := cfg.Endpoint + if endpoint == "" { + endpoint = "http://localhost:11434" + } + cfg.Endpoint = endpoint + + model := cfg.Model + if model == "" { + switch cfg.Provider { + case "openai": + model = "gpt-4o-mini" + case "anthropic": + model = "claude-haiku-4-5-20251001" + default: + model = "llama3" + } + } + cfg.Model = model + + return &Client{ + cfg: cfg, + cache: NewCache(), + http: &http.Client{Timeout: 30 * time.Second}, + } +} + +// EnrichFindings adds AI fix suggestions to findings in-place. +// Only HIGH and CRITICAL findings are enriched to keep noise low. +// Results are cached so identical rule+message pairs are only sent once. +func (c *Client) EnrichFindings(findings []scanners.Finding) { + for i := range findings { + f := &findings[i] + if f.Severity != "CRITICAL" && f.Severity != "HIGH" { + continue + } + + cacheKey := cacheKey(f.RuleID, f.Message) + if suggestion, ok := c.cache.Get(cacheKey); ok { + f.AISuggestion = suggestion + continue + } + + suggestion, err := c.getSuggestion(f) + if err != nil { + // Non-fatal: just skip this finding + continue + } + + c.cache.Set(cacheKey, suggestion) + f.AISuggestion = suggestion + } +} + +func (c *Client) getSuggestion(f *scanners.Finding) (string, error) { + prompt := buildPrompt(f) + + switch c.cfg.Provider { + case "openai": + return c.callOpenAI(prompt) + case "anthropic": + return c.callAnthropic(prompt) + default: + return c.callOllama(prompt) + } +} + +// buildPrompt constructs a concise, focused prompt for the finding +func buildPrompt(f *scanners.Finding) string { + return fmt.Sprintf( + "You are a security expert. Provide a concise fix suggestion (2-4 sentences max) for this security finding:\n\nTool: %s\nSeverity: %s\nRule: %s\nFile: %s\nIssue: %s\n\nRespond with only the fix suggestion, no preamble.", + f.Tool, f.Severity, f.RuleID, f.File, f.Message, + ) +} + +// --- Ollama --- + +type ollamaRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + Stream bool `json:"stream"` +} + +type ollamaResponse struct { + Response string `json:"response"` + Error string `json:"error,omitempty"` +} + +func (c *Client) callOllama(prompt string) (string, error) { + body, _ := json.Marshal(ollamaRequest{ + Model: c.cfg.Model, + Prompt: prompt, + Stream: false, + }) + + resp, err := c.http.Post(c.cfg.Endpoint+"/api/generate", "application/json", bytes.NewReader(body)) + if err != nil { + return "", fmt.Errorf("ollama request failed: %w", err) + } + defer resp.Body.Close() + + var result ollamaResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "", fmt.Errorf("ollama response decode failed: %w", err) + } + if result.Error != "" { + return "", fmt.Errorf("ollama error: %s", result.Error) + } + + return strings.TrimSpace(result.Response), nil +} + +// --- OpenAI --- + +type openAIRequest struct { + Model string `json:"model"` + Messages []openAIMessage `json:"messages"` +} + +type openAIMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type openAIResponse struct { + Choices []struct { + Message openAIMessage `json:"message"` + } `json:"choices"` + Error *struct { + Message string `json:"message"` + } `json:"error,omitempty"` +} + +func (c *Client) callOpenAI(prompt string) (string, error) { + body, _ := json.Marshal(openAIRequest{ + Model: c.cfg.Model, + Messages: []openAIMessage{ + {Role: "user", Content: prompt}, + }, + }) + + req, _ := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.cfg.APIKey) + + resp, err := c.http.Do(req) + if err != nil { + return "", fmt.Errorf("openai request failed: %w", err) + } + defer resp.Body.Close() + + rawBody, _ := io.ReadAll(resp.Body) + var result openAIResponse + if err := json.Unmarshal(rawBody, &result); err != nil { + return "", fmt.Errorf("openai response decode failed: %w", err) + } + if result.Error != nil { + return "", fmt.Errorf("openai error: %s", result.Error.Message) + } + if len(result.Choices) == 0 { + return "", fmt.Errorf("openai returned no choices") + } + + return strings.TrimSpace(result.Choices[0].Message.Content), nil +} + +// --- Anthropic --- + +type anthropicRequest struct { + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` + Messages []anthropicMessage `json:"messages"` +} + +type anthropicMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type anthropicResponse struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + Error *struct { + Message string `json:"message"` + } `json:"error,omitempty"` +} + +func (c *Client) callAnthropic(prompt string) (string, error) { + body, _ := json.Marshal(anthropicRequest{ + Model: c.cfg.Model, + MaxTokens: 256, + Messages: []anthropicMessage{ + {Role: "user", Content: prompt}, + }, + }) + + req, _ := http.NewRequest("POST", "https://api.anthropic.com/v1/messages", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-api-key", c.cfg.APIKey) + req.Header.Set("anthropic-version", "2023-06-01") + + resp, err := c.http.Do(req) + if err != nil { + return "", fmt.Errorf("anthropic request failed: %w", err) + } + defer resp.Body.Close() + + rawBody, _ := io.ReadAll(resp.Body) + var result anthropicResponse + if err := json.Unmarshal(rawBody, &result); err != nil { + return "", fmt.Errorf("anthropic response decode failed: %w", err) + } + if result.Error != nil { + return "", fmt.Errorf("anthropic error: %s", result.Error.Message) + } + if len(result.Content) == 0 { + return "", fmt.Errorf("anthropic returned empty content") + } + + return strings.TrimSpace(result.Content[0].Text), nil +} diff --git a/cli/ai/suggestions_test.go b/cli/ai/suggestions_test.go new file mode 100644 index 0000000..0a6d5d5 --- /dev/null +++ b/cli/ai/suggestions_test.go @@ -0,0 +1,133 @@ +package ai + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/edgarpsda/devsecops-kit/cli/scanners" +) + +func TestNewClientDefaults(t *testing.T) { + c := NewClient(Config{Provider: "ollama"}) + if c.cfg.Model != "llama3" { + t.Errorf("expected default ollama model 'llama3', got %q", c.cfg.Model) + } + if c.cfg.Endpoint != "http://localhost:11434" { + t.Errorf("expected default endpoint, got %q", c.cfg.Endpoint) + } + + c2 := NewClient(Config{Provider: "openai"}) + if c2.cfg.Model != "gpt-4o-mini" { + t.Errorf("expected default openai model 'gpt-4o-mini', got %q", c2.cfg.Model) + } + + c3 := NewClient(Config{Provider: "anthropic"}) + if c3.cfg.Model != "claude-haiku-4-5-20251001" { + t.Errorf("expected default anthropic model, got %q", c3.cfg.Model) + } +} + +func TestEnrichFindingsOnlyHighCritical(t *testing.T) { + // Mock server that returns a suggestion + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]string{"response": "Use parameterized queries."}) + })) + defer srv.Close() + + c := NewClient(Config{Provider: "ollama", Endpoint: srv.URL}) + + findings := []scanners.Finding{ + {RuleID: "r1", Severity: "CRITICAL", Message: "SQL injection", Tool: "semgrep", File: "a.go"}, + {RuleID: "r2", Severity: "HIGH", Message: "Hardcoded secret", Tool: "gitleaks", File: "b.go"}, + {RuleID: "r3", Severity: "MEDIUM", Message: "Missing header", Tool: "semgrep", File: "c.go"}, + {RuleID: "r4", Severity: "LOW", Message: "Info leak", Tool: "semgrep", File: "d.go"}, + } + + c.EnrichFindings(findings) + + if findings[0].AISuggestion == "" { + t.Error("CRITICAL finding should have a suggestion") + } + if findings[1].AISuggestion == "" { + t.Error("HIGH finding should have a suggestion") + } + if findings[2].AISuggestion != "" { + t.Error("MEDIUM finding should NOT have a suggestion") + } + if findings[3].AISuggestion != "" { + t.Error("LOW finding should NOT have a suggestion") + } +} + +func TestEnrichFindingsCachesResults(t *testing.T) { + callCount := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + json.NewEncoder(w).Encode(map[string]string{"response": "Fix suggestion."}) + })) + defer srv.Close() + + c := NewClient(Config{Provider: "ollama", Endpoint: srv.URL}) + + // Two findings with identical rule+message — should only call API once + findings := []scanners.Finding{ + {RuleID: "r1", Severity: "HIGH", Message: "SQL injection", Tool: "semgrep", File: "a.go"}, + {RuleID: "r1", Severity: "HIGH", Message: "SQL injection", Tool: "semgrep", File: "b.go"}, + } + + c.EnrichFindings(findings) + + if callCount != 1 { + t.Errorf("expected 1 API call due to caching, got %d", callCount) + } + if findings[0].AISuggestion == "" || findings[1].AISuggestion == "" { + t.Error("both findings should have suggestions") + } +} + +func TestBuildPromptContainsKeyFields(t *testing.T) { + f := &scanners.Finding{ + Tool: "trivy", + Severity: "HIGH", + RuleID: "CVE-2024-1234", + File: "go.sum", + Message: "Vulnerable dependency", + } + + prompt := buildPrompt(f) + + for _, expected := range []string{"trivy", "HIGH", "CVE-2024-1234", "go.sum", "Vulnerable dependency"} { + if !contains(prompt, expected) { + t.Errorf("prompt missing %q", expected) + } + } +} + +func TestCacheKeyStable(t *testing.T) { + k1 := cacheKey("rule1", "message1") + k2 := cacheKey("rule1", "message1") + k3 := cacheKey("rule1", "message2") + + if k1 != k2 { + t.Error("same inputs should produce same cache key") + } + if k1 == k3 { + t.Error("different inputs should produce different cache keys") + } +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsStr(s, substr)) +} + +func containsStr(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/cli/cmd/scan.go b/cli/cmd/scan.go index dc6d3b0..2e589eb 100644 --- a/cli/cmd/scan.go +++ b/cli/cmd/scan.go @@ -9,6 +9,7 @@ import ( "runtime" "github.com/spf13/cobra" + "github.com/edgarpsda/devsecops-kit/cli/ai" "github.com/edgarpsda/devsecops-kit/cli/config" "github.com/edgarpsda/devsecops-kit/cli/detectors" "github.com/edgarpsda/devsecops-kit/cli/reporters" @@ -98,6 +99,43 @@ func runScan() error { return fmt.Errorf("scan failed: %w", err) } + // Enrich findings with AI suggestions if enabled + if secConfig.AI.Enabled { + apiKey := secConfig.AI.APIKey + if apiKey == "" { + switch secConfig.AI.Provider { + case "openai": + apiKey = os.Getenv("OPENAI_API_KEY") + case "anthropic": + apiKey = os.Getenv("ANTHROPIC_API_KEY") + } + } + aiClient := ai.NewClient(ai.Config{ + Enabled: true, + Provider: secConfig.AI.Provider, + Model: secConfig.AI.Model, + Endpoint: secConfig.AI.Endpoint, + APIKey: apiKey, + }) + fmt.Println("🤖 Generating AI fix suggestions for HIGH/CRITICAL findings...") + aiClient.EnrichFindings(report.AllFindings) + // Sync enriched AllFindings back into per-tool Results + for i := range report.AllFindings { + f := &report.AllFindings[i] + if f.AISuggestion == "" { + continue + } + if result, ok := report.Results[f.Tool]; ok { + for j := range result.Findings { + if result.Findings[j].RuleID == f.RuleID && result.Findings[j].File == f.File { + result.Findings[j].AISuggestion = f.AISuggestion + break + } + } + } + } + } + // Output results switch scanOutputFormat { case "json": diff --git a/cli/config/loader.go b/cli/config/loader.go index 9e0c5c1..a0dab93 100644 --- a/cli/config/loader.go +++ b/cli/config/loader.go @@ -9,15 +9,25 @@ import ( // SecurityConfig represents the security-config.yml structure type SecurityConfig struct { - Version string `yaml:"version"` - Language string `yaml:"language"` - Framework string `yaml:"framework"` - SeverityThreshold string `yaml:"severity_threshold"` - Tools ToolsConfig `yaml:"tools"` - ExcludePaths []string `yaml:"exclude_paths"` - FailOn map[string]int `yaml:"fail_on"` - Licenses LicensesConfig `yaml:"licenses"` + Version string `yaml:"version"` + Language string `yaml:"language"` + Framework string `yaml:"framework"` + SeverityThreshold string `yaml:"severity_threshold"` + Tools ToolsConfig `yaml:"tools"` + ExcludePaths []string `yaml:"exclude_paths"` + FailOn map[string]int `yaml:"fail_on"` + Licenses LicensesConfig `yaml:"licenses"` Notifications NotificationsConfig `yaml:"notifications"` + AI AIConfig `yaml:"ai"` +} + +// AIConfig holds AI fix suggestion settings +type AIConfig struct { + Enabled bool `yaml:"enabled"` + Provider string `yaml:"provider"` // "ollama", "openai", "anthropic" + Model string `yaml:"model"` + Endpoint string `yaml:"endpoint"` // custom endpoint for ollama + APIKey string `yaml:"api_key"` // for openai/anthropic (prefer env vars) } // ToolsConfig represents the tools section diff --git a/cli/reporters/terminal.go b/cli/reporters/terminal.go index 7be0b14..986f15a 100644 --- a/cli/reporters/terminal.go +++ b/cli/reporters/terminal.go @@ -187,6 +187,10 @@ func (tr *TerminalReporter) printFinding(finding scanners.Finding) { fmt.Printf(" Rule: %s\n", colorGray(finding.RuleID)) } + if finding.AISuggestion != "" { + fmt.Printf(" %s %s\n", colorCyan("💡 Fix:"), colorGray(finding.AISuggestion)) + } + fmt.Println() } diff --git a/cli/scanners/types.go b/cli/scanners/types.go index dd4b207..14c9c53 100644 --- a/cli/scanners/types.go +++ b/cli/scanners/types.go @@ -11,14 +11,15 @@ type ScanResult struct { // Finding represents a single security finding type Finding struct { - File string `json:"file"` - Line int `json:"line"` - Column int `json:"column,omitempty"` - Severity string `json:"severity"` // "CRITICAL", "HIGH", "MEDIUM", "LOW", or rule ID - Message string `json:"message"` - RuleID string `json:"rule_id,omitempty"` - Tool string `json:"tool"` - RemoteURL string `json:"remote_url,omitempty"` // Link to rule documentation + File string `json:"file"` + Line int `json:"line"` + Column int `json:"column,omitempty"` + Severity string `json:"severity"` // "CRITICAL", "HIGH", "MEDIUM", "LOW", or rule ID + Message string `json:"message"` + RuleID string `json:"rule_id,omitempty"` + Tool string `json:"tool"` + RemoteURL string `json:"remote_url,omitempty"` // Link to rule documentation + AISuggestion string `json:"ai_suggestion,omitempty"` // AI-generated fix suggestion } // FindingSummary contains aggregated counts diff --git a/cli/templates/security-config.yml.tmpl b/cli/templates/security-config.yml.tmpl index 486118a..7a92477 100644 --- a/cli/templates/security-config.yml.tmpl +++ b/cli/templates/security-config.yml.tmpl @@ -35,3 +35,13 @@ notifications: pr_comment: true slack: false email: false + +# AI fix suggestions (opt-in) +# Providers: ollama (local, default), openai, anthropic +# For openai/anthropic set api_key here or via OPENAI_API_KEY / ANTHROPIC_API_KEY env vars +# ai: +# enabled: false +# provider: "ollama" # ollama | openai | anthropic +# model: "llama3" # ollama model, or gpt-4o-mini, claude-haiku-4-5-20251001 +# endpoint: "http://localhost:11434" # ollama endpoint (optional) +# api_key: "" # leave empty to use env var