Fix issues with aws_kms when working cross-account and with IDs (#60805)

* aws_kms: (integration tests) Test updating a key by ID rather than just my alias

* aws_kms: (integration tests) Test deletion of non-existent and keys that are already marked for deletion

* aws_kms: Ensure we can perform actions on a specific key_id rather than just aliases

In the process switch over to using get_key_details rather than listing all keys.

* aws_kms: When updating keys use the ARN rather than just the ID.

This is important when working with cross-account trusts.

3 files changed, 95 insertions, 44 deletions

diff --git a/changelogs/fragments/60805-aws_kms-ID-and-ARN-fixes.yml b/changelogs/fragments/60805-aws_kms-ID-and-ARN-fixes.yml
new file mode 100644
index 0000000000..08674ed275
--- /dev/null
+++ b/changelogs/fragments/60805-aws_kms-ID-and-ARN-fixes.yml
@@ -0,0 +1,3 @@
+bugfixes:
+- aws_kms - fix exception when only Key ID is passed
+- aws_kms - Use ARN rather than ID so that cross-account changes function
diff --git a/lib/ansible/modules/cloud/amazon/aws_kms.py b/lib/ansible/modules/cloud/amazon/aws_kms.py
index 03a6cc9392..1b4c4bef38 100644
--- a/lib/ansible/modules/cloud/amazon/aws_kms.py
+++ b/lib/ansible/modules/cloud/amazon/aws_kms.py
@@ -137,6 +137,7 @@ options:
 author:
   - Ted Timmons (@tedder)
   - Will Thames (@willthames)
+  - Mark Chappell (@tremble)
 extends_documentation_fragment:
 - aws
 - ec2
@@ -367,7 +368,9 @@ from ansible.module_utils.ec2 import AWSRetry, camel_dict_to_snake_dict
 from ansible.module_utils.ec2 import boto3_tag_list_to_ansible_dict, ansible_dict_to_boto3_tag_list
 from ansible.module_utils.ec2 import compare_aws_tags, compare_policies
 from ansible.module_utils.six import string_types
+from ansible.module_utils._text import to_native
 
+import traceback
 import json
 
 try:
@@ -536,7 +539,7 @@ def get_kms_facts(connection, module):
 
 
 def convert_grant_params(grant, key):
-    grant_params = dict(KeyId=key['key_id'],
+    grant_params = dict(KeyId=key['key_arn'],
                         GranteePrincipal=grant['grantee_principal'])
     if grant.get('operations'):
         grant_params['Operations'] = grant['operations']
@@ -594,39 +597,46 @@ def ensure_enabled_disabled(connection, module, key):
     changed = False
     if key['key_state'] == 'Disabled' and module.params['enabled']:
         try:
-            connection.enable_key(KeyId=key['key_id'])
+            connection.enable_key(KeyId=key['key_arn'])
             changed = True
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             module.fail_json_aws(e, msg="Failed to enable key")
 
     if key['key_state'] == 'Enabled' and not module.params['enabled']:
         try:
-            connection.disable_key(KeyId=key['key_id'])
+            connection.disable_key(KeyId=key['key_arn'])
             changed = True
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             module.fail_json_aws(e, msg="Failed to disable key")
     return changed
 
 
-def update_key(connection, module, key):
-    changed = False
-    alias = module.params['alias']
+def update_alias(connection, module, key_id, alias):
     if not alias.startswith('alias/'):
         alias = 'alias/' + alias
     aliases = get_kms_aliases_with_backoff(connection)['Aliases']
-    key_id = module.params.get('key_id')
     if key_id:
         # We will only add new aliases, not rename existing ones
         if alias not in [_alias['AliasName'] for _alias in aliases]:
             try:
                 connection.create_alias(KeyId=key_id, AliasName=alias)
-                changed = True
+                return True
             except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
                 module.fail_json_aws(e, msg="Failed create key alias")
+    return False
+
+
+def update_key(connection, module, key):
+    changed = False
+    alias = module.params.get('alias')
+    key_id = key['key_arn']
+
+    if alias:
+        changed = update_alias(connection, module, key_id, alias) or changed
 
     if key['key_state'] == 'PendingDeletion':
         try:
-            connection.cancel_key_deletion(KeyId=key['key_id'])
+            connection.cancel_key_deletion(KeyId=key_id)
             # key is disabled after deletion cancellation
             # set this so that ensure_enabled_disabled works correctly
             key['key_state'] = 'Disabled'
@@ -641,7 +651,7 @@ def update_key(connection, module, key):
     # (means you can't remove a description completely)
     if description and key['description'] != description:
         try:
-            connection.update_key_description(KeyId=key['key_id'], Description=description)
+            connection.update_key_description(KeyId=key_id, Description=description)
             changed = True
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             module.fail_json_aws(e, msg="Failed to update key description")
@@ -651,13 +661,13 @@ def update_key(connection, module, key):
                                          module.params.get('purge_tags'))
     if to_remove:
         try:
-            connection.untag_resource(KeyId=key['key_id'], TagKeys=to_remove)
+            connection.untag_resource(KeyId=key_id, TagKeys=to_remove)
             changed = True
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             module.fail_json_aws(e, msg="Unable to remove or update tag")
     if to_add:
         try:
-            connection.tag_resource(KeyId=key['key_id'],
+            connection.tag_resource(KeyId=key_id,
                                     Tags=[{'TagKey': tag_key, 'TagValue': desired_tags[tag_key]}
                                           for tag_key in to_add])
             changed = True
@@ -668,7 +678,7 @@ def update_key(connection, module, key):
     if module.params.get('policy'):
         policy = module.params.get('policy')
         try:
-            keyret = connection.get_key_policy(KeyId=key['key_id'], PolicyName='default')
+            keyret = connection.get_key_policy(KeyId=key_id, PolicyName='default')
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             # If we can't fetch the current policy assume we're making a change
             # Could occur if we have PutKeyPolicy without GetKeyPolicy
@@ -682,7 +692,7 @@ def update_key(connection, module, key):
             changed = True
             if not module.check_mode:
                 try:
-                    connection.put_key_policy(KeyId=key['key_id'], PolicyName='default', Policy=policy)
+                    connection.put_key_policy(KeyId=key_id, PolicyName='default', Policy=policy)
                 except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
                     module.fail_json_aws(e, msg="Unable to update key policy")
 
@@ -694,7 +704,7 @@ def update_key(connection, module, key):
     if to_remove:
         for grant in to_remove:
             try:
-                connection.retire_grant(KeyId=key['key_arn'], GrantId=grant['grant_id'])
+                connection.retire_grant(KeyId=key_id, GrantId=grant['grant_id'])
                 changed = True
             except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
                 module.fail_json_aws(e, msg="Unable to retire grant")
@@ -708,8 +718,8 @@ def update_key(connection, module, key):
             except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
                 module.fail_json_aws(e, msg="Unable to create grant")
 
-    # make results consistent with kms_facts
-    result = get_key_details(connection, module, key['key_id'])
+    # make results consistent with kms_facts before returning
+    result = get_key_details(connection, module, key_id)
     module.exit_json(changed=changed, **result)
 
 
@@ -750,17 +760,17 @@ def create_key(connection, module):
     module.exit_json(changed=True, **result)
 
 
-def delete_key(connection, module, key):
+def delete_key(connection, module, key_metadata, key_id):
     changed = False
 
-    if key['key_state'] != 'PendingDeletion':
+    if key_metadata['KeyState'] != 'PendingDeletion':
         try:
-            connection.schedule_key_deletion(KeyId=key['key_id'])
+            connection.schedule_key_deletion(KeyId=key_id)
             changed = True
         except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
             module.fail_json_aws(e, msg="Failed to schedule key for deletion")
 
-    result = get_key_details(connection, module, key['key_id'])
+    result = get_key_details(connection, module, key_id)
     module.exit_json(changed=changed, **result)
 
 
@@ -790,9 +800,9 @@ def get_arn_from_role_name(iam, rolename):
     raise Exception('could not find arn for name {0}.'.format(rolename))
 
 
-def do_policy_grant(kms, keyarn, role_arn, granttypes, mode='grant', dry_run=True, clean_invalid_entries=True):
+def do_policy_grant(module, kms, keyarn, role_arn, granttypes, mode='grant', dry_run=True, clean_invalid_entries=True):
     ret = {}
-    keyret = kms.get_key_policy(KeyId=keyarn, PolicyName='default')
+    keyret = get_key_policy_with_backoff(kms, keyarn, 'default')
     policy = json.loads(keyret['Policy'])
 
     changes_needed = {}
@@ -847,7 +857,8 @@ def do_policy_grant(kms, keyarn, role_arn, granttypes, mode='grant', dry_run=Tru
             kms.put_key_policy(KeyId=keyarn, PolicyName='default', Policy=policy_json_string)
             # returns nothing, so we have to just assume it didn't throw
             ret['changed'] = True
-    except Exception:
+    except Exception as e:
+        module.fail_json(msg='Could not update key_policy', new_policy=policy_json_string, details=to_native(e), exception=traceback.format_exc())
         raise
 
     ret['changes_needed'] = changes_needed
@@ -915,17 +926,30 @@ def main():
     kms = module.client('kms')
     iam = module.client('iam')
 
-    all_keys = get_kms_facts(kms, module)
     key_id = module.params.get('key_id')
     alias = module.params.get('alias')
-    if alias.startswith('alias/'):
+    if alias and alias.startswith('alias/'):
         alias = alias[6:]
-    if key_id:
-        filtr = ('key-id', key_id)
-    elif module.params.get('alias'):
-        filtr = ('alias', alias)
 
-    candidate_keys = [key for key in all_keys if key_matches_filter(key, filtr)]
+    # Fetch/Canonicalize key_id where possible
+    if key_id:
+        try:
+            # Don't use get_key_details it triggers module.fail when the key
+            # doesn't exist
+            key_metadata = get_kms_metadata_with_backoff(kms, key_id)['KeyMetadata']
+            key_id = key_metadata['Arn']
+        except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
+            # We can't create keys with a specific ID, if we can't access the
+            # key we'll have to fail
+            if module.params.get('state') == 'present':
+                module.fail_json(msg="Could not find key with id %s to update")
+            key_metadata = None
+    elif alias:
+        try:
+            key_metadata = get_kms_metadata_with_backoff(kms, 'alias/%s' % alias)['KeyMetadata']
+            key_id = key_metadata['Arn']
+        except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e:
+            key_metadata = None
 
     if module.params.get('policy_grant_types') or mode == 'deny':
         module.deprecate('Managing the KMS IAM Policy via policy_mode and policy_grant_types is fragile'
@@ -941,8 +965,8 @@ def main():
                 if g not in statement_label:
                     module.fail_json(msg='{0} is an unknown grant type.'.format(g))
 
-        ret = do_policy_grant(kms,
-                              candidate_keys[0]['key_arn'],
+        ret = do_policy_grant(module, kms,
+                              key_id,
                               module.params['policy_role_arn'],
                               module.params['policy_grant_types'],
                               mode=mode,
@@ -951,19 +975,20 @@ def main():
         result.update(ret)
 
         module.exit_json(**result)
-    else:
 
+    else:
         if module.params.get('state') == 'present':
-            if candidate_keys:
-                update_key(kms, module, candidate_keys[0])
+            if key_metadata:
+                key_details = get_key_details(kms, module, key_id)
+                update_key(kms, module, key_details)
             else:
-                if module.params.get('key_id'):
-                    module.fail_json(msg="Could not find key with id %s to update")
+                if key_id:
+                    module.fail_json(msg="Could not find key with id %s to update" % key_id)
                 else:
                     create_key(kms, module)
         else:
-            if candidate_keys:
-                delete_key(kms, module, candidate_keys[0])
+            if key_metadata:
+                delete_key(kms, module, key_metadata, key_id)
             else:
                 module.exit_json(changed=False)
 
diff --git a/test/integration/targets/aws_kms/tasks/main.yml b/test/integration/targets/aws_kms/tasks/main.yml
index 2e8643a50e..7bada3eef0 100644
--- a/test/integration/targets/aws_kms/tasks/main.yml
+++ b/test/integration/targets/aws_kms/tasks/main.yml
@@ -57,7 +57,7 @@
 
     - name: Update Policy on key to match AWS Console generate policy
       aws_kms:
-        key_alias: "alias/{{ resource_prefix }}-kms"
+        key_id: '{{ new_key["keys"][0]["key_id"] }}'
         policy: "{{ lookup('template', 'console-policy.j2') | to_json }}"
       register: kms_policy_changed
 
@@ -68,7 +68,7 @@
 
     - name: Attempt to re-assert the same policy
       aws_kms:
-        key_alias: "alias/{{ resource_prefix }}-kms"
+        alias: "alias/{{ resource_prefix }}-kms"
         policy: "{{ lookup('template', 'console-policy.j2') | to_json }}"
       register: kms_policy_changed
 
@@ -80,7 +80,7 @@
     - name: grant user-style access to production secrets
       aws_kms:
         mode: grant
-        key_alias: "alias/{{ resource_prefix }}-kms"
+        alias: "alias/{{ resource_prefix }}-kms"
         role_name: "{{ resource_prefix }}-kms-role"
         grant_types: "role,role grant"
 
@@ -93,7 +93,7 @@
     - name: remove access to production secrets from role
       aws_kms:
         mode: deny
-        key_alias: "alias/{{ resource_prefix }}-kms"
+        alias: "alias/{{ resource_prefix }}-kms"
         role_arn: "{{ iam_role_result.iam_role.arn }}"
 
     - name: find facts about the key
@@ -300,6 +300,18 @@
           - delete_kms.key_state == "PendingDeletion"
           - delete_kms.changed
 
+    - name: re-delete the key
+      aws_kms:
+        alias: "{{ resource_prefix }}-kms"
+        state: absent
+      register: delete_kms
+
+    - name: assert that state is pending deletion
+      assert:
+        that:
+          - delete_kms.key_state == "PendingDeletion"
+          - delete_kms is not changed
+
     - name: undelete and enable the key
       aws_kms:
         alias: "{{ resource_prefix }}-kms"
@@ -313,6 +325,17 @@
           - undelete_kms.key_state == "Enabled"
           - undelete_kms.changed
 
+    - name: delete a non-existant key
+      aws_kms:
+        key_id: '00000000-0000-0000-0000-000000000000'
+        state: absent
+      register: delete_kms
+
+    - name: assert that state is unchanged
+      assert:
+        that:
+          - delete_kms is not changed
+
   always:
     # ============================================================
     #   CLEAN-UP