From 0200e2aac73aebaf284f4998000ce475565a7b50 Mon Sep 17 00:00:00 2001
From: Aidan Goldthorpe <BlueDragon23@users.noreply.github.com>
Date: Mon, 6 Jun 2022 20:57:28 +1000
Subject: [PATCH] Bug fix for selecting KMS key via alias ARN (#5182)

---
 moto/kms/models.py               |  2 +-
 tests/test_kms/test_kms_boto3.py |  2 +-
 tests/test_kms/test_model.py     | 52 ++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_kms/test_model.py

diff --git a/moto/kms/models.py b/moto/kms/models.py
index b60b32dde..470b25f7a 100644
--- a/moto/kms/models.py
+++ b/moto/kms/models.py
@@ -302,7 +302,7 @@ class KmsBackend(BaseBackend):
     def get_alias_name(alias_name):
         # Allow use of ARN as well as alias name
         if alias_name.startswith("arn:") and ":alias/" in alias_name:
-            return alias_name.split(":alias/")[1]
+            return "alias/" + alias_name.split(":alias/")[1]
 
         return alias_name
 
diff --git a/tests/test_kms/test_kms_boto3.py b/tests/test_kms/test_kms_boto3.py
index cb03a525b..ff1006078 100644
--- a/tests/test_kms/test_kms_boto3.py
+++ b/tests/test_kms/test_kms_boto3.py
@@ -647,7 +647,7 @@ def test_generate_data_key_all_valid_key_ids(prefix, append_key_id):
     if append_key_id:
         target_id += key_id
 
-    client.generate_data_key(KeyId=key_id, NumberOfBytes=32)
+    client.generate_data_key(KeyId=target_id, NumberOfBytes=32)
 
 
 @mock_kms
diff --git a/tests/test_kms/test_model.py b/tests/test_kms/test_model.py
new file mode 100644
index 000000000..5d0ffc097
--- /dev/null
+++ b/tests/test_kms/test_model.py
@@ -0,0 +1,52 @@
+import pytest
+
+from moto.kms.models import KmsBackend
+
+PLAINTEXT = b"text"
+REGION = "us-east-1"
+
+
+@pytest.fixture
+def backend():
+    return KmsBackend(REGION)
+
+
+@pytest.fixture
+def key(backend):
+    return backend.create_key(
+        None, "ENCRYPT_DECRYPT", "SYMMETRIC_DEFAULT", "Test key", None, REGION
+    )
+
+
+def test_encrypt_key_id(backend, key):
+    ciphertext, arn = backend.encrypt(key.id, PLAINTEXT, {})
+
+    assert ciphertext is not None
+    assert arn == key.arn
+
+
+def test_encrypt_key_arn(backend, key):
+    ciphertext, arn = backend.encrypt(key.arn, PLAINTEXT, {})
+
+    assert ciphertext is not None
+    assert arn == key.arn
+
+
+def test_encrypt_alias_name(backend, key):
+    backend.add_alias(key.id, "alias/test/test")
+
+    ciphertext, arn = backend.encrypt("alias/test/test", PLAINTEXT, {})
+
+    assert ciphertext is not None
+    assert arn == key.arn
+
+
+def test_encrypt_alias_arn(backend, key):
+    backend.add_alias(key.id, "alias/test/test")
+
+    ciphertext, arn = backend.encrypt(
+        f"arn:aws:kms:{REGION}:{key.account_id}:alias/test/test", PLAINTEXT, {}
+    )
+
+    assert ciphertext is not None
+    assert arn == key.arn