feat: Added for encrypting the database and blob storage with WB-mana…

…ged key (#49) Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Aastha Gupta <[email protected]> Co-authored-by: Aastha Gupta <[email protected]> Co-authored-by: Zachary Blasczyk <[email protected]>
wandb · Aug 1, 2024 · 519c340 · 519c340
1 parent f6ad2d7
commit 519c340
Show file tree

Hide file tree

Showing 22 changed files with 412 additions and 68 deletions.
diff --git a/README.md b/README.md
@@ -70,19 +70,24 @@ resources that lack official modules.
 
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
-| <a name="input_allowed_ip_ranges"></a> [allowed\_ip\_ranges](#input\_allowed\_ip\_ranges) | allowed public IP addresses or CIDR ranges. | `list(string)` | `[]` | no |
+| <a name="input_allowed_ip_ranges"></a> [allowed\_ip\_ranges](#input\_allowed\_ip\_ranges) | Allowed public IP addresses or CIDR ranges. | `list(string)` | `[]` | no |
 | <a name="input_allowed_subscriptions"></a> [allowed\_subscriptions](#input\_allowed\_subscriptions) | List of allowed customer subscriptions coma seperated values | `string` | `""` | no |
 | <a name="input_app_wandb_env"></a> [app\_wandb\_env](#input\_app\_wandb\_env) | Extra environment variables for W&B | `map(string)` | `{}` | no |
 | <a name="input_azuremonitor"></a> [azuremonitor](#input\_azuremonitor) | # To support otel azure monitor sql and redis metrics need operator-wandb chart minimum version 0.14.0 | `bool` | `false` | no |
 | <a name="input_blob_container"></a> [blob\_container](#input\_blob\_container) | Use an existing bucket. | `string` | `""` | no |
 | <a name="input_cluster_sku_tier"></a> [cluster\_sku\_tier](#input\_cluster\_sku\_tier) | The Azure AKS SKU Tier to use for this cluster (https://learn.microsoft.com/en-us/azure/aks/free-standard-pricing-tiers) | `string` | `"Free"` | no |
 | <a name="input_create_private_link"></a> [create\_private\_link](#input\_create\_private\_link) | Use for the azure private link. | `bool` | `false` | no |
 | <a name="input_create_redis"></a> [create\_redis](#input\_create\_redis) | Boolean indicating whether to provision an redis instance (true) or not (false). | `bool` | `false` | no |
+| <a name="input_customer_storage_vault_key_id"></a> [customer\_storage\_vault\_key\_id](#input\_customer\_storage\_vault\_key\_id) | The Azure Key Vault key ID for customer-provided storage encryption keys. Must match the pattern 'https://<vault-name>.vault.azure.net/keys/<key-name>/<key-version>', or be null. | `string` | `null` | no |
 | <a name="input_database_availability_mode"></a> [database\_availability\_mode](#input\_database\_availability\_mode) | n/a | `string` | `"SameZone"` | no |
 | <a name="input_database_sku_name"></a> [database\_sku\_name](#input\_database\_sku\_name) | Specifies the SKU Name for this MySQL Server | `string` | `"GP_Standard_D4ds_v4"` | no |
 | <a name="input_database_version"></a> [database\_version](#input\_database\_version) | Version for MySQL | `string` | `"5.7"` | no |
 | <a name="input_deletion_protection"></a> [deletion\_protection](#input\_deletion\_protection) | If the instance should have deletion protection enabled. The database / Bucket can't be deleted when this value is set to `true`. | `bool` | `true` | no |
+| <a name="input_disable_storage_vault_key_id"></a> [disable\_storage\_vault\_key\_id](#input\_disable\_storage\_vault\_key\_id) | Flag to disable the `customer_managed_key` block, the properties 'encryption.identity, encryption.keyvaultproperties' cannot be updated in a single operation. | `bool` | `false` | no |
 | <a name="input_domain_name"></a> [domain\_name](#input\_domain\_name) | Domain for accessing the Weights & Biases UI. | `string` | `null` | no |
+| <a name="input_enable_database_vault_key"></a> [enable\_database\_vault\_key](#input\_enable\_database\_vault\_key) | Flag to enable managed key encryption for the database. Once enabled, cannot be disabled. | `bool` | `false` | no |
+| <a name="input_enable_purge_protection"></a> [enable\_purge\_protection](#input\_enable\_purge\_protection) | Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled. | `bool` | `false` | no |
+| <a name="input_enable_storage_vault_key"></a> [enable\_storage\_vault\_key](#input\_enable\_storage\_vault\_key) | Flag to enable managed key encryption for the storage account. | `bool` | `false` | no |
 | <a name="input_external_bucket"></a> [external\_bucket](#input\_external\_bucket) | config an external bucket | `any` | `null` | no |
 | <a name="input_kubernetes_instance_type"></a> [kubernetes\_instance\_type](#input\_kubernetes\_instance\_type) | Use for the Kubernetes cluster. | `string` | `"Standard_D4a_v4"` | no |
 | <a name="input_kubernetes_node_count"></a> [kubernetes\_node\_count](#input\_kubernetes\_node\_count) | n/a | `number` | `2` | no |
@@ -116,6 +121,7 @@ resources that lack official modules.
 | <a name="output_address"></a> [address](#output\_address) | n/a |
 | <a name="output_aks_node_count"></a> [aks\_node\_count](#output\_aks\_node\_count) | n/a |
 | <a name="output_aks_node_instance_type"></a> [aks\_node\_instance\_type](#output\_aks\_node\_instance\_type) | n/a |
+| <a name="output_client_id"></a> [client\_id](#output\_client\_id) | n/a |
 | <a name="output_cluster_ca_certificate"></a> [cluster\_ca\_certificate](#output\_cluster\_ca\_certificate) | n/a |
 | <a name="output_cluster_client_certificate"></a> [cluster\_client\_certificate](#output\_cluster\_client\_certificate) | n/a |
 | <a name="output_cluster_client_key"></a> [cluster\_client\_key](#output\_cluster\_client\_key) | n/a |
@@ -128,5 +134,6 @@ resources that lack official modules.
 | <a name="output_standardized_size"></a> [standardized\_size](#output\_standardized\_size) | n/a |
 | <a name="output_storage_account"></a> [storage\_account](#output\_storage\_account) | n/a |
 | <a name="output_storage_container"></a> [storage\_container](#output\_storage\_container) | n/a |
+| <a name="output_tenant_id"></a> [tenant\_id](#output\_tenant\_id) | n/a |
 | <a name="output_url"></a> [url](#output\_url) | The URL to the W&B application |
 <!-- END_TF_DOCS -->
diff --git a/examples/byob/README.md b/examples/byob/README.md
@@ -23,13 +23,53 @@ storage_key = <sensitive>
 
 To retrieve the storage key, you can use the Azure CLI installed previously like the example below.
 
-```basb
-az storage account keys list --account-name rgnamestorage --query '[].{key: value}' --output tsv
+```bash
+az storage account keys list --account-name <storage_account_name> --resource-group <resource_group_name> --query '[0].value' -o tsv
 1111111111111122222222222333333333334444444555555555
-5555555554444444333333333332222222222211111111111111
 ```
 
-You only need to provide one key.
+This command will return the storage key, which you can then use for your deployment needs. Ensure you handle the storage key securely as it contains sensitive information.
+
+# Customer Managed Key Encryption
+
+The following section provides details on enabling Customer Managed Key (CMK) encryption for the Azure Blob Storage container which is disabled by default.
+
+To configure Customer Managed Key encryption, ensure you are using the latest version of out terraform which has the following added to the `variables.tf` file:
+
+- create_cmk
+- disable_storage_vault_key_id
+- enable_purge_protection
+- tenant_id
+- client_id
+
+You need to obtain the `tenant_id` and `client_id` from a Solutions Architect at W&B for an already instantiated instance of a Weights & Biases managed deployment.
+
+Set the follwoing new variabels to enable the CMK:
+
+```ini terraform.tfvars
+create_cmk = true
+
+enable_purge_protection       = true
+disable_storage_vault_key_id  = false
+
+tenant_id = "<tenant_id>"
+client_id = "<client_id>"
+```
+
+After updating your `terraform.tfvars` configuration, run the Terraform commands to apply the changes:
+
+```bash
+terraform init -upgrade
+terraform apply -var-file=terraform.tfvars
+```
+
+Upon successful execution, you will receive the following output which needs to the shared with the SA at Weights & Biases.
 
-* Note that all information about Storage Account and keys are mere examples, they are not valid.
+```bash
+blob_container = "<storage_account_name>/wandb"
+command_to_get_storage_key = "az storage account keys list --account-name <storage_account_name> --resource-group <resource_group_name> --query '[0].value' -o tsv"
+storage_key = <sensitive>
+storage_vault_key_id = "https://<key_vault_name>.vault.azure.net/keys/<key_name>/<key_version>"
+```
 
+Retrieve the storage key as shown above and also share that with the SA at Weights & Biases.
diff --git a/examples/byob/main.tf b/examples/byob/main.tf
@@ -3,14 +3,19 @@ provider "azurerm" {
 }
 
 module "byob" {
-  source              = "../../modules/byob"
-  resource_group_name = var.resource_group_name
-  location            = var.location
-  prefix              = var.prefix
-  deletion_protection = var.deletion_protection
+  source                       = "../../modules/byob"
+  resource_group_name          = var.resource_group_name
+  location                     = var.location
+  prefix                       = var.prefix
+  deletion_protection          = var.deletion_protection
+  create_cmk                   = var.create_cmk
+  enable_purge_protection      = var.enable_purge_protection
+  client_id                    = var.client_id
+  tenant_id                    = var.tenant_id
+  tags                         = var.tags
+  disable_storage_vault_key_id = var.disable_storage_vault_key_id
 }
 
-
 output "blob_container" {
   value = module.byob.blob_container
 }
@@ -19,3 +24,11 @@ output "storage_key" {
   value     = module.byob.azure_storage_key
   sensitive = true
 }
+
+output "storage_vault_key_id" {
+  value = module.byob.vault_key_id
+}
+
+output "command_to_get_storage_key" {
+  value = "az storage account keys list --account-name ${module.byob.storage_account_name} --resource-group ${module.byob.resource_group_name.name} --query '[0].value' -o tsv"
+}
diff --git a/examples/byob/terraform.tfvars b/examples/byob/terraform.tfvars
@@ -1,3 +1,16 @@
-resource_group_name = "rg-name"
-location            = "westeurope"
-prefix              = "byob-wandb"
+resource_group_name = "<rg-name>"
+
+location = "<westeurope>"
+prefix   = "<byob-wandb>"
+tags = {
+  "name" = "wandb"
+}
+
+#To enable Azure Key Vault encryption uncomment the below lines
+# create_cmk = true
+
+# enable_purge_protection       = true
+# disable_storage_vault_key_id  = false
+
+# tenant_id = "<tenant_id>"
+# client_id = "<client_id>"
diff --git a/examples/byob/variables.tf b/examples/byob/variables.tf
@@ -1,6 +1,6 @@
 variable "resource_group_name" {
-  description = "Resource Group name"
   type        = string
+  description = "Resource Group"
 }
 
 variable "location" {
@@ -18,3 +18,37 @@ variable "deletion_protection" {
   type        = bool
   default     = false
 }
+
+variable "create_cmk" {
+  type        = bool
+  default     = false
+  description = "Flag to create a Customer Managed Key for the Key Vault to encrypt the storage account."
+}
+
+variable "disable_storage_vault_key_id" {
+  type        = bool
+  default     = false
+  description = "Flag to disable the `customer_managed_key` block, the properties 'encryption.identity, encryption.keyvaultproperties' cannot be updated in a single operation."
+}
+
+variable "enable_purge_protection" {
+  type        = bool
+  default     = false
+  description = "Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled."
+}
+
+variable "tenant_id" {
+  type        = string
+  description = "The tenant ID for the Key Vault Access Policy. Get from W&B SA"
+}
+
+variable "client_id" {
+  type        = string
+  description = "The client ID (object id) for the Key Vault Access Policy. Get from W&B SA"
+}
+
+variable "tags" {
+  type        = map(string)
+  description = "Map of tags for resource"
+  default     = {}
+}
diff --git a/main.tf b/main.tf
@@ -2,6 +2,8 @@ locals {
   fqdn       = var.subdomain == null ? var.domain_name : "${var.subdomain}.${var.domain_name}"
   url_prefix = var.ssl ? "https" : "http"
   url        = "${local.url_prefix}://${local.fqdn}"
+
+  enable_internal_storage = var.blob_container == "" && var.external_bucket == null
 }
 
 resource "azurerm_resource_group" "default" {
@@ -27,23 +29,24 @@ module "networking" {
   private_link        = var.create_private_link
   allowed_ip_ranges   = var.allowed_ip_ranges
   tags                = var.tags
-
 }
 
 module "database" {
-  source              = "./modules/database"
-  namespace           = var.namespace
+  source    = "./modules/database"
+  namespace = var.namespace
+
   resource_group_name = azurerm_resource_group.default.name
   location            = azurerm_resource_group.default.location
 
   database_availability_mode   = var.database_availability_mode
   database_version             = var.database_version
   database_private_dns_zone_id = module.networking.database_private_dns_zone.id
   database_subnet_id           = module.networking.database_subnet.id
+  sku_name                     = try(local.deployment_size[var.size].db, var.database_sku_name)
+  deletion_protection          = var.deletion_protection
 
-  sku_name            = try(local.deployment_size[var.size].db, var.database_sku_name)
-  deletion_protection = var.deletion_protection
-
+  database_key_id = try(module.vault.vault_internal_keys[module.vault.vault_key_map.database].id, null)
+  identity_ids    = module.identity.identity.id
   tags = {
     "customer-ns" = var.namespace,
     "env"         = "managed-install"
@@ -64,23 +67,33 @@ module "redis" {
 module "vault" {
   source = "./modules/vault"
 
-  identity_object_id = module.identity.identity.principal_id
-  location           = azurerm_resource_group.default.location
-  namespace          = var.namespace
-  resource_group     = azurerm_resource_group.default
+  identity_object_id      = module.identity.identity.principal_id
+  location                = azurerm_resource_group.default.location
+  namespace               = var.namespace
+  resource_group          = azurerm_resource_group.default
+  enable_purge_protection = var.enable_purge_protection
+
+  enable_database_vault_key = var.enable_database_vault_key
+  enable_storage_vault_key  = var.enable_storage_vault_key && local.enable_internal_storage
 
   tags = var.tags
 }
 
 module "storage" {
-  count               = (var.blob_container == "" && var.external_bucket == null) ? 1 : 0
-  source              = "./modules/storage"
+  source = "./modules/storage"
+
+  count               = local.enable_internal_storage ? 1 : 0
   namespace           = var.namespace
   resource_group_name = azurerm_resource_group.default.name
   location            = azurerm_resource_group.default.location
   create_queue        = !var.use_internal_queue
+  identity_ids        = module.identity.identity.id
   deletion_protection = var.deletion_protection
-  tags                = var.tags
+
+  storage_key_id               = var.customer_storage_vault_key_id == null ? try(module.vault.vault_internal_keys[module.vault.vault_key_map.storage].id, null) : var.customer_storage_vault_key_id
+  disable_storage_vault_key_id = var.disable_storage_vault_key_id
+
+  tags = var.tags
 }
 
 module "app_lb" {
@@ -176,8 +189,9 @@ resource "azurerm_role_assignment" "gateway_role" {
 }
 
 module "cron_job" {
-  count                  = length(var.allowed_subscriptions) > 0 ? 1 : 0 # private endpoint approval application deployed as a cronjob in default namespace
-  source                 = "./modules/cron_job"
+  count  = length(var.allowed_subscriptions) > 0 ? 1 : 0 # private endpoint approval application deployed as a cronjob in default namespace
+  source = "./modules/cron_job"
+
   namespace              = "default"
   client_id              = module.pod_identity[0].identity.client_id
   serviceaccountName     = local.private_endpoint_approval_sa
@@ -186,15 +200,13 @@ module "cron_job" {
   applicationGatewayName = module.app_lb.gateway.name
   allowedSubscriptions   = var.allowed_subscriptions
   depends_on             = [module.app_lb, module.pod_identity]
-
 }
 
 resource "azurerm_role_assignment" "otel_role" {
   count                = var.azuremonitor ? 1 : 0
   scope                = azurerm_resource_group.default.id
   role_definition_name = "Contributor"
   principal_id         = module.identity.otel_identity.principal_id
-
 }
 
 resource "azurerm_federated_identity_credential" "otel_app" {
@@ -235,8 +247,8 @@ module "wandb" {
   spec = {
     values = {
       global = {
-        host    = local.url
-        license = var.license
+        host          = local.url
+        license       = var.license
         cloudProvider = "azure"
         bucket = var.external_bucket == null ? {
           provider  = "az"

diff --git a/modules/byob/main.tf b/modules/byob/main.tf
@@ -1,10 +1,47 @@
-module "storage" {
-  source = "../storage"
+module "identity" {
+  count  = var.create_cmk ? 1 : 0
+  source = "../identity"
+
+  namespace      = var.prefix
+  resource_group = { name = "${var.resource_group_name}", id = "byob" }
+  location       = var.location
+}
+
+module "vault" {
+  count          = var.create_cmk ? 1 : 0
+  source         = "../vault"
+  namespace      = var.prefix
+  resource_group = { name = "${var.resource_group_name}", id = "byob" }
+  location       = var.location
+
+  identity_object_id       = module.identity[0].identity.principal_id
+  depends_on               = [module.identity]
+  tags                     = var.tags
+  enable_purge_protection  = var.enable_purge_protection
+  enable_storage_vault_key = var.create_cmk
+}
+
+resource "azurerm_key_vault_access_policy" "wandb" {
+  count        = var.create_cmk ? 1 : 0
+  key_vault_id = module.vault[0].vault.id
+  tenant_id    = var.tenant_id
+  object_id    = var.client_id
 
+  key_permissions         = ["Get", "List", "Update", "Create", "Import", "Delete", "Recover", "Backup", "Restore"]
+  certificate_permissions = ["Get", "List", "Update", "Create", "Import", "Delete", "Recover", "Backup", "Restore"]
+  secret_permissions      = ["Get", "List", "Set", "Delete", "Recover", "Backup", "Restore"]
+}
+
+module "storage" {
+  source              = "../storage"
   create_queue        = false
   namespace           = var.prefix
   resource_group_name = var.resource_group_name
   location            = var.location
-
   deletion_protection = var.deletion_protection
+  storage_key_id      = try(module.vault[0].vault_internal_keys[module.vault[0].vault_key_map.storage].id, null)
+  identity_ids        = try(module.identity[0].identity.id, null)
+
+  disable_storage_vault_key_id = var.disable_storage_vault_key_id
 }
+
diff --git a/modules/byob/outputs.tf b/modules/byob/outputs.tf
@@ -2,6 +2,14 @@ output "azure_storage_key" {
   value = module.storage.account.primary_access_key
 }
 
+output "storage_key_id" {
+  value = try(module.vault[0].vault_internal_keys[module.vault[0].vault_key_map.storage].id, null)
+}
+
 output "blob_container" {
   value = "${module.storage.account.name}/${module.storage.container.name}"
-}
+}
+
+output "storage_account_name" {
+  value = module.storage.account.name
+}