From ce69cdd1d5f3f9200dcf474670910406bcb9394d Mon Sep 17 00:00:00 2001
From: Zachary Blasczyk <77289967+zacharyblasczyk@users.noreply.github.com>
Date: Thu, 12 Sep 2024 11:17:20 -0500
Subject: [PATCH] fix: Allow the admin arn to be set for the KMS key (#277)

---
 main.tf                  | 2 ++
 modules/kms/main.tf      | 6 +++++-
 modules/kms/variables.tf | 6 ++++++
 variables.tf             | 6 ++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/main.tf b/main.tf
index aaa16a62..4600c08c 100644
--- a/main.tf
+++ b/main.tf
@@ -6,6 +6,8 @@ module "kms" {
   key_alias  = var.kms_key_alias == null ? "${var.namespace}-kms-alias" : var.kms_key_alias
   key_policy = var.kms_key_policy
 
+  policy_administrator_arn = var.kms_key_policy_administrator_arn
+
   create_clickhouse_key = var.enable_clickhouse
   clickhouse_key_alias  = var.kms_clickhouse_key_alias == null ? "${var.namespace}-kms-clickhouse-alias" : var.kms_clickhouse_key_alias
   clickhouse_key_policy = var.kms_clickhouse_key_policy
diff --git a/modules/kms/main.tf b/modules/kms/main.tf
index 4b0f6edd..bc0519da 100644
--- a/modules/kms/main.tf
+++ b/modules/kms/main.tf
@@ -1,5 +1,9 @@
 data "aws_caller_identity" "current" {}
 
+locals {
+  policy_administrator_arn = var.policy_administrator_arn != "" ? var.policy_administrator_arn : data.aws_caller_identity.current.arn
+}
+
 resource "aws_kms_key" "key" {
   deletion_window_in_days = var.key_deletion_window
   description             = "AWS KMS Customer-managed key to encrypt Weights & Biases resources"
@@ -11,7 +15,7 @@ resource "aws_kms_key" "key" {
       {
         "Sid" : "Allow administration of the key",
         "Effect" : "Allow",
-        "Principal" : { "AWS" : "${data.aws_caller_identity.current.arn}" },
+        "Principal" : { "AWS" : "${local.policy_administrator_arn}" },
         "Action" : "kms:*",
         "Resource" : "*"
       },
diff --git a/modules/kms/variables.tf b/modules/kms/variables.tf
index bcafbd3f..ced4039c 100644
--- a/modules/kms/variables.tf
+++ b/modules/kms/variables.tf
@@ -14,6 +14,12 @@ variable "iam_principal_arn" {
   default     = ""
 }
 
+variable "policy_administrator_arn" {
+  description = "The IAM principal (role or user) ARN that will be authorized to manage the key."
+  type        = string
+  default     = ""
+}
+
 variable "key_policy" {
   description = "The policy that will define the permissions for the kms key."
   type        = string
diff --git a/variables.tf b/variables.tf
index cd152ab3..10b2d793 100644
--- a/variables.tf
+++ b/variables.tf
@@ -199,6 +199,12 @@ variable "kms_key_policy" {
   default     = ""
 }
 
+variable "kms_key_policy_administrator_arn" {
+  type        = string
+  description = "The principal that will be allowed to manage the kms key."
+  default     = ""
+}
+
 variable "kms_clickhouse_key_alias" {
   type        = string
   description = "KMS key alias for AWS KMS Customer managed key used by Clickhouse CMEK."