@@ -3,15 +3,15 @@ package commands
33import (
44 "fmt"
55 "os"
6+ "strings"
67
8+ "github.com/Checkmarx/gen-ai-wrapper/pkg/message"
9+ "github.com/Checkmarx/gen-ai-wrapper/pkg/role"
10+ gptWrapper "github.com/Checkmarx/gen-ai-wrapper/pkg/wrapper"
711 "github.com/checkmarx/ast-cli/internal/commands/util/printer"
812 "github.com/checkmarx/ast-cli/internal/logger"
913 "github.com/checkmarx/ast-cli/internal/params"
1014 "github.com/checkmarx/ast-cli/internal/wrappers"
11- "github.com/checkmarxDev/gpt-wrapper/pkg/connector"
12- "github.com/checkmarxDev/gpt-wrapper/pkg/message"
13- "github.com/checkmarxDev/gpt-wrapper/pkg/role"
14- "github.com/checkmarxDev/gpt-wrapper/pkg/wrapper"
1515 "github.com/google/uuid"
1616 "github.com/pkg/errors"
1717 "github.com/spf13/cobra"
@@ -41,18 +41,24 @@ const dropLen = 4
4141
4242const 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."
4343
44+ // chatModel model to use when calling the CheckmarxAI
45+ const checkmarxAiChatModel = "gpt-4"
46+ const tenantIDClaimKey = "tenant_id"
47+ const guidedRemediationFeatureNameKics = "cli-guided-remediation-kics"
48+ const guidedRemediationFeatureNameSast = "cli-guided-remediation-sast"
49+
4450type OutputModel struct {
4551 ConversationID string `json:"conversationId"`
4652 Response []string `json:"response"`
4753}
4854
49- func ChatKicsSubCommand (chatWrapper wrappers.ChatWrapper ) * cobra.Command {
55+ func ChatKicsSubCommand (chatWrapper wrappers.ChatWrapper , tenantWrapper wrappers. TenantConfigurationWrapper ) * cobra.Command {
5056 chatKicsCmd := & cobra.Command {
5157 Use : "kics" ,
5258 Short : "Chat about KICS result with OpenAI models" ,
5359 Long : "Chat about KICS result with OpenAI models" ,
5460 Hidden : true ,
55- RunE : runChatKics (chatWrapper ),
61+ RunE : runChatKics (chatWrapper , tenantWrapper ),
5662 }
5763
5864 chatKicsCmd .Flags ().String (params .ChatAPIKey , "" , "OpenAI API key" )
@@ -65,7 +71,6 @@ func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper) *cobra.Command {
6571 chatKicsCmd .Flags ().String (params .ChatKicsResultVulnerability , "" , "IaC result vulnerability name" )
6672
6773 _ = chatKicsCmd .MarkFlagRequired (params .ChatUserInput )
68- _ = chatKicsCmd .MarkFlagRequired (params .ChatAPIKey )
6974 _ = chatKicsCmd .MarkFlagRequired (params .ChatKicsResultFile )
7075 _ = chatKicsCmd .MarkFlagRequired (params .ChatKicsResultLine )
7176 _ = chatKicsCmd .MarkFlagRequired (params .ChatKicsResultSeverity )
@@ -74,27 +79,29 @@ func ChatKicsSubCommand(chatWrapper wrappers.ChatWrapper) *cobra.Command {
7479 return chatKicsCmd
7580}
7681
77- func runChatKics (chatKicsWrapper wrappers.ChatWrapper ) func (cmd * cobra.Command , args []string ) error {
82+ func runChatKics (
83+ chatKicsWrapper wrappers.ChatWrapper , tenantWrapper wrappers.TenantConfigurationWrapper ,
84+ ) func (cmd * cobra.Command , args []string ) error {
7885 return func (cmd * cobra.Command , args []string ) error {
79- chatAPIKey , _ := cmd .Flags ().GetString (params .ChatAPIKey )
8086 chatConversationID , _ := cmd .Flags ().GetString (params .ChatConversationID )
81- chatModel , _ := cmd .Flags ().GetString (params .ChatModel )
8287 chatResultFile , _ := cmd .Flags ().GetString (params .ChatKicsResultFile )
8388 chatResultLine , _ := cmd .Flags ().GetString (params .ChatKicsResultLine )
8489 chatResultSeverity , _ := cmd .Flags ().GetString (params .ChatKicsResultSeverity )
8590 chatResultVulnerability , _ := cmd .Flags ().GetString (params .ChatKicsResultVulnerability )
8691 userInput , _ := cmd .Flags ().GetString (params .ChatUserInput )
8792
88- statefulWrapper := wrapper .NewStatefulWrapper (connector .NewFileSystemConnector ("" ), chatAPIKey , chatModel , dropLen , 0 )
89-
90- if chatConversationID == "" {
91- chatConversationID = statefulWrapper .GenerateId ().String ()
93+ azureAiEnabled , checkmarxAiEnabled , tenantConfigurationResponses , err := getEngineSelection (cmd , tenantWrapper )
94+ if err != nil {
95+ return err
9296 }
9397
94- id , err := uuid .Parse (chatConversationID )
98+ statefulWrapper , customerToken := CreateStatefulWrapper (cmd , azureAiEnabled , checkmarxAiEnabled , tenantConfigurationResponses )
99+
100+ tenantID := getTenantID (customerToken )
101+
102+ id , err := getKicsConversationID (cmd , chatConversationID , statefulWrapper )
95103 if err != nil {
96- logger .PrintIfVerbose (err .Error ())
97- return outputError (cmd , id , errors .Errorf (ConversationIDErrorFormat , chatConversationID ))
104+ return err
98105 }
99106
100107 chatResultCode , err := os .ReadFile (chatResultFile )
@@ -103,22 +110,96 @@ func runChatKics(chatKicsWrapper wrappers.ChatWrapper) func(cmd *cobra.Command,
103110 return outputError (cmd , id , errors .Errorf (FileErrorFormat , chatResultFile ))
104111 }
105112
106- newMessages := buildMessages (chatResultCode , chatResultVulnerability , chatResultLine , chatResultSeverity , userInput )
107- response , err := chatKicsWrapper .Call (statefulWrapper , id , newMessages )
113+ newMessages := buildKicsMessages (chatResultCode , chatResultVulnerability , chatResultLine , chatResultSeverity , userInput )
114+
115+ responseContent , err := sendRequest (statefulWrapper , azureAiEnabled , checkmarxAiEnabled , tenantID , chatKicsWrapper , id , newMessages , customerToken , guidedRemediationFeatureNameKics )
108116 if err != nil {
109117 return outputError (cmd , id , err )
110118 }
111119
112- responseContent := getMessageContents (response )
113-
114120 return printer .Print (cmd .OutOrStdout (), & OutputModel {
115121 ConversationID : id .String (),
116122 Response : responseContent ,
117123 }, printer .FormatJSON )
118124 }
119125}
120126
121- func buildMessages (chatResultCode []byte ,
127+ func getKicsConversationID (cmd * cobra.Command , chatConversationID string , statefulWrapper gptWrapper.StatefulWrapper ) (uuid.UUID , error ) {
128+ if chatConversationID == "" {
129+ chatConversationID = statefulWrapper .GenerateId ().String ()
130+ }
131+
132+ id , err := uuid .Parse (chatConversationID )
133+ if err != nil {
134+ logger .PrintIfVerbose (err .Error ())
135+ return uuid.UUID {}, outputError (cmd , id , errors .Errorf (ConversationIDErrorFormat , chatConversationID ))
136+ }
137+ return id , err
138+ }
139+
140+ func getTenantID (customerToken string ) string {
141+ tenantID , _ := wrappers .ExtractFromTokenClaims (customerToken , tenantIDClaimKey )
142+ // remove from tenant id all the string before ::
143+ if strings .Contains (tenantID , "::" ) {
144+ tenantID = tenantID [strings .LastIndex (tenantID , "::" )+ 2 :]
145+ }
146+ return tenantID
147+ }
148+
149+ func sendRequest (statefulWrapper gptWrapper.StatefulWrapper , azureAiEnabled bool , checkmarxAiEnabled bool , tenantID string , chatKicsWrapper wrappers.ChatWrapper ,
150+ id uuid.UUID , newMessages []message.Message , customerToken string , featureName string ) (responseContent []string , err error ) {
151+ requestID := statefulWrapper .GenerateId ().String ()
152+
153+ var response []message.Message
154+
155+ if azureAiEnabled || checkmarxAiEnabled {
156+ metadata := message.MetaData {
157+ TenantID : tenantID ,
158+ RequestID : requestID ,
159+ UserAgent : params .DefaultAgent ,
160+ Feature : featureName ,
161+ }
162+ if azureAiEnabled {
163+ logger .Printf ("Sending message to Azure AI model for " + featureName + " guided remediation. RequestID: " + requestID )
164+ } else {
165+ logger .Printf ("Sending message to Checkmarx AI model for " + featureName + " guided remediation. RequestID: " + requestID )
166+ }
167+ response , err = chatKicsWrapper .SecureCall (statefulWrapper , id , newMessages , & metadata , customerToken )
168+ if err != nil {
169+ return nil , err
170+ }
171+ } else { // if chatgpt is enabled or no engine is enabled
172+ logger .Printf ("Sending message to ChatGPT model for " + featureName + " guided remediation. RequestID: " + requestID )
173+ response , err = chatKicsWrapper .Call (statefulWrapper , id , newMessages )
174+ if err != nil {
175+ return nil , err
176+ }
177+ }
178+
179+ responseContent = getMessageContents (response )
180+ return responseContent , nil
181+ }
182+
183+ func getEngineSelection (cmd * cobra.Command , tenantWrapper wrappers.TenantConfigurationWrapper ) (azureAiEnabled , checkmarxAiEnabled bool ,
184+ tenantConfigurationResponses * []* wrappers.TenantConfigurationResponse , err error ) {
185+ if ! isCxOneAPIKeyAvailable () {
186+ azureAiEnabled = false
187+ checkmarxAiEnabled = false
188+ logger .Printf ("CxOne API key is not available, ChatGPT model will be used for guided remediation." )
189+ } else {
190+ var err error
191+ tenantConfigurationResponses , err = GetTenantConfigurationResponses (tenantWrapper )
192+ if err != nil {
193+ return false , false , nil , outputError (cmd , uuid .Nil , err )
194+ }
195+
196+ azureAiEnabled = isAzureAiGuidedRemediationEnabled (tenantConfigurationResponses )
197+ checkmarxAiEnabled = isCheckmarxAiGuidedRemediationEnabled (tenantConfigurationResponses )
198+ }
199+ return azureAiEnabled , checkmarxAiEnabled , tenantConfigurationResponses , nil
200+ }
201+
202+ func buildKicsMessages (chatResultCode []byte ,
122203 chatResultVulnerability , chatResultLine , chatResultSeverity , userInput string ) []message.Message {
123204 var newMessages []message.Message
124205 newMessages = append (newMessages , message.Message {
0 commit comments