Option for user to automatically add Amazon Bedrock integrations when…

… read only instance role (IAM) is set #1759

deepfence_server/apiDocs/operation.go

@@ -726,9 +726,9 @@ func (d *OpenAPIDocs) AddIntegrationOperations() {
 	d.AddOperation("addGenerativeAiIntegrationBedrock", http.MethodPost, "/deepfence/generative-ai-integration/bedrock",
 		"Add AWS Bedrock Generative AI Integration", "Add a new AWS Bedrock Generative AI Integration",
 		http.StatusOK, []string{tagGenerativeAi}, bearerToken, new(AddGenerativeAiBedrockIntegration), new(MessageResponse))
-	d.AddOperation("autoAddGenerativeAiIntegrationBedrock", http.MethodPost, "/deepfence/generative-ai-integration/bedrock/auto-add",
-		"Automatically add AWS Bedrock Generative AI Integration", "Automatically add AWS Bedrock Generative AI Integrations using IAM role",
-		http.StatusOK, []string{tagGenerativeAi}, bearerToken, new(AutoAddGenerativeAiBedrockIntegration), new(MessageResponse))
+	d.AddOperation("autoAddGenerativeAiIntegration", http.MethodPost, "/deepfence/generative-ai-integration/auto-add",
+		"Automatically add Generative AI Integration", "Automatically add Generative AI Integrations using IAM role",
+		http.StatusAccepted, []string{tagGenerativeAi}, bearerToken, nil, nil)
 
 	d.AddOperation("listGenerativeAiIntegration", http.MethodGet, "/deepfence/generative-ai-integration",
 		"List Generative AI Integrations", "List all the added Generative AI Integrations",

deepfence_server/handler/generative_ai_integration.go

@@ -16,6 +16,7 @@ import (
 	"github.com/deepfence/ThreatMapper/deepfence_utils/encryption"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
 	postgresqlDb "github.com/deepfence/ThreatMapper/deepfence_utils/postgresql/postgresql-db"
+	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
 	"github.com/go-chi/chi/v5"
 	httpext "github.com/go-playground/pkg/v5/net/http"
 )
@@ -25,97 +26,67 @@ var (
 	ErrGenerativeAIIntegrationExists = BadDecoding{
 		err: errors.New("similar integration already exists"),
 	}
-	ErrBedrockNoActiveModel = BadDecoding{
-		err: bedrock.ErrBedrockNoActiveModel,
-	}
 )
 
-func (h *Handler) AddBedrockIntegrationUsingIAMRole(w http.ResponseWriter, r *http.Request) {
-	defer r.Body.Close()
-	var req model.AutoAddGenerativeAiBedrockIntegration
-	err := httpext.DecodeJSON(r, httpext.NoQueryParams, MaxPostRequestSize, &req)
+func (h *Handler) AddGenerativeAIIntegrationUsingIAMRole(w http.ResponseWriter, r *http.Request) {
+	// Only AWS at the moment
+	foundModel, err := bedrock.CheckBedrockModelAvailability()
 	if err != nil {
 		log.Error().Msgf("%v", err)
 		h.respondError(&BadDecoding{err}, w)
 		return
 	}
-
-	addModelRequests, err := bedrock.ListBedrockModels(req.AWSRegion)
-	if err != nil {
-		log.Error().Msgf("%v", err)
-		h.respondError(&BadDecoding{err}, w)
+	if !foundModel {
+		h.respondError(&BadDecoding{err: bedrock.ErrBedrockNoActiveModel}, w)
 		return
 	}
 
-	if len(addModelRequests) == 0 {
-		log.Error().Msgf("%v", ErrBedrockNoActiveModel)
-		h.respondError(&ErrBedrockNoActiveModel, w)
-		return
-	}
-
-	ctx := r.Context()
-	pgClient, err := directory.PostgresClient(ctx)
+	worker, err := directory.Worker(r.Context())
 	if err != nil {
-		h.respondError(&InternalServerError{err}, w)
+		h.respondError(err, w)
 		return
 	}
-
-	// encrypt secret
-	aesValue, err := model.GetAESValueForEncryption(ctx, pgClient)
+	user, statusCode, _, err := h.GetUserFromJWT(r.Context())
 	if err != nil {
-		log.Error().Msgf(err.Error())
-		h.respondError(&InternalServerError{err}, w)
+		h.respondWithErrorCode(err, w, statusCode)
 		return
 	}
-	aes := encryption.AES{}
-	err = json.Unmarshal(aesValue, &aes)
+	data := utils.AutoFetchGenerativeAIIntegrationsParameters{
+		CloudProvider: "aws",
+		UserID:        user.ID,
+	}
+	dataJson, err := json.Marshal(data)
 	if err != nil {
-		log.Error().Msgf(err.Error())
-		h.respondError(&InternalServerError{err}, w)
+		h.respondError(err, w)
 		return
 	}
-	user, statusCode, _, err := h.GetUserFromJWT(ctx)
+	err = worker.Enqueue(utils.AutoFetchGenerativeAIIntegrations, dataJson, utils.DefaultTaskOpts()...)
 	if err != nil {
-		h.respondWithErrorCode(err, w, statusCode)
+		h.respondError(err, w)
 		return
 	}
 
-	for _, addModelRequest := range addModelRequests {
-		err = h.AddGenerativeAiIntegrationHelper(ctx, addModelRequest, aes, user, pgClient)
-		if err != nil {
-			log.Warn().Msgf(err.Error())
-			continue
-		}
-	}
-
-	httpext.JSON(w, http.StatusOK, model.MessageResponse{Message: api_messages.SuccessIntegrationCreated})
+	w.WriteHeader(http.StatusAccepted)
 }
 
 func (h *Handler) AddOpenAiIntegration(w http.ResponseWriter, r *http.Request) {
-	defer r.Body.Close()
-	var req model.AddGenerativeAiOpenAIIntegration
-	err := httpext.DecodeJSON(r, httpext.NoQueryParams, MaxPostRequestSize, &req)
-	if err != nil {
-		log.Error().Msgf("%v", err)
-		h.respondError(&BadDecoding{err}, w)
-		return
-	}
-	h.AddGenerativeAiIntegration(req, w, r)
+	AddGenerativeAiIntegration[model.AddGenerativeAiOpenAIIntegration](w, r, h)
 }
 
 func (h *Handler) AddBedrockIntegration(w http.ResponseWriter, r *http.Request) {
+	AddGenerativeAiIntegration[model.AddGenerativeAiBedrockIntegration](w, r, h)
+}
+
+func AddGenerativeAiIntegration[T model.AddGenerativeAiIntegrationRequest](w http.ResponseWriter, r *http.Request, h *Handler) {
 	defer r.Body.Close()
-	var req model.AddGenerativeAiBedrockIntegration
+	var req T
 	err := httpext.DecodeJSON(r, httpext.NoQueryParams, MaxPostRequestSize, &req)
 	if err != nil {
 		log.Error().Msgf("%v", err)
 		h.respondError(&BadDecoding{err}, w)
 		return
 	}
-	h.AddGenerativeAiIntegration(req, w, r)
-}
 
-func (h *Handler) AddGenerativeAiIntegration(req model.AddGenerativeAiIntegrationRequest, w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	pgClient, err := directory.PostgresClient(ctx)
 	if err != nil {

deepfence_server/model/generative_ai_integration.go

@@ -215,7 +215,3 @@ type GenerativeAiIntegrationListResponse struct {
 	LastErrorMsg       string `json:"last_error_msg"`
 	DefaultIntegration bool   `json:"default_integration"`
 }
-
-type AutoAddGenerativeAiBedrockIntegration struct {
-	AWSRegion string `json:"aws_region" validate:"required,oneof=us-east-1 us-east-2 us-west-1 us-west-2 af-south-1 ap-east-1 ap-south-1 ap-northeast-1 ap-northeast-2 ap-northeast-3 ap-southeast-1 ap-southeast-2 ap-southeast-3 ca-central-1 eu-central-1 eu-west-1 eu-west-2 eu-west-3 eu-south-1 eu-north-1 me-south-1 me-central-1 sa-east-1 us-gov-east-1 us-gov-west-1" required:"true" enum:"us-east-1,us-east-2,us-west-1,us-west-2,af-south-1,ap-east-1,ap-south-1,ap-northeast-1,ap-northeast-2,ap-northeast-3,ap-southeast-1,ap-southeast-2,ap-southeast-3,ca-central-1,eu-central-1,eu-west-1,eu-west-2,eu-west-3,eu-south-1,eu-north-1,me-south-1,me-central-1,sa-east-1,us-gov-east-1,us-gov-west-1"`
-}

deepfence_server/pkg/generative-ai-integration/bedrock/aimodels.go

@@ -1,12 +1,45 @@
 package bedrock
 
 import (
+	"encoding/json"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/bedrock"
+	"github.com/aws/aws-sdk-go/service/bedrockruntime"
 	"github.com/deepfence/ThreatMapper/deepfence_server/model"
+	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
 )
 
+func CheckBedrockModelAvailability() (bool, error) {
+	foundModel := false
+	for _, region := range BedrockRegions {
+		sess, err := session.NewSession(&aws.Config{
+			Region: aws.String(region),
+		})
+		if err != nil {
+			return false, err
+		}
+		svc := bedrock.New(sess)
+
+		models, err := svc.ListFoundationModels(&bedrock.ListFoundationModelsInput{
+			ByOutputModality: &textModality,
+		})
+		for _, modelSummary := range models.ModelSummaries {
+			if *modelSummary.ModelLifecycle.Status == modelLifecycleActive {
+				if _, ok := BedrockModelBody[*modelSummary.ModelId]; ok {
+					foundModel = true
+					break
+				}
+			}
+		}
+		if foundModel == true {
+			break
+		}
+	}
+	return foundModel, nil
+}
+
 // ListBedrockModels Fetch enabled Bedrock models using IAM roles
 func ListBedrockModels(region string) ([]model.AddGenerativeAiIntegrationRequest, error) {
 	sess, err := session.NewSession(&aws.Config{
@@ -15,24 +48,47 @@ func ListBedrockModels(region string) ([]model.AddGenerativeAiIntegrationRequest
 	if err != nil {
 		return nil, err
 	}
+
 	svc := bedrock.New(sess)
+	bedrockRuntimeSvc := bedrockruntime.New(sess)
+
 	models, err := svc.ListFoundationModels(&bedrock.ListFoundationModelsInput{
 		ByOutputModality: &textModality,
 	})
 	if err != nil {
 		return nil, err
 	}
-	resp := []model.AddGenerativeAiIntegrationRequest{}
+
+	var bedrockModels []model.AddGenerativeAiIntegrationRequest
+	message := "hello"
+
 	for _, modelSummary := range models.ModelSummaries {
 		if *modelSummary.ModelLifecycle.Status == modelLifecycleActive {
-			if _, ok := BedrockModelBody[*modelSummary.ModelId]; ok {
-				resp = append(resp, model.AddGenerativeAiBedrockIntegration{
+			if body, ok := BedrockModelBody[*modelSummary.ModelId]; ok {
+				body[BedrockModelBodyInputKey[*modelSummary.ModelId]] = body[BedrockModelBodyInputKey[*modelSummary.ModelId]].(string) + message + BedrockModelBodyInputSuffix[*modelSummary.ModelId]
+				bodyBytes, err := json.Marshal(body)
+				if err != nil {
+					log.Warn().Msg(err.Error())
+					continue
+				}
+				_, err = bedrockRuntimeSvc.InvokeModel(&bedrockruntime.InvokeModelInput{
+					Accept:      &acceptHeader,
+					Body:        bodyBytes,
+					ContentType: &contentTypeHeader,
+					ModelId:     modelSummary.ModelId,
+				})
+				if err != nil {
+					log.Warn().Msg(err.Error())
+					continue
+				}
+
+				bedrockModels = append(bedrockModels, model.AddGenerativeAiBedrockIntegration{
 					AWSRegion:  region,
 					UseIAMRole: true,
 					ModelID:    *modelSummary.ModelId,
 				})
 			}
 		}
 	}
-	return resp, nil
+	return bedrockModels, nil
 }

deepfence_server/pkg/generative-ai-integration/bedrock/types.go

@@ -24,6 +24,8 @@ const (
 )
 
 var (
+	BedrockRegions = []string{"us-east-1", "us-west-2", "ap-southeast-1", "ap-northeast-1", "eu-central-1"}
+
 	textModality = "TEXT"
 
 	contentTypeHeader = "application/json"

deepfence_server/router/router.go

@@ -511,7 +511,7 @@ func SetupRoutes(r *chi.Mux, serverPort string, serveOpenapiDocs bool, ingestC c
 			r.Route("/generative-ai-integration", func(r chi.Router) {
 				r.Post("/openai", dfHandler.AuthHandler(ResourceIntegration, PermissionWrite, dfHandler.AddOpenAiIntegration))
 				r.Post("/bedrock", dfHandler.AuthHandler(ResourceIntegration, PermissionWrite, dfHandler.AddBedrockIntegration))
-				r.Post("/bedrock/auto-add", dfHandler.AuthHandler(ResourceIntegration, PermissionWrite, dfHandler.AddBedrockIntegrationUsingIAMRole))
+				r.Post("/auto-add", dfHandler.AuthHandler(ResourceIntegration, PermissionWrite, dfHandler.AddGenerativeAIIntegrationUsingIAMRole))
 
 				r.Get("/", dfHandler.AuthHandler(ResourceIntegration, PermissionRead, dfHandler.GetGenerativeAiIntegrations))
 				r.Route("/{integration_id}", func(r chi.Router) {

deepfence_utils/utils/constants.go

@@ -54,6 +54,7 @@ const (
 	UpdateCloudResourceScanStatusTask = "update_cloud_resource_scan_status"
 	UpdatePodScanStatusTask           = "update_pod_scan_status"
 	BulkDeleteScans                   = "bulk_delete_scans"
+	AutoFetchGenerativeAIIntegrations = "auto_fetch_generative_ai_integrations"
 )
 
 const (

deepfence_utils/utils/structs.go

@@ -42,6 +42,11 @@ type SbomBody struct {
 	SBOM string `json:"sbom" required:"true"`
 }
 
+type AutoFetchGenerativeAIIntegrationsParameters struct {
+	CloudProvider string `json:"cloud_provider"`
+	UserID        int64  `json:"user_id"`
+}
+
 type SecretScanParameters struct {
 	ImageName             string `json:"image_name"`
 	ImageID               string `json:"image_id"`

deepfence_worker/go.mod

@@ -35,7 +35,7 @@ require (
 	github.com/deepfence/ThreatMapper/deepfence_utils v0.0.0-00010101000000-000000000000
 	github.com/deepfence/YaraHunter v0.0.0-00010101000000-000000000000
 	github.com/deepfence/agent-plugins-grpc v1.1.0
-	github.com/deepfence/golang_deepfence_sdk/utils v0.0.0-20231201162830-c15be6d96147
+	github.com/deepfence/golang_deepfence_sdk/utils v0.0.0-20231201173641-092afefd00a2
 	github.com/deepfence/package-scanner v0.0.0-00010101000000-000000000000
 	github.com/hibiken/asynq v0.24.1
 	github.com/kelseyhightower/envconfig v1.4.0
@@ -92,7 +92,7 @@ require (
 	github.com/containerd/typeurl/v2 v2.1.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.2.0 // indirect
-	github.com/deepfence/golang_deepfence_sdk/client v0.0.0-20231030062708-5506162b00b7 // indirect
+	github.com/deepfence/golang_deepfence_sdk/client v0.0.0-20231201173641-092afefd00a2 // indirect
 	github.com/deepfence/vessel v0.12.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/distribution/reference v0.5.0 // indirect