Option for user to automatically add Amazon Bedrock integrations (#1813)

Option for user to automatically add Amazon Bedrock integrations when read only instance role (IAM) is set #1759

deepfence_server/apiDocs/operation.go

@@ -718,6 +718,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("listGenerativeAiIntegration", http.MethodGet, "/deepfence/generative-ai-integration",
 		"List Generative AI Integrations", "List all the added Generative AI Integrations",

deepfence_server/go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/Masterminds/sprig/v3 v3.2.0
 	github.com/PagerDuty/go-pagerduty v1.7.0
 	github.com/andygrunwald/go-jira v1.16.0
-	github.com/aws/aws-sdk-go v1.48.1
+	github.com/aws/aws-sdk-go v1.48.12
 	github.com/casbin/casbin/v2 v2.75.0
 	github.com/deepfence/ThreatMapper/deepfence_utils v0.0.0-00010101000000-000000000000
 	github.com/docker/docker v24.0.5+incompatible

deepfence_server/go.sum

@@ -16,8 +16,8 @@ github.com/XSAM/otelsql v0.25.0 h1:ji1G+O45lrmZV9pXv2jQNRzYVFIwEB0jlY0XXdgpuNk=
 github.com/XSAM/otelsql v0.25.0/go.mod h1:VfWJ7nRF1t74mSL36s0ksIohT4nmFH5/opajHcmXPFc=
 github.com/andygrunwald/go-jira v1.16.0 h1:PU7C7Fkk5L96JvPc6vDVIrd99vdPnYudHu4ju2c2ikQ=
 github.com/andygrunwald/go-jira v1.16.0/go.mod h1:UQH4IBVxIYWbgagc0LF/k9FRs9xjIiQ8hIcC6HfLwFU=
-github.com/aws/aws-sdk-go v1.48.1 h1:OXPUVL4cLdsDsqkVIuhwY+D389tjI7e1xu0lsDYyeMk=
-github.com/aws/aws-sdk-go v1.48.1/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3TjupRn3Eqk=
+github.com/aws/aws-sdk-go v1.48.12 h1:n+eGzflzzvYubu2cOjqpVll7lF+Ci0ThyCpg5kzfzbo=
+github.com/aws/aws-sdk-go v1.48.12/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3TjupRn3Eqk=
 github.com/bool64/dev v0.2.29 h1:x+syGyh+0eWtOzQ1ItvLzOGIWyNWnyjXpHIcpF2HvL4=
 github.com/bool64/dev v0.2.29/go.mod h1:iJbh1y/HkunEPhgebWRNcs8wfGq7sjvJ6W5iabL8ACg=
 github.com/bool64/shared v0.1.5 h1:fp3eUhBsrSjNCQPcSdQqZxxh9bBwrYiZ+zOKFkM0/2E=

deepfence_server/handler/generative_ai_integration.go

@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"errors"
@@ -10,6 +11,7 @@ import (
 	api_messages "github.com/deepfence/ThreatMapper/deepfence_server/constants/api-messages"
 	"github.com/deepfence/ThreatMapper/deepfence_server/model"
 	generative_ai_integration "github.com/deepfence/ThreatMapper/deepfence_server/pkg/generative-ai-integration"
+	"github.com/deepfence/ThreatMapper/deepfence_server/pkg/generative-ai-integration/bedrock"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/encryption"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
@@ -18,42 +20,106 @@ import (
 	httpext "github.com/go-playground/pkg/v5/net/http"
 )
 
-var ErrStreamUnsupported = errors.New("streaming unsupported")
-
-func (h *Handler) AddOpenAiIntegration(w http.ResponseWriter, r *http.Request) {
-	AddGenerativeAiIntegration[model.AddGenerativeAiOpenAIIntegration](w, r, h)
-}
-
-func (h *Handler) AddBedrockIntegration(w http.ResponseWriter, r *http.Request) {
-	AddGenerativeAiIntegration[model.AddGenerativeAiBedrockIntegration](w, r, h)
-}
+var (
+	ErrStreamUnsupported             = errors.New("streaming unsupported")
+	ErrGenerativeAIIntegrationExists = BadDecoding{
+		err: errors.New("similar integration already exists"),
+	}
+	ErrBedrockNoActiveModel = BadDecoding{
+		err: bedrock.ErrBedrockNoActiveModel,
+	}
+)
 
-func AddGenerativeAiIntegration[T model.AddGenerativeAiIntegrationRequest](w http.ResponseWriter, r *http.Request, h *Handler) {
+func (h *Handler) AddBedrockIntegrationUsingIAMRole(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
-	var req T
+	var req model.AutoAddGenerativeAiBedrockIntegration
 	err := httpext.DecodeJSON(r, httpext.NoQueryParams, MaxPostRequestSize, &req)
 	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)
+		return
+	}
+
+	if len(addModelRequests) == 0 {
+		log.Error().Msgf("%v", ErrBedrockNoActiveModel)
+		h.respondError(&ErrBedrockNoActiveModel, w)
+		return
+	}
+
 	ctx := r.Context()
 	pgClient, err := directory.PostgresClient(ctx)
 	if err != nil {
 		h.respondError(&InternalServerError{err}, w)
 		return
 	}
 
-	obj, err := generative_ai_integration.NewGenerativeAiIntegration(ctx, req)
+	// encrypt secret
+	aesValue, err := model.GetAESValueForEncryption(ctx, pgClient)
+	if err != nil {
+		log.Error().Msgf(err.Error())
+		h.respondError(&InternalServerError{err}, w)
+		return
+	}
+	aes := encryption.AES{}
+	err = json.Unmarshal(aesValue, &aes)
+	if err != nil {
+		log.Error().Msgf(err.Error())
+		h.respondError(&InternalServerError{err}, w)
+		return
+	}
+	user, statusCode, _, err := h.GetUserFromJWT(ctx)
+	if err != nil {
+		h.respondWithErrorCode(err, w, statusCode)
+		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})
+}
+
+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
 	}
-	err = obj.ValidateConfig(h.Validator)
+	h.AddGenerativeAiIntegration(req, w, r)
+}
+
+func (h *Handler) AddBedrockIntegration(w http.ResponseWriter, r *http.Request) {
+	defer r.Body.Close()
+	var req model.AddGenerativeAiBedrockIntegration
+	err := httpext.DecodeJSON(r, httpext.NoQueryParams, MaxPostRequestSize, &req)
 	if err != nil {
-		h.respondError(&ValidatorError{err: err}, w)
+		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 {
+		h.respondError(&InternalServerError{err}, w)
 		return
 	}
 
@@ -71,40 +137,60 @@ func AddGenerativeAiIntegration[T model.AddGenerativeAiIntegrationRequest](w htt
 		h.respondError(&InternalServerError{err}, w)
 		return
 	}
-	err = obj.EncryptSecret(aes)
+	user, statusCode, _, err := h.GetUserFromJWT(ctx)
 	if err != nil {
-		log.Error().Msgf(err.Error())
-		h.respondError(&InternalServerError{err}, w)
+		h.respondWithErrorCode(err, w, statusCode)
 		return
 	}
 
-	// add integration to database
-	// before that check if integration already exists
-	integrationExists, err := req.IntegrationExists(ctx, pgClient)
+	err = h.AddGenerativeAiIntegrationHelper(ctx, req, aes, user, pgClient)
 	if err != nil {
 		log.Error().Msgf(err.Error())
-		h.respondError(&InternalServerError{err}, w)
+		h.respondError(err, w)
 		return
 	}
-	if integrationExists {
-		err = httpext.JSON(w, http.StatusBadRequest, model.ErrorResponse{Message: api_messages.ErrIntegrationExists})
-		if err != nil {
-			log.Error().Msg(err.Error())
-		}
-		return
+
+	h.AuditUserActivity(r, EventGenerativeAIIntegration, ActionCreate, map[string]interface{}{"integration_type": req.GetIntegrationType()}, true)
+
+	err = httpext.JSON(w, http.StatusOK, model.MessageResponse{Message: api_messages.SuccessIntegrationCreated})
+	if err != nil {
+		log.Error().Msg(err.Error())
 	}
+}
 
-	user, statusCode, _, err := h.GetUserFromJWT(ctx)
+func (h *Handler) AddGenerativeAiIntegrationHelper(ctx context.Context, req model.AddGenerativeAiIntegrationRequest, aes encryption.AES, user *model.User, pgClient *postgresqlDb.Queries) error {
+	obj, err := generative_ai_integration.NewGenerativeAiIntegration(ctx, req)
 	if err != nil {
-		h.respondWithErrorCode(err, w, statusCode)
-		return
+		return &BadDecoding{err}
+	}
+	err = obj.ValidateConfig(h.Validator)
+	if err != nil {
+		return &ValidatorError{err: err}
+	}
+	err = obj.VerifyAuth(ctx)
+	if err != nil {
+		return &BadDecoding{err: err}
+	}
+
+	err = obj.EncryptSecret(aes)
+	if err != nil {
+		return err
+	}
+
+	// add integration to database
+	// before that check if integration already exists
+	integrationExists, err := req.IntegrationExists(ctx, pgClient)
+	if err != nil {
+		return err
+	}
+	if integrationExists {
+		return &ErrGenerativeAIIntegrationExists
 	}
 
 	// store the integration in db
 	bConfig, err := json.Marshal(obj)
 	if err != nil {
-		h.respondWithErrorCode(err, w, statusCode)
-		return
+		return err
 	}
 
 	arg := postgresqlDb.CreateGenerativeAiIntegrationParams{
@@ -115,22 +201,14 @@ func AddGenerativeAiIntegration[T model.AddGenerativeAiIntegrationRequest](w htt
 	}
 	dbIntegration, err := pgClient.CreateGenerativeAiIntegration(ctx, arg)
 	if err != nil {
-		log.Error().Msgf(err.Error())
-		h.respondError(&InternalServerError{err}, w)
-		return
+		return err
 	}
 
-	h.AuditUserActivity(r, EventGenerativeAIIntegration, ActionCreate, map[string]interface{}{"integration_type": req.GetIntegrationType()}, true)
-
 	err = pgClient.UpdateGenerativeAiIntegrationDefault(ctx, dbIntegration.ID)
 	if err != nil {
 		log.Warn().Msgf(err.Error())
 	}
-
-	err = httpext.JSON(w, http.StatusOK, model.MessageResponse{Message: api_messages.SuccessIntegrationCreated})
-	if err != nil {
-		log.Error().Msg(err.Error())
-	}
+	return nil
 }
 
 func (h *Handler) GetGenerativeAiIntegrations(w http.ResponseWriter, r *http.Request) {

deepfence_server/model/generative_ai_integration.go

@@ -215,3 +215,7 @@ 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

@@ -0,0 +1,38 @@
+package bedrock
+
+import (
+	"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/deepfence/ThreatMapper/deepfence_server/model"
+)
+
+// ListBedrockModels Fetch enabled Bedrock models using IAM roles
+func ListBedrockModels(region string) ([]model.AddGenerativeAiIntegrationRequest, error) {
+	sess, err := session.NewSession(&aws.Config{
+		Region: aws.String(region),
+	})
+	if err != nil {
+		return nil, err
+	}
+	svc := bedrock.New(sess)
+	models, err := svc.ListFoundationModels(&bedrock.ListFoundationModelsInput{
+		ByOutputModality: &textModality,
+	})
+	if err != nil {
+		return nil, err
+	}
+	resp := []model.AddGenerativeAiIntegrationRequest{}
+	for _, modelSummary := range models.ModelSummaries {
+		if *modelSummary.ModelLifecycle.Status == modelLifecycleActive {
+			if _, ok := BedrockModelBody[*modelSummary.ModelId]; ok {
+				resp = append(resp, model.AddGenerativeAiBedrockIntegration{
+					AWSRegion:  region,
+					UseIAMRole: true,
+					ModelID:    *modelSummary.ModelId,
+				})
+			}
+		}
+	}
+	return resp, nil
+}

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

@@ -37,6 +37,24 @@ func (b *Bedrock) ValidateConfig(validate *validator.Validate) error {
 	return validate.Struct(b)
 }
 
+func (b *Bedrock) VerifyAuth(ctx context.Context) error {
+	var err error
+	if b.UseIAMRole {
+		_, err = session.NewSession(&aws.Config{
+			Region: aws.String(b.AWSRegion),
+		})
+	} else {
+		_, err = session.NewSession(&aws.Config{
+			Region:      aws.String(b.AWSRegion),
+			Credentials: credentials.NewStaticCredentials(b.AWSAccessKey, b.AWSSecretKey, ""),
+		})
+	}
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (b *Bedrock) EncryptSecret(aes encryption.AES) error {
 	var err error
 	b.AWSSecretKey, err = aes.Encrypt(b.AWSSecretKey)

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

@@ -1,6 +1,9 @@
 package bedrock
 
 import (
+	"fmt"
+	"strings"
+
 	"github.com/deepfence/ThreatMapper/deepfence_server/pkg/generative-ai-integration/common"
 )
 
@@ -16,13 +19,27 @@ const (
 	MetaLlama2Chat13BV1           = "meta.llama2-13b-chat-v1"
 	CohereCommandTextV14          = "cohere.command-text-v14"
 	CohereCommandLightTextV14     = "cohere.command-light-text-v14"
+
+	modelLifecycleActive = "ACTIVE"
 )
 
 var (
+	textModality = "TEXT"
+
 	contentTypeHeader = "application/json"
 	acceptHeader      = "*/*"
+
+	ErrBedrockNoActiveModel error
 )
 
+func init() {
+	supportedModels := ""
+	for modelID, _ := range BedrockModelBody {
+		supportedModels += modelID + ", "
+	}
+	ErrBedrockNoActiveModel = fmt.Errorf("no active model, one of these should be active: %s", strings.Trim(supportedModels, ", "))
+}
+
 type Bedrock struct {
 	common.GenerativeAiIntegrationCommon
 	AWSAccessKey string `json:"aws_access_key" validate:"omitempty,min=16,max=128"`

deepfence_server/pkg/generative-ai-integration/generative_ai_integration.go

@@ -40,6 +40,7 @@ func NewGenerativeAiIntegrationFromDBEntry(ctx context.Context, integrationType
 // GenerativeAiIntegration is the interface for all integrations
 type GenerativeAiIntegration interface {
 	ValidateConfig(*validator.Validate) error
+	VerifyAuth(ctx context.Context) error
 	GenerateCloudPostureQuery(model.GenerativeAiIntegrationRequest) (string, error)
 	GenerateLinuxPostureQuery(model.GenerativeAiIntegrationRequest) (string, error)
 	GenerateKubernetesPostureQuery(model.GenerativeAiIntegrationRequest) (string, error)