Merge pull request #733 from Checkmarx/avis/azure-support

AST-45352 - Add support for Checkmarx Azure AI

.golangci.yml

@@ -11,7 +11,7 @@ linters-settings:
           - github.com/CheckmarxDev/containers-resolver/pkg/containerResolver
           - github.com/Checkmarx/gen-ai-prompts/prompts/sast_result_remediation
           - github.com/spf13/viper
-          - github.com/checkmarxDev/gpt-wrapper
+          - github.com/Checkmarx/gen-ai-wrapper
           - github.com/spf13/cobra
           - github.com/pkg/errors
           - github.com/google

go.mod

@@ -4,10 +4,10 @@ go 1.22.7
 
 require (
 	github.com/Checkmarx/gen-ai-prompts v0.0.0-20240807143411-708ceec12b63
+	github.com/Checkmarx/gen-ai-wrapper v1.0.2
 	github.com/CheckmarxDev/containers-resolver v1.0.14
 	github.com/MakeNowJust/heredoc v1.0.0
 	github.com/bouk/monkey v1.0.0
-	github.com/checkmarxDev/gpt-wrapper v0.0.0-20230721160222-85da2fd1cc4c
 	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/gomarkdown/markdown v0.0.0-20241102151059-6bc1ffdc6e8c
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
@@ -17,7 +17,7 @@ require (
 	github.com/mssola/user_agent v0.6.0
 	github.com/pkg/errors v0.9.1
 	github.com/spf13/cobra v1.8.1
-	github.com/spf13/viper v1.18.2
+	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
 	github.com/tomnomnom/linkheader v0.0.0-20180905144013-02ca5825eb80
 	github.com/xeipuuv/gojsonschema v1.2.0
@@ -195,7 +195,7 @@ require (
 	github.com/opencontainers/selinux v1.11.0 // indirect
 	github.com/pborman/indent v1.2.1 // indirect
 	github.com/pelletier/go-toml v1.9.5 // indirect
-	github.com/pelletier/go-toml/v2 v2.2.1 // indirect
+	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/peterbourgon/diskv v2.0.1+incompatible // indirect
 	github.com/pierrec/lz4/v4 v4.1.21 // indirect
 	github.com/pjbgf/sha1cd v0.3.0 // indirect

go.sum

@@ -62,6 +62,8 @@ github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/Checkmarx/gen-ai-prompts v0.0.0-20240807143411-708ceec12b63 h1:SCuTcE+CFvgjbIxUNL8rsdB2sAhfuNx85HvxImKta3g=
 github.com/Checkmarx/gen-ai-prompts v0.0.0-20240807143411-708ceec12b63/go.mod h1:MI6lfLerXU+5eTV/EPTDavgnV3owz3GPT4g/msZBWPo=
+github.com/Checkmarx/gen-ai-wrapper v1.0.2 h1:T6X40+4hYnwfDsvkjWs9VIcE6s1O+8DUu0+sDdCY3GI=
+github.com/Checkmarx/gen-ai-wrapper v1.0.2/go.mod h1:xwRLefezwNNnRGu1EjGS6wNiR9FVV/eP9D+oXwLViVM=
 github.com/CheckmarxDev/containers-resolver v1.0.14 h1:jOc1POn6XVETZtXMyer6Rp6K/MLfkuGCKgMk6ZLEf/w=
 github.com/CheckmarxDev/containers-resolver v1.0.14/go.mod h1:ne5YunT/hKQ7fnZejFVGodlfOUReNE7hZW2KLbpQi48=
 github.com/CycloneDX/cyclonedx-go v0.9.0 h1:inaif7qD8bivyxp7XLgxUYtOXWtDez7+j72qKTMQTb8=
@@ -195,8 +197,6 @@ github.com/charmbracelet/x/term v0.1.1 h1:3cosVAiPOig+EV4X9U+3LDgtwwAoEzJjNdwbXD
 github.com/charmbracelet/x/term v0.1.1/go.mod h1:wB1fHt5ECsu3mXYusyzcngVWWlu1KKUmmLhfgr/Flxw=
 github.com/charmbracelet/x/windows v0.1.0 h1:gTaxdvzDM5oMa/I2ZNF7wN78X/atWemG9Wph7Ika2k4=
 github.com/charmbracelet/x/windows v0.1.0/go.mod h1:GLEO/l+lizvFDBPLIOk+49gdX49L9YWMB5t+DZd0jkQ=
-github.com/checkmarxDev/gpt-wrapper v0.0.0-20230721160222-85da2fd1cc4c h1:oKI4C1dXYpi0B8pltDDzp1ZRiyeILv5enbp9h4ASQ3s=
-github.com/checkmarxDev/gpt-wrapper v0.0.0-20230721160222-85da2fd1cc4c/go.mod h1:l+0rISRGaps2HWkpvKbYPE1nsNx28vBj6bKorEm1M5o=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
 github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
@@ -763,8 +763,8 @@ github.com/pborman/indent v1.2.1/go.mod h1:FitS+t35kIYtB5xWTZAPhnmrxcciEEOdbyrrp
 github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
 github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
 github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
-github.com/pelletier/go-toml/v2 v2.2.1 h1:9TA9+T8+8CUCO2+WYnDLCgrYi9+omqKXyjDtosvtEhg=
-github.com/pelletier/go-toml/v2 v2.2.1/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
+github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
+github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
 github.com/peterbourgon/diskv v2.0.1+incompatible h1:UBdAOUP5p4RWqPBg048CAvpKN+vxiaj6gdUUzhl4XmI=
 github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=
 github.com/phayes/freeport v0.0.0-20220201140144-74d24b5ae9f5 h1:Ii+DKncOVM8Cu1Hc+ETb5K+23HdAMvESYE3ZJ5b5cMI=
@@ -885,8 +885,8 @@ github.com/spf13/jwalterweatherman v1.1.0/go.mod h1:aNWZUN0dPAAO/Ljvb5BEdw96iTZ0
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.10.0/go.mod h1:SoyBPwAtKDzypXNDFKN5kzH7ppppbGZtls1UpIy5AsM=
-github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ=
-github.com/spf13/viper v1.18.2/go.mod h1:EKmWIqdnk5lOcmR72yw6hS+8OPYcwD0jteitLMVB+yk=
+github.com/spf13/viper v1.19.0 h1:RWq5SEjt8o25SROyN3z2OrDB9l7RPd3lwTWU8EcEdcI=
+github.com/spf13/viper v1.19.0/go.mod h1:GQUN9bilAbhU/jgc1bKs99f/suXKeUMct8Adx5+Ntkg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=

internal/commands/chat-kics.go

@@ -3,15 +3,15 @@ package commands
 import (
 	"fmt"
 	"os"
+	"strings"
 
+	"github.com/Checkmarx/gen-ai-wrapper/pkg/message"
+	"github.com/Checkmarx/gen-ai-wrapper/pkg/role"
+	gptWrapper "github.com/Checkmarx/gen-ai-wrapper/pkg/wrapper"
 	"github.com/checkmarx/ast-cli/internal/commands/util/printer"
 	"github.com/checkmarx/ast-cli/internal/logger"
 	"github.com/checkmarx/ast-cli/internal/params"
 	"github.com/checkmarx/ast-cli/internal/wrappers"
-	"github.com/checkmarxDev/gpt-wrapper/pkg/connector"
-	"github.com/checkmarxDev/gpt-wrapper/pkg/message"
-	"github.com/checkmarxDev/gpt-wrapper/pkg/role"
-	"github.com/checkmarxDev/gpt-wrapper/pkg/wrapper"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
@@ -41,18 +41,24 @@ const dropLen = 4
 
 const FileErrorFormat = "It seems that %s is not available for AI Guided Remediation. Please ensure that you have opened the correct workspace or the relevant file."
 
+// chatModel model to use when calling the CheckmarxAI
+const checkmarxAiChatModel = "gpt-4"
+const tenantIDClaimKey = "tenant_id"
+const guidedRemediationFeatureNameKics = "cli-guided-remediation-kics"
+const guidedRemediationFeatureNameSast = "cli-guided-remediation-sast"
+
 type OutputModel struct {
 	ConversationID string   `json:"conversationId"`
 	Response       []string `json:"response"`
 }
 
-func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper) *cobra.Command {
+func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper, tenantWrapper wrappers.TenantConfigurationWrapper) *cobra.Command {
 	chatKicsCmd := &cobra.Command{
 		Use:    "kics",
 		Short:  "Chat about KICS result with OpenAI models",
 		Long:   "Chat about KICS result with OpenAI models",
 		Hidden: true,
-		RunE:   runChatKics(chatWrapper),
+		RunE:   runChatKics(chatWrapper, tenantWrapper),
 	}
 
 	chatKicsCmd.Flags().String(params.ChatAPIKey, "", "OpenAI API key")
@@ -65,7 +71,6 @@ func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper) *cobra.Command {
 	chatKicsCmd.Flags().String(params.ChatKicsResultVulnerability, "", "IaC result vulnerability name")
 
 	_ = chatKicsCmd.MarkFlagRequired(params.ChatUserInput)
-	_ = chatKicsCmd.MarkFlagRequired(params.ChatAPIKey)
 	_ = chatKicsCmd.MarkFlagRequired(params.ChatKicsResultFile)
 	_ = chatKicsCmd.MarkFlagRequired(params.ChatKicsResultLine)
 	_ = chatKicsCmd.MarkFlagRequired(params.ChatKicsResultSeverity)
@@ -74,27 +79,29 @@ func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper) *cobra.Command {
 	return chatKicsCmd
 }
 
-func runChatKics(chatKicsWrapper wrappers.ChatWrapper) func(cmd *cobra.Command, args []string) error {
+func runChatKics(
+	chatKicsWrapper wrappers.ChatWrapper, tenantWrapper wrappers.TenantConfigurationWrapper,
+) func(cmd *cobra.Command, args []string) error {
 	return func(cmd *cobra.Command, args []string) error {
-		chatAPIKey, _ := cmd.Flags().GetString(params.ChatAPIKey)
 		chatConversationID, _ := cmd.Flags().GetString(params.ChatConversationID)
-		chatModel, _ := cmd.Flags().GetString(params.ChatModel)
 		chatResultFile, _ := cmd.Flags().GetString(params.ChatKicsResultFile)
 		chatResultLine, _ := cmd.Flags().GetString(params.ChatKicsResultLine)
 		chatResultSeverity, _ := cmd.Flags().GetString(params.ChatKicsResultSeverity)
 		chatResultVulnerability, _ := cmd.Flags().GetString(params.ChatKicsResultVulnerability)
 		userInput, _ := cmd.Flags().GetString(params.ChatUserInput)
 
-		statefulWrapper := wrapper.NewStatefulWrapper(connector.NewFileSystemConnector(""), chatAPIKey, chatModel, dropLen, 0)
-
-		if chatConversationID == "" {
-			chatConversationID = statefulWrapper.GenerateId().String()
+		azureAiEnabled, checkmarxAiEnabled, tenantConfigurationResponses, err := getEngineSelection(cmd, tenantWrapper)
+		if err != nil {
+			return err
 		}
 
-		id, err := uuid.Parse(chatConversationID)
+		statefulWrapper, customerToken := CreateStatefulWrapper(cmd, azureAiEnabled, checkmarxAiEnabled, tenantConfigurationResponses)
+
+		tenantID := getTenantID(customerToken)
+
+		id, err := getKicsConversationID(cmd, chatConversationID, statefulWrapper)
 		if err != nil {
-			logger.PrintIfVerbose(err.Error())
-			return outputError(cmd, id, errors.Errorf(ConversationIDErrorFormat, chatConversationID))
+			return err
 		}
 
 		chatResultCode, err := os.ReadFile(chatResultFile)
@@ -103,22 +110,96 @@ func runChatKics(chatKicsWrapper wrappers.ChatWrapper) func(cmd *cobra.Command,
 			return outputError(cmd, id, errors.Errorf(FileErrorFormat, chatResultFile))
 		}
 
-		newMessages := buildMessages(chatResultCode, chatResultVulnerability, chatResultLine, chatResultSeverity, userInput)
-		response, err := chatKicsWrapper.Call(statefulWrapper, id, newMessages)
+		newMessages := buildKicsMessages(chatResultCode, chatResultVulnerability, chatResultLine, chatResultSeverity, userInput)
+
+		responseContent, err := sendRequest(statefulWrapper, azureAiEnabled, checkmarxAiEnabled, tenantID, chatKicsWrapper, id, newMessages, customerToken, guidedRemediationFeatureNameKics)
 		if err != nil {
 			return outputError(cmd, id, err)
 		}
 
-		responseContent := getMessageContents(response)
-
 		return printer.Print(cmd.OutOrStdout(), &OutputModel{
 			ConversationID: id.String(),
 			Response:       responseContent,
 		}, printer.FormatJSON)
 	}
 }
 
-func buildMessages(chatResultCode []byte,
+func getKicsConversationID(cmd *cobra.Command, chatConversationID string, statefulWrapper gptWrapper.StatefulWrapper) (uuid.UUID, error) {
+	if chatConversationID == "" {
+		chatConversationID = statefulWrapper.GenerateId().String()
+	}
+
+	id, err := uuid.Parse(chatConversationID)
+	if err != nil {
+		logger.PrintIfVerbose(err.Error())
+		return uuid.UUID{}, outputError(cmd, id, errors.Errorf(ConversationIDErrorFormat, chatConversationID))
+	}
+	return id, err
+}
+
+func getTenantID(customerToken string) string {
+	tenantID, _ := wrappers.ExtractFromTokenClaims(customerToken, tenantIDClaimKey)
+	// remove from tenant id all the string before ::
+	if strings.Contains(tenantID, "::") {
+		tenantID = tenantID[strings.LastIndex(tenantID, "::")+2:]
+	}
+	return tenantID
+}
+
+func sendRequest(statefulWrapper gptWrapper.StatefulWrapper, azureAiEnabled bool, checkmarxAiEnabled bool, tenantID string, chatKicsWrapper wrappers.ChatWrapper,
+	id uuid.UUID, newMessages []message.Message, customerToken string, featureName string) (responseContent []string, err error) {
+	requestID := statefulWrapper.GenerateId().String()
+
+	var response []message.Message
+
+	if azureAiEnabled || checkmarxAiEnabled {
+		metadata := message.MetaData{
+			TenantID:  tenantID,
+			RequestID: requestID,
+			UserAgent: params.DefaultAgent,
+			Feature:   featureName,
+		}
+		if azureAiEnabled {
+			logger.Printf("Sending message to Azure AI model for " + featureName + " guided remediation. RequestID: " + requestID)
+		} else {
+			logger.Printf("Sending message to Checkmarx AI model for " + featureName + " guided remediation. RequestID: " + requestID)
+		}
+		response, err = chatKicsWrapper.SecureCall(statefulWrapper, id, newMessages, &metadata, customerToken)
+		if err != nil {
+			return nil, err
+		}
+	} else { // if chatgpt is enabled or no engine is enabled
+		logger.Printf("Sending message to ChatGPT model for " + featureName + " guided remediation. RequestID: " + requestID)
+		response, err = chatKicsWrapper.Call(statefulWrapper, id, newMessages)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	responseContent = getMessageContents(response)
+	return responseContent, nil
+}
+
+func getEngineSelection(cmd *cobra.Command, tenantWrapper wrappers.TenantConfigurationWrapper) (azureAiEnabled, checkmarxAiEnabled bool,
+	tenantConfigurationResponses *[]*wrappers.TenantConfigurationResponse, err error) {
+	if !isCxOneAPIKeyAvailable() {
+		azureAiEnabled = false
+		checkmarxAiEnabled = false
+		logger.Printf("CxOne API key is not available, ChatGPT model will be used for guided remediation.")
+	} else {
+		var err error
+		tenantConfigurationResponses, err = GetTenantConfigurationResponses(tenantWrapper)
+		if err != nil {
+			return false, false, nil, outputError(cmd, uuid.Nil, err)
+		}
+
+		azureAiEnabled = isAzureAiGuidedRemediationEnabled(tenantConfigurationResponses)
+		checkmarxAiEnabled = isCheckmarxAiGuidedRemediationEnabled(tenantConfigurationResponses)
+	}
+	return azureAiEnabled, checkmarxAiEnabled, tenantConfigurationResponses, nil
+}
+
+func buildKicsMessages(chatResultCode []byte,
 	chatResultVulnerability, chatResultLine, chatResultSeverity, userInput string) []message.Message {
 	var newMessages []message.Message
 	newMessages = append(newMessages, message.Message{

internal/commands/chat-kics_test.go

@@ -6,7 +6,11 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/checkmarx/ast-cli/internal/params"
+	"github.com/checkmarx/ast-cli/internal/wrappers"
+	"github.com/checkmarx/ast-cli/internal/wrappers/mock"
 	"github.com/google/uuid"
+	"github.com/spf13/viper"
 	"gotest.tools/assert"
 )
 
@@ -49,7 +53,6 @@ func TestChatKicsInvalidFile(t *testing.T) {
 func TestChatKicsCorrectResponse(t *testing.T) {
 	buffer, err := executeRedirectedTestCommand("chat", "kics",
 		"--conversation-id", uuid.New().String(),
-		"--chat-apikey", "apiKey",
 		"--user-input", "userInput",
 		"--result-file", "./data/Dockerfile",
 		"--result-line", "0",
@@ -61,3 +64,76 @@ func TestChatKicsCorrectResponse(t *testing.T) {
 	s := strings.ToLower(string(output))
 	assert.Assert(t, strings.Contains(s, "mock"), s)
 }
+
+func TestChatKicsAzureAICorrectResponse(t *testing.T) {
+	mock.TenantConfiguration = []*wrappers.TenantConfigurationResponse{
+		{
+			Key:   "scan.config.plugins.ideScans",
+			Value: "true",
+		},
+		{
+			Key:   "scan.config.plugins.azureAiGuidedRemediation",
+			Value: "true",
+		},
+		{
+			Key:   "scan.config.plugins.aiGuidedRemediationAiEngine",
+			Value: "azureai",
+		},
+	}
+	origAPIKey := viper.GetString(params.AstAPIKey)
+	viper.Set(params.AstAPIKey, "SomeKey")
+
+	buffer, err := executeRedirectedTestCommand("chat", "kics",
+		"--conversation-id", uuid.New().String(),
+		"--user-input", "userInput",
+		"--result-file", "./data/Dockerfile",
+		"--result-line", "0",
+		"--result-severity", "LOW",
+		"--result-vulnerability", "Vulnerability")
+	assert.NilError(t, err)
+	output, err := io.ReadAll(buffer)
+	assert.NilError(t, err)
+	s := strings.ToLower(string(output))
+
+	mock.TenantConfiguration = []*wrappers.TenantConfigurationResponse{}
+	viper.Set(params.AstAPIKey, origAPIKey)
+
+	assert.Assert(t, strings.Contains(s, "mock message from securecall"), s)
+}
+
+func TestChatKicsCheckmarxAICorrectResponse(t *testing.T) {
+	mock.TenantConfiguration = []*wrappers.TenantConfigurationResponse{
+		{
+			Key:   "scan.config.plugins.ideScans",
+			Value: "true",
+		},
+		{
+			Key:   "scan.config.plugins.checkmarxAiGuidedRemediation",
+			Value: "true",
+		},
+		{
+			Key:   "scan.config.plugins.aiGuidedRemediationAiEngine",
+			Value: "checkmarxai",
+		},
+	}
+	origAPIKey := viper.GetString(params.AstAPIKey)
+	viper.Set(params.AstAPIKey, "SomeKey")
+
+	buffer, err := executeRedirectedTestCommand("chat", "kics",
+		"--conversation-id", uuid.New().String(),
+		"--chat-apikey", "apiKey",
+		"--user-input", "userInput",
+		"--result-file", "./data/Dockerfile",
+		"--result-line", "0",
+		"--result-severity", "LOW",
+		"--result-vulnerability", "Vulnerability")
+	assert.NilError(t, err)
+	output, err := io.ReadAll(buffer)
+	assert.NilError(t, err)
+	s := strings.ToLower(string(output))
+
+	mock.TenantConfiguration = []*wrappers.TenantConfigurationResponse{}
+	viper.Set(params.AstAPIKey, origAPIKey)
+
+	assert.Assert(t, strings.Contains(s, "mock message from securecall"), s)
+}