Add support to generate auto fixes using LLM (AI)

README.md

@@ -274,6 +274,12 @@ gosec can ignore generated go files with default generated code comment.
 gosec -exclude-generated ./...
 ```
 
+### Autofix vulnerabilities
+gosec can suggest fixes based on AI. To use this feature, please provide the following two parameters `ai-api-provider` and `ai-api-key`. In the current version, gosec only supports gemini. [How to create gemini key.](https://ai.google.dev/gemini-api/docs/api-key)
+
+```bash
+gosec -ai-api-provider="your_provider" -ai-api-key="your_key" ./...
+```
 
 ### Annotating code
 

autofix/ai.go

@@ -0,0 +1,128 @@
+package autofix
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/google/generative-ai-go/genai"
+	"google.golang.org/api/option"
+
+	"github.com/securego/gosec/v2/issue"
+)
+
+const (
+	GeminiModel    = "gemini-1.5-flash"
+	AIPrompt       = "In Go, what is the solution to fix the error %q. Answer limited to 200 words"
+	GeminiProvider = "gemini"
+
+	timeout = 30 * time.Second
+)
+
+// GenAIClient defines the interface for the GenAI client
+type GenAIClient interface {
+	Close() error
+	GenerativeModel(name string) GenAIGenerativeModel
+}
+
+// GenAIGenerativeModel defines the interface for the Generative Model
+type GenAIGenerativeModel interface {
+	GenerateContent(ctx context.Context, prompt string) (string, error)
+}
+
+// genAIClientWrapper wraps the genai.Client to implement GenAIClient
+type genAIClientWrapper struct {
+	client *genai.Client
+}
+
+func (w *genAIClientWrapper) Close() error {
+	return w.client.Close()
+}
+
+func (w *genAIClientWrapper) GenerativeModel(name string) GenAIGenerativeModel {
+	return &genAIGenerativeModelWrapper{model: w.client.GenerativeModel(name)}
+}
+
+// genAIGenerativeModelWrapper wraps the genai.GenerativeModel to implement GenAIGenerativeModel
+type genAIGenerativeModelWrapper struct {
+	// model is the underlying generative model
+	model *genai.GenerativeModel
+}
+
+func (w *genAIGenerativeModelWrapper) GenerateContent(ctx context.Context, prompt string) (string, error) {
+	resp, err := w.model.GenerateContent(ctx, genai.Text(prompt))
+	if err != nil {
+		return "", fmt.Errorf("generate content error: %w", err)
+	}
+	if len(resp.Candidates) == 0 {
+		return "", errors.New("no candidates found")
+	}
+	// Return the first candidate
+	return fmt.Sprintf("%+v", resp.Candidates[0].Content.Parts[0]), nil
+}
+
+func NewGenAIClient(ctx context.Context, aiApiKey, endpoint string) (GenAIClient, error) {
+	clientOptions := []option.ClientOption{option.WithAPIKey(aiApiKey)}
+	if endpoint != "" {
+		clientOptions = append(clientOptions, option.WithEndpoint(endpoint))
+	}
+
+	client, err := genai.NewClient(ctx, clientOptions...)
+	if err != nil {
+		return nil, fmt.Errorf("calling gemini API: %w", err)
+	}
+
+	return &genAIClientWrapper{client: client}, nil
+}
+
+func generateSolutionByGemini(client GenAIClient, issues []*issue.Issue) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	model := client.GenerativeModel(GeminiModel)
+	cachedAutofix := make(map[string]string)
+	for _, issue := range issues {
+		if val, ok := cachedAutofix[issue.What]; ok {
+			issue.Autofix = val
+			continue
+		}
+
+		prompt := fmt.Sprintf(AIPrompt, issue.What)
+		resp, err := model.GenerateContent(ctx, prompt)
+		if err != nil {
+			return fmt.Errorf("gemini generating content: %w", err)
+		}
+
+		if resp == "" {
+			return errors.New("gemini no candidates found")
+		}
+
+		issue.Autofix = resp
+		cachedAutofix[issue.What] = issue.Autofix
+	}
+	return nil
+}
+
+// GenerateSolution generates a solution for the given issues using the specified AI provider
+func GenerateSolution(aiApiProvider, aiApiKey, endpoint string, issues []*issue.Issue) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	var client GenAIClient
+
+	switch aiApiProvider {
+	case GeminiProvider:
+		var err error
+		client, err = NewGenAIClient(ctx, aiApiKey, endpoint)
+		if err != nil {
+			return fmt.Errorf("generate solution error: %w", err)
+		}
+	default:
+		return errors.New("ai provider not supported")
+	}
+
+	defer client.Close()
+
+	return generateSolutionByGemini(client, issues)
+}

autofix/ai_test.go

@@ -0,0 +1,114 @@
+package autofix
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/securego/gosec/v2/issue"
+)
+
+// MockGenAIClient is a mock of the GenAIClient interface
+type MockGenAIClient struct {
+	mock.Mock
+}
+
+func (m *MockGenAIClient) Close() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+func (m *MockGenAIClient) GenerativeModel(name string) GenAIGenerativeModel {
+	args := m.Called(name)
+	return args.Get(0).(GenAIGenerativeModel)
+}
+
+// MockGenAIGenerativeModel is a mock of the GenAIGenerativeModel interface
+type MockGenAIGenerativeModel struct {
+	mock.Mock
+}
+
+func (m *MockGenAIGenerativeModel) GenerateContent(ctx context.Context, prompt string) (string, error) {
+	args := m.Called(ctx, prompt)
+	return args.String(0), args.Error(1)
+}
+
+func TestGenerateSolutionByGemini_Success(t *testing.T) {
+	// Arrange
+	issues := []*issue.Issue{
+		{What: "Example issue 1"},
+	}
+
+	mockClient := new(MockGenAIClient)
+	mockModel := new(MockGenAIGenerativeModel)
+	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel)
+	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("Autofix for issue 1", nil)
+
+	// Act
+	err := generateSolutionByGemini(mockClient, issues)
+
+	// Assert
+	assert.NoError(t, err)
+	assert.Equal(t, "Autofix for issue 1", issues[0].Autofix)
+	mockClient.AssertExpectations(t)
+	mockModel.AssertExpectations(t)
+}
+
+func TestGenerateSolutionByGemini_NoCandidates(t *testing.T) {
+	// Arrange
+	issues := []*issue.Issue{
+		{What: "Example issue 2"},
+	}
+
+	mockClient := new(MockGenAIClient)
+	mockModel := new(MockGenAIGenerativeModel)
+	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel)
+	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("", nil)
+
+	// Act
+	err := generateSolutionByGemini(mockClient, issues)
+
+	// Assert
+	assert.Error(t, err)
+	assert.Equal(t, "gemini no candidates found", err.Error())
+	mockClient.AssertExpectations(t)
+	mockModel.AssertExpectations(t)
+}
+
+func TestGenerateSolutionByGemini_APIError(t *testing.T) {
+	// Arrange
+	issues := []*issue.Issue{
+		{What: "Example issue 3"},
+	}
+
+	mockClient := new(MockGenAIClient)
+	mockModel := new(MockGenAIGenerativeModel)
+	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel)
+	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("", errors.New("API error"))
+
+	// Act
+	err := generateSolutionByGemini(mockClient, issues)
+
+	// Assert
+	assert.Error(t, err)
+	assert.Equal(t, "gemini generating content: API error", err.Error())
+	mockClient.AssertExpectations(t)
+	mockModel.AssertExpectations(t)
+}
+
+func TestGenerateSolution_UnsupportedProvider(t *testing.T) {
+	// Arrange
+	issues := []*issue.Issue{
+		{What: "Example issue 4"},
+	}
+
+	// Act
+	err := GenerateSolution("unsupported-provider", "test-api-key", "", issues)
+
+	// Assert
+	assert.Error(t, err)
+	assert.Equal(t, "ai provider not supported", err.Error())
+}

cmd/gosec/main.go

@@ -25,6 +25,7 @@ import (
 	"strings"
 
 	"github.com/securego/gosec/v2"
+	"github.com/securego/gosec/v2/autofix"
 	"github.com/securego/gosec/v2/cmd/vflag"
 	"github.com/securego/gosec/v2/issue"
 	"github.com/securego/gosec/v2/report"
@@ -149,6 +150,15 @@ var (
 	// flagTerse shows only the summary of scan discarding all the logs
 	flagTerse = flag.Bool("terse", false, "Shows only the results and summary")
 
+	// AI platform provider to generate solutions to issues
+	flagAiApiProvider = flag.String("ai-api-provider", "", "AI platform provider to generate solutions to issues")
+
+	// key to implementing AI provider services
+	flagAiApiKey = flag.String("ai-api-key", "", "key to implementing AI provider services")
+
+	// endpoint to the AI provider
+	flagAiEndpoint = flag.String("ai-endpoint", "", "endpoint to the AI provider")
+
 	// exclude the folders from scan
 	flagDirsExclude arrayFlags
 
@@ -457,6 +467,14 @@ func main() {
 
 	reportInfo := gosec.NewReportInfo(issues, metrics, errors).WithVersion(Version)
 
+	// Call AI request to solve the issues
+	if *flagAiApiProvider != "" && *flagAiApiKey != "" {
+		err := autofix.GenerateSolution(*flagAiApiProvider, *flagAiApiKey, *flagAiEndpoint, issues)
+		if err != nil {
+			logger.Print(err)
+		}
+	}
+
 	if *flagOutput == "" || *flagStdOut {
 		fileFormat := getPrintedFormat(*flagFormat, *flagVerbose)
 		if err := printReport(fileFormat, *flagColor, rootPaths, reportInfo); err != nil {

go.mod

@@ -2,30 +2,63 @@ module github.com/securego/gosec/v2
 
 require (
 	github.com/ccojocar/zxcvbn-go v1.0.2
+	github.com/google/generative-ai-go v0.17.0
 	github.com/google/uuid v1.6.0
 	github.com/gookit/color v1.5.4
 	github.com/lib/pq v1.10.9
 	github.com/mozilla/tls-observatory v0.0.0-20210609171429-7bc42856d2e5
 	github.com/onsi/ginkgo/v2 v2.20.0
 	github.com/onsi/gomega v1.34.1
+	github.com/stretchr/testify v1.9.0
 	golang.org/x/crypto v0.26.0
 	golang.org/x/lint v0.0.0-20210508222113-6edffad5e616
 	golang.org/x/text v0.17.0
 	golang.org/x/tools v0.24.0
+	google.golang.org/api v0.186.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
+	cloud.google.com/go v0.115.0 // indirect
+	cloud.google.com/go/ai v0.8.0 // indirect
+	cloud.google.com/go/auth v0.6.0 // indirect
+	cloud.google.com/go/auth/oauth2adapt v0.2.2 // indirect
+	cloud.google.com/go/compute/metadata v0.3.0 // indirect
+	cloud.google.com/go/longrunning v0.5.7 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/pprof v0.0.0-20240727154555-813a5fbdbec8 // indirect
+	github.com/google/s2a-go v0.1.7 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
+	github.com/googleapis/gax-go/v2 v2.12.5 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 // indirect
+	go.opencensus.io v0.24.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.51.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect
+	go.opentelemetry.io/otel v1.26.0 // indirect
+	go.opentelemetry.io/otel/metric v1.26.0 // indirect
+	go.opentelemetry.io/otel/trace v1.26.0 // indirect
 	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
 	golang.org/x/mod v0.20.0 // indirect
 	golang.org/x/net v0.28.0 // indirect
+	golang.org/x/oauth2 v0.21.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.23.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4 // indirect
+	google.golang.org/grpc v1.64.1 // indirect
+	google.golang.org/protobuf v1.34.2 // indirect
 )
 
-go 1.20
+go 1.21
+
+toolchain go1.22.5