From 9e9522cbf0b5488579624c68400634827f998fa1 Mon Sep 17 00:00:00 2001
From: Mark Chappell <mchappel@redhat.com>
Date: Sat, 23 Oct 2021 20:37:30 +0200
Subject: [PATCH 1/2] ec2_key - add support for tagging and returning id

---
 plugins/module_utils/tagging.py               |   5 +
 plugins/modules/ec2_key.py                    |  82 +++--
 .../targets/ec2_key/tasks/main.yml            | 290 +++++++++++++++++-
 3 files changed, 343 insertions(+), 34 deletions(-)

diff --git a/plugins/module_utils/tagging.py b/plugins/module_utils/tagging.py
index e5042725963..99252bc4ee4 100644
--- a/plugins/module_utils/tagging.py
+++ b/plugins/module_utils/tagging.py
@@ -94,6 +94,9 @@ def ansible_dict_to_boto3_tag_list(tags_dict, tag_name_key_name='Key', tag_value
         ]
     """
 
+    if not tags_dict:
+        return []
+
     tags_list = []
     for k, v in tags_dict.items():
         tags_list.append({tag_name_key_name: k, tag_value_key_name: to_native(v)})
@@ -127,6 +130,8 @@ def boto3_tag_specifications(tags_dict, types=None):
     Returns:
         List: List of dictionaries representing an AWS Tag Specification
     """
+    if not tags_dict:
+        return None
     specifications = list()
     tag_list = ansible_dict_to_boto3_tag_list(tags_dict)
 
diff --git a/plugins/modules/ec2_key.py b/plugins/modules/ec2_key.py
index e9e8660b7cf..de7f9fcbc0b 100644
--- a/plugins/modules/ec2_key.py
+++ b/plugins/modules/ec2_key.py
@@ -47,6 +47,17 @@
       - This option has no effect since version 2.5 and will be removed after 2022-06-01.
     type: int
     required: false
+  tags:
+    description:
+      - A dictionary of tags to set on the key pair.
+    type: dict
+    version_added: 2.1.0
+  purge_tags:
+    description:
+      - Delete any tags not specified in I(tags).
+    default: false
+    type: bool
+    version_added: 2.1.0
 
 extends_documentation_fragment:
 - amazon.aws.aws
@@ -114,6 +125,16 @@
       returned: when state is present
       type: str
       sample: my_keypair
+    id:
+      description: id of the keypair
+      returned: when state is present
+      type: str
+      sample: key-123456789abc
+    tags:
+      description: a dictionary representing the tags attached to the key pair
+      returned: when state is present
+      type: dict
+      sample: '{"my_key": "my value"}'
     private_key:
       description: private key of a newly created keypair
       returned: when a new keypair is created by AWS (key_material is not provided)
@@ -135,14 +156,21 @@
 from ..module_utils.core import AnsibleAWSModule
 from ..module_utils.core import is_boto3_error_code
 from ..module_utils.ec2 import AWSRetry
+from ..module_utils.ec2 import ensure_ec2_tags
+from ..module_utils.tagging import boto3_tag_specifications
+from ..module_utils.tagging import boto3_tag_list_to_ansible_dict
 
 
 def extract_key_data(key):
 
     data = {
         'name': key['KeyName'],
-        'fingerprint': key['KeyFingerprint']
+        'fingerprint': key['KeyFingerprint'],
+        'id': key['KeyPairId'],
+        'tags': {},
     }
+    if 'Tags' in key:
+        data['tags'] = boto3_tag_list_to_ansible_dict(key['Tags'])
     if 'KeyMaterial' in key:
         data['private_key'] = key['KeyMaterial']
     return data
@@ -162,7 +190,7 @@ def get_key_fingerprint(module, ec2_client, key_material):
         random_name = "ansible-" + str(uuid.uuid4())
         name_in_use = find_key_pair(module, ec2_client, random_name)
 
-    temp_key = import_key_pair(module, ec2_client, random_name, key_material)
+    temp_key = _import_key_pair(module, ec2_client, random_name, key_material)
     delete_key_pair(module, ec2_client, random_name, finish_task=False)
     return temp_key['KeyFingerprint']
 
@@ -182,40 +210,54 @@ def find_key_pair(module, ec2_client, name):
 
 def create_key_pair(module, ec2_client, name, key_material, force):
 
+    tags = module.params.get('tags')
+    purge_tags = module.params.get('purge_tags')
     key = find_key_pair(module, ec2_client, name)
+    tag_spec = boto3_tag_specifications(tags, ['key-pair'])
+    changed = False
     if key:
         if key_material and force:
-            if not module.check_mode:
-                new_fingerprint = get_key_fingerprint(module, ec2_client, key_material)
-                if key['KeyFingerprint'] != new_fingerprint:
+            new_fingerprint = get_key_fingerprint(module, ec2_client, key_material)
+            if key['KeyFingerprint'] != new_fingerprint:
+                changed = True
+                if not module.check_mode:
                     delete_key_pair(module, ec2_client, name, finish_task=False)
-                    key = import_key_pair(module, ec2_client, name, key_material)
-                    key_data = extract_key_data(key)
-                    module.exit_json(changed=True, key=key_data, msg="key pair updated")
-            else:
-                # Assume a change will be made in check mode since a comparison can't be done
-                module.exit_json(changed=True, key=extract_key_data(key), msg="key pair updated")
+                    key = _import_key_pair(module, ec2_client, name, key_material, tag_spec)
+                key_data = extract_key_data(key)
+                module.exit_json(changed=True, key=key_data, msg="key pair updated")
+        changed |= ensure_ec2_tags(ec2_client, module, key['KeyPairId'], tags=tags, purge_tags=purge_tags)
+        key = find_key_pair(module, ec2_client, name)
         key_data = extract_key_data(key)
-        module.exit_json(changed=False, key=key_data, msg="key pair already exists")
+        module.exit_json(changed=changed, key=key_data, msg="key pair already exists")
     else:
         # key doesn't exist, create it now
         key_data = None
         if not module.check_mode:
             if key_material:
-                key = import_key_pair(module, ec2_client, name, key_material)
+                key = _import_key_pair(module, ec2_client, name, key_material, tag_spec)
             else:
-                try:
-                    key = ec2_client.create_key_pair(aws_retry=True, KeyName=name)
-                except botocore.exceptions.ClientError as err:
-                    module.fail_json_aws(err, msg="error creating key")
+                key = _create_key_pair(module, ec2_client, name, tag_spec)
             key_data = extract_key_data(key)
         module.exit_json(changed=True, key=key_data, msg="key pair created")
 
 
-def import_key_pair(module, ec2_client, name, key_material):
+def _create_key_pair(module, ec2_client, name, tag_spec):
+    params = dict(KeyName=name)
+    if tag_spec:
+        params['TagSpecifications'] = tag_spec
+    try:
+        key = ec2_client.create_key_pair(aws_retry=True, **params)
+    except botocore.exceptions.ClientError as err:
+        module.fail_json_aws(err, msg="error creating key")
+    return key
+
 
+def _import_key_pair(module, ec2_client, name, key_material, tag_spec=None):
+    params = dict(KeyName=name, PublicKeyMaterial=to_bytes(key_material))
+    if tag_spec:
+        params['TagSpecifications'] = tag_spec
     try:
-        key = ec2_client.import_key_pair(aws_retry=True, KeyName=name, PublicKeyMaterial=to_bytes(key_material))
+        key = ec2_client.import_key_pair(aws_retry=True, **params)
     except botocore.exceptions.ClientError as err:
         module.fail_json_aws(err, msg="error importing key")
     return key
@@ -243,6 +285,8 @@ def main():
         key_material=dict(no_log=False),
         force=dict(type='bool', default=True),
         state=dict(default='present', choices=['present', 'absent']),
+        tags=dict(type='dict'),
+        purge_tags=dict(type='bool', default=False),
         wait=dict(type='bool', removed_at_date='2022-06-01', removed_from_collection='amazon.aws'),
         wait_timeout=dict(type='int', removed_at_date='2022-06-01', removed_from_collection='amazon.aws')
     )
diff --git a/tests/integration/targets/ec2_key/tasks/main.yml b/tests/integration/targets/ec2_key/tasks/main.yml
index 69e7edcb3f9..3f7816f2fce 100644
--- a/tests/integration/targets/ec2_key/tasks/main.yml
+++ b/tests/integration/targets/ec2_key/tasks/main.yml
@@ -24,6 +24,18 @@
            - 'result.msg == "missing required arguments: name"'
 
     # ============================================================
+    - name: test removing a non-existent key pair (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: absent
+      register: result
+      check_mode: true
+
+    - name: assert removing a non-existent key pair
+      assert:
+        that:
+          - 'not result.changed'
+
     - name: test removing a non-existent key pair
       ec2_key:
         name: '{{ ec2_key_name }}'
@@ -36,23 +48,269 @@
           - 'not result.changed'
 
     # ============================================================
+    - name: test creating a new key pair (check_mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          snake_case: 'a_snake_case_value'
+          CamelCase: 'CamelCaseValue'
+          "spaced key": 'Spaced value'
+      register: result
+      check_mode: true
+
+    - name: assert creating a new key pair
+      assert:
+        that:
+           - result is changed
+
     - name: test creating a new key pair
       ec2_key:
         name: '{{ ec2_key_name }}'
         state: present
+        tags:
+          snake_case: 'a_snake_case_value'
+          CamelCase: 'CamelCaseValue'
+          "spaced key": 'Spaced value'
       register: result
 
     - name: assert creating a new key pair
       assert:
         that:
-           - 'result.changed'
-           - '"key" in result'
-           - '"name" in result.key'
-           - '"fingerprint" in result.key'
-           - '"private_key" in result.key'
-           - 'result.key.name == "{{ec2_key_name}}"'
+          - result is changed
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" in result.key'
+          - '"id" in result.key'
+          - '"tags" in result.key'
+          - result.key.name == ec2_key_name
+          - result.key.id.startswith('key-')
+          - '"snake_case" in result.key.tags'
+          - result.key.tags['snake_case'] == 'a_snake_case_value'
+          - '"CamelCase" in result.key.tags'
+          - result.key.tags['CamelCase'] == 'CamelCaseValue'
+          - '"spaced key" in result.key.tags'
+          - result.key.tags['spaced key'] == 'Spaced value'
+
+    - set_fact:
+        key_id_1: '{{ result.key.id }}'
+
+    - name: 'test re-"creating" the same key (check_mode)'
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          snake_case: 'a_snake_case_value'
+          CamelCase: 'CamelCaseValue'
+          "spaced key": 'Spaced value'
+      register: result
+      check_mode: true
+
+    - name: assert re-creating the same key
+      assert:
+        that:
+           - result is not changed
+
+    - name: 'test re-"creating" the same key'
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          snake_case: 'a_snake_case_value'
+          CamelCase: 'CamelCaseValue'
+          "spaced key": 'Spaced value'
+      register: result
 
     # ============================================================
+    - name: test updating tags without purge (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: false
+      register: result
+      check_mode: true
+
+    - name: assert updated tags
+      assert:
+        that:
+           - result is changed
+
+    - name: test updating tags without purge
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: false
+      register: result
+
+    - name: assert updated tags
+      assert:
+        that:
+          - result is changed
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" not in result.key'
+          - '"id" in result.key'
+          - result.key.id == key_id_1
+          - '"tags" in result.key'
+          - result.key.name == ec2_key_name
+          - '"snake_case" in result.key.tags'
+          - result.key.tags['snake_case'] == 'a_snake_case_value'
+          - '"CamelCase" in result.key.tags'
+          - result.key.tags['CamelCase'] == 'CamelCaseValue'
+          - '"spaced key" in result.key.tags'
+          - result.key.tags['spaced key'] == 'Spaced value'
+          - '"newKey" in result.key.tags'
+          - result.key.tags['newKey'] == 'Another value'
+
+    - name: test updating tags without purge - idempotency (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: false
+      register: result
+      check_mode: true
+
+    - name: assert updated tags
+      assert:
+        that:
+           - result is not changed
+
+    - name: test updating tags without purge - idempotency
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: false
+      register: result
+
+    - name: assert updated tags
+      assert:
+        that:
+          - result is not changed
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" not in result.key'
+          - '"id" in result.key'
+          - '"tags" in result.key'
+          - result.key.name == ec2_key_name
+          - result.key.id == key_id_1
+          - '"snake_case" in result.key.tags'
+          - result.key.tags['snake_case'] == 'a_snake_case_value'
+          - '"CamelCase" in result.key.tags'
+          - result.key.tags['CamelCase'] == 'CamelCaseValue'
+          - '"spaced key" in result.key.tags'
+          - result.key.tags['spaced key'] == 'Spaced value'
+          - '"newKey" in result.key.tags'
+          - result.key.tags['newKey'] == 'Another value'
+
+    # ============================================================
+    - name: test updating tags with purge (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: true
+      register: result
+      check_mode: true
+
+    - name: assert updated tags
+      assert:
+        that:
+           - result is changed
+
+    - name: test updating tags with purge
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: true
+      register: result
+
+    - name: assert updated tags
+      assert:
+        that:
+          - result is changed
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" not in result.key'
+          - '"id" in result.key'
+          - result.key.id == key_id_1
+          - '"tags" in result.key'
+          - result.key.name == ec2_key_name
+          - '"snake_case" not in result.key.tags'
+          - '"CamelCase" not in result.key.tags'
+          - '"spaced key" not in result.key.tags'
+          - '"newKey" in result.key.tags'
+          - result.key.tags['newKey'] == 'Another value'
+
+    - name: test updating tags with purge - idempotency (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: true
+      register: result
+      check_mode: true
+
+    - name: assert updated tags
+      assert:
+        that:
+           - result is not changed
+
+    - name: test updating tags with purge - idempotency
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: present
+        tags:
+          newKey: 'Another value'
+        purge_tags: true
+      register: result
+
+    - name: assert updated tags
+      assert:
+        that:
+          - result is not changed
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" not in result.key'
+          - '"id" in result.key'
+          - '"tags" in result.key'
+          - result.key.name == ec2_key_name
+          - result.key.id == key_id_1
+          - '"snake_case" not in result.key.tags'
+          - '"CamelCase" not in result.key.tags'
+          - '"spaced key" not in result.key.tags'
+          - '"newKey" in result.key.tags'
+          - result.key.tags['newKey'] == 'Another value'
+
+    # ============================================================
+    - name: test removing an existent key (check mode)
+      ec2_key:
+        name: '{{ ec2_key_name }}'
+        state: absent
+      register: result
+      check_mode: true
+
+    - name: assert removing an existent key
+      assert:
+        that:
+           - result is changed
+
     - name: test removing an existent key
       ec2_key:
         name: '{{ ec2_key_name }}'
@@ -62,9 +320,9 @@
     - name: assert removing an existent key
       assert:
         that:
-           - 'result.changed'
+           - result is changed
            - '"key" in result'
-           - 'result.key == None'
+           - result.key == None
 
     # ============================================================
     - name: test state=present with key_material
@@ -77,13 +335,15 @@
     - name: assert state=present with key_material
       assert:
         that:
-           - 'result.changed == True'
-           - '"key" in result'
-           - '"name" in result.key'
-           - '"fingerprint" in result.key'
-           - '"private_key" not in result.key'
-           - 'result.key.name == "{{ec2_key_name}}"'
-           - 'result.key.fingerprint == "{{fingerprint}}"'
+          - 'result.changed == True'
+          - '"key" in result'
+          - '"name" in result.key'
+          - '"fingerprint" in result.key'
+          - '"private_key" not in result.key'
+          - '"id" in result.key'
+          - '"tags" in result.key'
+          - 'result.key.name == "{{ec2_key_name}}"'
+          - 'result.key.fingerprint == "{{fingerprint}}"'
 
     # ============================================================
 

From 4a34fad57af254baad0a7f147bb801391f80ff2e Mon Sep 17 00:00:00 2001
From: Mark Chappell <mchappel@redhat.com>
Date: Sat, 23 Oct 2021 20:55:17 +0200
Subject: [PATCH 2/2] changelog

---
 changelogs/fragments/548-ec2_key-tagging.yml | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 changelogs/fragments/548-ec2_key-tagging.yml

diff --git a/changelogs/fragments/548-ec2_key-tagging.yml b/changelogs/fragments/548-ec2_key-tagging.yml
new file mode 100644
index 00000000000..24b0382b295
--- /dev/null
+++ b/changelogs/fragments/548-ec2_key-tagging.yml
@@ -0,0 +1,2 @@
+minor_changes:
+- ec2_key - add support for tagging key pairs (https://github.com/ansible-collections/amazon.aws/pull/548).