From 79522357eb81a97881ea5e0c377c7ad8671c2933 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 12 Sep 2023 14:12:25 -0400
Subject: [PATCH] Add better error messages for KMS Key lookup on imported
 dataset

---
 .../modules/dataset_sharing/aws/kms_client.py   | 17 +++++++++++++++++
 .../datasets/services/dataset_service.py        | 11 ++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/backend/dataall/modules/dataset_sharing/aws/kms_client.py b/backend/dataall/modules/dataset_sharing/aws/kms_client.py
index 645df769a..5642a9013 100644
--- a/backend/dataall/modules/dataset_sharing/aws/kms_client.py
+++ b/backend/dataall/modules/dataset_sharing/aws/kms_client.py
@@ -52,3 +52,20 @@ def get_key_id(self, key_alias: str):
             return None
         else:
             return response['KeyMetadata']['KeyId']
+
+    def check_key_exists(self, key_alias: str):
+        try:
+            key_exist = False
+            paginator = self._client.get_paginator('list_aliases')
+            for page in paginator.paginate():
+                key_aliases = [alias["AliasName"] for alias in page['Aliases']]
+                if key_alias in key_aliases:
+                    key_exist = True
+                    break
+        except Exception as e:
+            log.error(
+                f'Failed to list kms key aliases in account {self._account_id}: {e}'
+            )
+            return None
+        else:
+            return key_exist
diff --git a/backend/dataall/modules/datasets/services/dataset_service.py b/backend/dataall/modules/datasets/services/dataset_service.py
index 4bf756b68..be6a4ed47 100644
--- a/backend/dataall/modules/datasets/services/dataset_service.py
+++ b/backend/dataall/modules/datasets/services/dataset_service.py
@@ -54,13 +54,22 @@ def check_dataset_account(session, environment):
     def check_imported_resources(environment, data):
         kms_alias = data.get('KmsKeyAlias')
         if kms_alias not in [None, "Undefined", "", "SSE-S3"]:
+            key_exists = KmsClient(environment.AwsAccountId, environment.region).check_key_exists(
+                key_alias=f"alias/{kms_alias}"
+            )
+            if not key_exists:
+                raise exceptions.AWSResourceNotFound(
+                    action=IMPORT_DATASET,
+                    message=f'KMS key with alias={kms_alias} cannot be found - Please check if KMS Key Alias exists in account {environment.AwsAccountId}',
+                )
+
             key_id = KmsClient(environment.AwsAccountId, environment.region).get_key_id(
                 key_alias=f"alias/{kms_alias}"
             )
             if not key_id:
                 raise exceptions.AWSResourceNotFound(
                     action=IMPORT_DATASET,
-                    message=f'KMS key with alias={kms_alias} cannot be found',
+                    message=f'Data.all Environment Pivot Role does not have kms:DescribeKey Permission to KMS key with alias={kms_alias}',
                 )
         return True