diff --git a/deepfence_server/apiDocs/operation.go b/deepfence_server/apiDocs/operation.go index c347ac6216..92378b1ed5 100644 --- a/deepfence_server/apiDocs/operation.go +++ b/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", diff --git a/deepfence_server/handler/generative_ai_integration.go b/deepfence_server/handler/generative_ai_integration.go index 05c801479c..a1c89479e0 100644 --- a/deepfence_server/handler/generative_ai_integration.go +++ b/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 { diff --git a/deepfence_server/model/generative_ai_integration.go b/deepfence_server/model/generative_ai_integration.go index 59235d4e21..13c04477b5 100644 --- a/deepfence_server/model/generative_ai_integration.go +++ b/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"` -} diff --git a/deepfence_server/pkg/generative-ai-integration/bedrock/aimodels.go b/deepfence_server/pkg/generative-ai-integration/bedrock/aimodels.go index 84f790ea2f..ae60f0369f 100644 --- a/deepfence_server/pkg/generative-ai-integration/bedrock/aimodels.go +++ b/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,18 +48,41 @@ 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, @@ -34,5 +90,5 @@ func ListBedrockModels(region string) ([]model.AddGenerativeAiIntegrationRequest } } } - return resp, nil + return bedrockModels, nil } diff --git a/deepfence_server/pkg/generative-ai-integration/bedrock/types.go b/deepfence_server/pkg/generative-ai-integration/bedrock/types.go index 9fa0c3f81c..c4dac0b0ef 100644 --- a/deepfence_server/pkg/generative-ai-integration/bedrock/types.go +++ b/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" diff --git a/deepfence_server/router/router.go b/deepfence_server/router/router.go index cbc4ca28c4..5fee18fc0a 100644 --- a/deepfence_server/router/router.go +++ b/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) { diff --git a/deepfence_utils/utils/constants.go b/deepfence_utils/utils/constants.go index 6fac530393..9a0c4541e5 100644 --- a/deepfence_utils/utils/constants.go +++ b/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 ( diff --git a/deepfence_utils/utils/structs.go b/deepfence_utils/utils/structs.go index 91c81045d1..987975f77e 100644 --- a/deepfence_utils/utils/structs.go +++ b/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"` diff --git a/deepfence_worker/go.mod b/deepfence_worker/go.mod index 937fc0a7ae..a9a2cc0d75 100644 --- a/deepfence_worker/go.mod +++ b/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 diff --git a/deepfence_worker/tasks/generativeai/generativeai.go b/deepfence_worker/tasks/generativeai/generativeai.go new file mode 100644 index 0000000000..b68e2bf428 --- /dev/null +++ b/deepfence_worker/tasks/generativeai/generativeai.go @@ -0,0 +1,134 @@ +package generativeai + +import ( + "context" + "encoding/json" + "errors" + + "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" + postgresqlDb "github.com/deepfence/ThreatMapper/deepfence_utils/postgresql/postgresql-db" + "github.com/deepfence/ThreatMapper/deepfence_utils/utils" + "github.com/hibiken/asynq" +) + +var ( + ErrBedrockNoModels = errors.New("could not automatically fetch Amazon Bedrock integrations, no active models found") + ErrGenerativeAIIntegrationExists = errors.New("similar integration already exists") +) + +func AutoFetchGenerativeAIIntegrations(ctx context.Context, task *asynq.Task) error { + var params utils.AutoFetchGenerativeAIIntegrationsParameters + + log.Info().Msgf("AutoFetchGenerativeAIIntegrations, payload: %s ", string(task.Payload())) + + if err := json.Unmarshal(task.Payload(), ¶ms); err != nil { + log.Error().Msgf("AutoFetchGenerativeAIIntegrations, error in Unmarshal: %s", err.Error()) + return nil + } + + var err error + switch string(task.Payload()) { + case "aws": + err = AutoFetchBedrockIntegrations(ctx, params) + } + if err != nil { + log.Error().Msg(err.Error()) + } + return nil +} + +func AutoFetchBedrockIntegrations(ctx context.Context, params utils.AutoFetchGenerativeAIIntegrationsParameters) error { + var foundRegion string + var err error + var models []model.AddGenerativeAiIntegrationRequest + for _, region := range bedrock.BedrockRegions { + models, err = bedrock.ListBedrockModels(region) + if err != nil { + log.Warn().Msg(err.Error()) + continue + } + if len(models) == 0 { + continue + } + foundRegion = region + break + } + if foundRegion == "" { + return ErrBedrockNoModels + } + + pgClient, err := directory.PostgresClient(ctx) + if err != nil { + return err + } + // encrypt secret + aesValue, err := model.GetAESValueForEncryption(ctx, pgClient) + if err != nil { + return err + } + aes := encryption.AES{} + err = json.Unmarshal(aesValue, &aes) + if err != nil { + return err + } + + for _, req := range models { + err = CreateGenerativeAIModel(ctx, req, aes, params.UserID, pgClient) + if err != nil { + log.Error().Msg(err.Error()) + continue + } + } + + return nil +} + +func CreateGenerativeAIModel(ctx context.Context, req model.AddGenerativeAiIntegrationRequest, aes encryption.AES, userID int64, pgClient *postgresqlDb.Queries) error { + obj, err := generative_ai_integration.NewGenerativeAiIntegration(ctx, req) + if err != nil { + return 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 { + return err + } + + arg := postgresqlDb.CreateGenerativeAiIntegrationParams{ + IntegrationType: req.GetIntegrationType(), + Label: req.GetLabel(), + Config: bConfig, + CreatedByUserID: userID, + } + dbIntegration, err := pgClient.CreateGenerativeAiIntegration(ctx, arg) + if err != nil { + return err + } + + err = pgClient.UpdateGenerativeAiIntegrationDefault(ctx, dbIntegration.ID) + if err != nil { + log.Warn().Msgf(err.Error()) + } + return nil +} diff --git a/deepfence_worker/worker.go b/deepfence_worker/worker.go index 9227b435f3..2da4aa2a91 100644 --- a/deepfence_worker/worker.go +++ b/deepfence_worker/worker.go @@ -12,6 +12,7 @@ import ( "github.com/deepfence/ThreatMapper/deepfence_utils/telemetry" "github.com/deepfence/ThreatMapper/deepfence_utils/utils" "github.com/deepfence/ThreatMapper/deepfence_worker/cronjobs" + "github.com/deepfence/ThreatMapper/deepfence_worker/tasks/generativeai" "github.com/deepfence/ThreatMapper/deepfence_worker/tasks/malwarescan" "github.com/deepfence/ThreatMapper/deepfence_worker/tasks/reports" "github.com/deepfence/ThreatMapper/deepfence_worker/tasks/sbom" @@ -199,6 +200,8 @@ func NewWorker(ns directory.NamespaceID, cfg config) (Worker, context.CancelFunc worker.AddOneShotHandler(utils.LinkNodesTask, cronjobs.LinkNodes) + worker.AddOneShotHandler(utils.AutoFetchGenerativeAIIntegrations, generativeai.AutoFetchGenerativeAIIntegrations) + // sbom worker.AddRetryableHandler(utils.ScanSBOMTask, sbom.NewSBOMScanner(ingestC).ScanSBOM)