From 8e7fb779c77e2782f016f9edd394d694d04143fa Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:18:39 +0200
Subject: [PATCH 01/21] fix typo, add imds support attribute

---
 plugins/modules/ec2_ami.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index 2213e93af5c..78eadfe55c1 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -553,12 +553,15 @@ def get_image_by_id(connection, image_id):
 
     result = images[0]
     try:
-        image_attribue = describe_image_attribute(connection, attribute="launchPermission", image_id=image_id)
-        if image_attribue:
-            result["LaunchPermissions"] = image_attribue["LaunchPermissions"]
-        image_attribue = describe_image_attribute(connection, attribute="productCodes", image_id=image_id)
-        if image_attribue:
-            result["ProductCodes"] = image_attribue["ProductCodes"]
+        image_attribute = describe_image_attribute(connection, attribute="launchPermission", image_id=image_id)
+        if image_attribute:
+            result["LaunchPermissions"] = image_attribute["LaunchPermissions"]
+        image_attribute = describe_image_attribute(connection, attribute="productCodes", image_id=image_id)
+        if image_attribute:
+            result["ProductCodes"] = image_attribute["ProductCodes"]
+        image_attribute = describe_image_attribute(connection, attribute="imdsSupport", image_id=image_id)
+        if image_attribute:
+            result["ImdsSupport"] = image_attribute["Value"]
     except AnsibleEC2Error as e:
         raise Ec2AmiFailure(f"Error retrieving image attributes for image {image_id}", e)
     return result

From 2e7ce11b8afa2e562af62e7b677445af0a428217 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:19:21 +0200
Subject: [PATCH 02/21] ec2_ami: add option to enable imdsv2

---
 plugins/modules/ec2_ami.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index 78eadfe55c1..6ca1c8d5fa7 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -186,6 +186,12 @@
       - See the AWS documentation for more detail U(https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/uefi-secure-boot.html).
     type: str
     version_added: 5.5.0
+  imdsv2_enable:
+    description:
+      - Force IMDS v2 on the AMI
+      - See the AWS documentation for more detail U(https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances-ami-configuration).
+    type: bool
+    default: false
 author:
   - "Evan Duffield (@scicoin-project) <eduffield@iacquire.com>"
   - "Constantin Bugneac (@Constantin07) <constantin.bugneac@endava.com>"
@@ -769,6 +775,21 @@ def set_description(connection, module, image, description):
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting description for image {image['ImageId']}", e)
 
+    @staticmethod
+    def set_imdsv2(connection, image, imdsv2_enable):
+        if not imdsv2_enable and image["ImdsSupport"] != "v2.0":
+            return False
+
+        if image["ImdsSupport"] == "v2.0":
+            return False
+
+        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": "v2.0"}
+        try:
+            modify_image_attribute(connection, image_id=image["imageId"], **imds_params)
+            return True
+        except AnsibleEC2Error as e:
+            raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image['imageId']}", e)
+
     @classmethod
     def do(cls, module, connection, image_id):
         """Entry point to update an image"""
@@ -785,6 +806,7 @@ def do(cls, module, connection, image_id):
         changed |= cls.set_launch_permission(connection, image, launch_permissions, module.check_mode)
         changed |= cls.set_tags(connection, module, image_id, module.params["tags"], module.params["purge_tags"])
         changed |= cls.set_description(connection, module, image, module.params["description"])
+        changed |= cls.set_imdsv2(connection, image, module.params["imdsv2_enable"])
 
         if changed and module.check_mode:
             module.exit_json(changed=True, msg="Would have updated AMI if not in check mode.")
@@ -847,6 +869,14 @@ def set_launch_permissions(connection, launch_permissions, image_id):
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting launch permissions for image {image_id}", e)
 
+    @staticmethod
+    def set_imdsv2(connection, image_id):
+        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": "v2.0"}
+        try:
+            modify_image_attribute(connection, image_id=image_id, **imds_params)
+        except AnsibleEC2Error as e:
+            raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image_id}", e)
+
     @staticmethod
     def create_or_register(create_image_parameters):
         create_from_instance = "InstanceId" in create_image_parameters
@@ -955,6 +985,8 @@ def do(cls, module, connection, _image_id):
             CreateImage.set_tags(connection, module, module.params.get("tags"), image_id)
 
         cls.set_launch_permissions(connection, module.params.get("launch_permissions"), image_id)
+        if module.params.get("imdsv2_enable"):
+            cls.set_imdsv2(connection, image_id)
 
         module.exit_json(
             msg="AMI creation operation complete.", changed=True, **get_ami_info(get_image_by_id(connection, image_id))
@@ -1005,6 +1037,7 @@ def main():
         tpm_support={"type": "str"},
         uefi_data={"type": "str"},
         virtualization_type={"default": "hvm"},
+        imdsv2_enable={"type": "bool", "default": False},
         wait={"type": "bool", "default": False},
         wait_timeout={"default": 1200, "type": "int"},
     )

From 3893d3d002d042d1d2189d6e296b5481d3d8a48b Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:20:03 +0200
Subject: [PATCH 03/21] add test for imdsv2 enable

---
 .../targets/ec2_ami/tasks/main.yml            | 29 +++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/tests/integration/targets/ec2_ami/tasks/main.yml b/tests/integration/targets/ec2_ami/tasks/main.yml
index 267e52abb8d..7eaf2f825fe 100644
--- a/tests/integration/targets/ec2_ami/tasks/main.yml
+++ b/tests/integration/targets/ec2_ami/tasks/main.yml
@@ -708,8 +708,12 @@
         tags:
           Name: "{{ ec2_ami_name }}_permissions"
         launch_permissions:
-          org_arns: ["arn:aws:organizations::123456789012:organization/o-123ab4cdef"]
-          org_unit_arns: ["arn:aws:organizations::123456789012:ou/o-123example/ou-1234-5exampld"]
+          org_arns:
+            ["arn:aws:organizations::123456789012:organization/o-123ab4cdef"]
+          org_unit_arns:
+            [
+              "arn:aws:organizations::123456789012:ou/o-123example/ou-1234-5exampld",
+            ]
       register: permissions_update_result
 
     - name: Get ami info
@@ -727,6 +731,27 @@
           - "'organizational_unit_arn' in permissions_info_result.images[0].launch_permissions[1]"
           - permissions_info_result.images[0].launch_permissions[1]['organizational_unit_arn'] == 'arn:aws:organizations::123456789012:ou/o-123example/ou-1234-5exampld'
 
+    - name: Create image with imdsv2 enabled
+      amazon.aws.ec2_ami:
+        state: present
+        instance_id: "{{ ec2_instance_id }}"
+        name: "{{ ec2_ami_name }}_imdsv2"
+        tags:
+          Name: "{{ ec2_ami_name }}_imdsv2"
+        imdsv2_enable: true
+      register: imdsv2_image
+
+    - name: Get ami info
+      amazon.aws.ec2_ami_info:
+        image_ids: "{{ imdsv2_image.image_id }}"
+        describe_image_attributes: true
+      register: imdsv2_info
+
+    - name: Assert ami IMDSV2 is set
+      ansible.builtin.assert:
+        that:
+          - imdsv2_info.images[0].imdsv2_enabled
+
   # ============================================================
 
   always:

From 17d7b03b3d44d2978007e383c42160de06a3cd1b Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:20:17 +0200
Subject: [PATCH 04/21] ec2_ami_info: return imdsv2_enabled attribute

---
 plugins/modules/ec2_ami_info.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/plugins/modules/ec2_ami_info.py b/plugins/modules/ec2_ami_info.py
index 92e59679a06..d277bb74f51 100644
--- a/plugins/modules/ec2_ami_info.py
+++ b/plugins/modules/ec2_ami_info.py
@@ -152,6 +152,10 @@
             description: An AWS account ID with permissions to launch the AMI.
             type: str
       sample: [{"group": "all"}, {"user_id": "123456789012"}]
+    imdsv2_enabled:
+      description: Whether the image has IMDSv2 enabled or not.
+      returned: When O(describe_image_attributes=true).
+      type: bool
     name:
       description: The name of the AMI that was provided during image creation.
       returned: always
@@ -277,6 +281,10 @@ def list_ec2_images(ec2_client, module, request_args):
                     ec2_client, attribute="launchPermission", image_id=image["image_id"]
                 ).get("LaunchPermissions", [])
                 image["launch_permissions"] = [camel_dict_to_snake_dict(perm) for perm in launch_permissions]
+                imdsv2_enabled = describe_image_attribute(
+                    ec2_client, attribute="ImdsSupport", image_id=image["image_id"]
+                ).get("ImdsSupport", "")
+                image["imdsv2_enabled"] = True if imdsv2_enabled == "v2.0" else False
         except is_ansible_aws_error_code("AuthFailure"):
             # describing launch permissions of images owned by others is not permitted, but shouldn't cause failures
             pass

From 778808d41bc7fd2e1af0f8cf7e2b9a50b181f0e2 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:20:37 +0200
Subject: [PATCH 05/21] add imdsv2 changelog

---
 changelogs/fragments/add_imdsv2.yml | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 changelogs/fragments/add_imdsv2.yml

diff --git a/changelogs/fragments/add_imdsv2.yml b/changelogs/fragments/add_imdsv2.yml
new file mode 100644
index 00000000000..319b914deb3
--- /dev/null
+++ b/changelogs/fragments/add_imdsv2.yml
@@ -0,0 +1,4 @@
+---
+minor_changes:
+  - ec2_ami - adds ``imdsv2_enable`` parameter to enable IMDSv2 when creating/updating an image.
+  - ec2_ami_info - add ``imdsv2_enable`` return when ``describe_image_attributes`` is set.

From d9d386a4bd77f8a3305a863e54cf41898cfba431 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 09:43:07 +0200
Subject: [PATCH 06/21] split aws docs link on new line

---
 plugins/modules/ec2_ami.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index 6ca1c8d5fa7..5f09ddf3363 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -189,7 +189,8 @@
   imdsv2_enable:
     description:
       - Force IMDS v2 on the AMI
-      - See the AWS documentation for more detail U(https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances-ami-configuration).
+      - See the AWS documentation for more detail
+      - U(https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances-ami-configuration).
     type: bool
     default: false
 author:

From 204636f8eac634d591744ec224a1b04cb2c46496 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <andrea@braingap.uk>
Date: Fri, 27 Sep 2024 10:04:46 +0200
Subject: [PATCH 07/21] update unit tests to account for ImdsSupport calls

---
 tests/unit/plugins/modules/test_ec2_ami.py      | 2 +-
 tests/unit/plugins/modules/test_ec2_ami_info.py | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/unit/plugins/modules/test_ec2_ami.py b/tests/unit/plugins/modules/test_ec2_ami.py
index dfec225e240..bc5e9a65b86 100644
--- a/tests/unit/plugins/modules/test_ec2_ami.py
+++ b/tests/unit/plugins/modules/test_ec2_ami.py
@@ -63,7 +63,7 @@ def test_get_image_by_id_found(m_describe_images, m_describe_image_attribute):
     image = ec2_ami.get_image_by_id(connection, "ami-0c7a795306730b288")
     assert image["ImageId"] == "ami-0c7a795306730b288"
     assert m_describe_images.call_count == 1
-    assert m_describe_image_attribute.call_count == 2
+    assert m_describe_image_attribute.call_count == 3
     m_describe_images.assert_has_calls(
         [
             call(
diff --git a/tests/unit/plugins/modules/test_ec2_ami_info.py b/tests/unit/plugins/modules/test_ec2_ami_info.py
index 087c55077cd..12c7e175acd 100644
--- a/tests/unit/plugins/modules/test_ec2_ami_info.py
+++ b/tests/unit/plugins/modules/test_ec2_ami_info.py
@@ -160,11 +160,15 @@ def test_list_ec2_images(m_get_images, m_describe_image_attribute):
     assert m_get_images.call_count == 1
     m_get_images.assert_called_with(ec2_client, request_args)
 
-    assert m_describe_image_attribute.call_count == 2
+    assert m_describe_image_attribute.call_count == 4
     m_describe_image_attribute.assert_has_calls(
         [call(ec2_client, attribute="launchPermission", image_id=images[0]["image_id"])],
         [call(ec2_client, attribute="launchPermission", image_id=images[1]["image_id"])],
     )
+    m_describe_image_attribute.assert_has_calls(
+        [call(ec2_client, attribute="ImdsSupport", image_id=images[0]["image_id"])],
+        [call(ec2_client, attribute="ImdsSupport", image_id=images[1]["image_id"])],
+    )
 
     assert len(list_ec2_images_result) == 2
     assert list_ec2_images_result[0]["image_id"] == "ami-1234567890"

From 2fdc65883493728bd50519b7f333df00ff1b8959 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Fri, 27 Sep 2024 12:22:30 +0200
Subject: [PATCH 08/21] add `version_added` for imdsv2

---
 plugins/modules/ec2_ami.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index 5f09ddf3363..c0e9d8f48a2 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -193,6 +193,7 @@
       - U(https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-IMDS-new-instances.html#configure-IMDS-new-instances-ami-configuration).
     type: bool
     default: false
+    version_added: 9.0.0
 author:
   - "Evan Duffield (@scicoin-project) <eduffield@iacquire.com>"
   - "Constantin Bugneac (@Constantin07) <constantin.bugneac@endava.com>"

From f250befcd7962da31b222ea4e1585f992aaffe85 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Fri, 4 Oct 2024 11:34:22 +0200
Subject: [PATCH 09/21] fix imds_params values

---
 plugins/modules/ec2_ami.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index c0e9d8f48a2..3bc80835163 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -785,7 +785,7 @@ def set_imdsv2(connection, image, imdsv2_enable):
         if image["ImdsSupport"] == "v2.0":
             return False
 
-        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": "v2.0"}
+        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": {"Value": "v2.0"}}
         try:
             modify_image_attribute(connection, image_id=image["imageId"], **imds_params)
             return True
@@ -873,7 +873,7 @@ def set_launch_permissions(connection, launch_permissions, image_id):
 
     @staticmethod
     def set_imdsv2(connection, image_id):
-        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": "v2.0"}
+        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": {"Value": "v2.0"}}
         try:
             modify_image_attribute(connection, image_id=image_id, **imds_params)
         except AnsibleEC2Error as e:

From 182c871990a3188922b5fcb980ea54e630730853 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Fri, 4 Oct 2024 11:36:18 +0200
Subject: [PATCH 10/21] add PR to changelog

---
 changelogs/fragments/add_imdsv2.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/changelogs/fragments/add_imdsv2.yml b/changelogs/fragments/add_imdsv2.yml
index 319b914deb3..517aee4548e 100644
--- a/changelogs/fragments/add_imdsv2.yml
+++ b/changelogs/fragments/add_imdsv2.yml
@@ -1,4 +1,4 @@
 ---
 minor_changes:
-  - ec2_ami - adds ``imdsv2_enable`` parameter to enable IMDSv2 when creating/updating an image.
-  - ec2_ami_info - add ``imdsv2_enable`` return when ``describe_image_attributes`` is set.
+  - ec2_ami - adds ``imdsv2_enable`` parameter to enable IMDSv2 when creating/updating an image (https://github.com/ansible-collections/amazon.aws/pull/2310).
+  - ec2_ami_info - add ``imdsv2_enable`` return when ``describe_image_attributes`` is set (https://github.com/ansible-collections/amazon.aws/pull/2310).

From 31f6e2d187f2653197a0fa113997c0e4c669a1e2 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Wed, 9 Oct 2024 12:25:52 +0200
Subject: [PATCH 11/21] correct imds support attribute value get

---
 plugins/modules/ec2_ami.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index 3bc80835163..d2c7906714b 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -569,7 +569,7 @@ def get_image_by_id(connection, image_id):
             result["ProductCodes"] = image_attribute["ProductCodes"]
         image_attribute = describe_image_attribute(connection, attribute="imdsSupport", image_id=image_id)
         if image_attribute:
-            result["ImdsSupport"] = image_attribute["Value"]
+            result["ImdsSupport"] = image_attribute["ImdsSupport"]["Value"]
     except AnsibleEC2Error as e:
         raise Ec2AmiFailure(f"Error retrieving image attributes for image {image_id}", e)
     return result

From 7a51457bd9c2991ad3c26fab3cd06045e3b4d7b8 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Wed, 9 Oct 2024 12:26:42 +0200
Subject: [PATCH 12/21] retry applying attribute when image is unavailable

---
 plugins/module_utils/ec2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/plugins/module_utils/ec2.py b/plugins/module_utils/ec2.py
index 9690e0d5cc8..05625c8dbc1 100644
--- a/plugins/module_utils/ec2.py
+++ b/plugins/module_utils/ec2.py
@@ -860,7 +860,7 @@ def deregister_image(client, image_id: str) -> bool:
 
 
 @EC2ImageErrorHandler.common_error_handler("modify image attribute")
-@AWSRetry.jittered_backoff()
+@AWSRetry.jittered_backoff(catch_extra_error_codes=["InvalidAMIID.Unavailable"])
 def modify_image_attribute(client, image_id: str, **params: Dict[str, Any]) -> bool:
     client.modify_image_attribute(ImageId=image_id, **params)
     return True

From 0db9eae72c147d8df1ec85ae284f05170ac7b632 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Wed, 9 Oct 2024 12:27:26 +0200
Subject: [PATCH 13/21] update ec2_ami_instance tests to include imds v2
 support

---
 .../targets/ec2_ami_instance/tasks/main.yml   | 35 +++++++++++++++----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/tests/integration/targets/ec2_ami_instance/tasks/main.yml b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
index 2b3c44c38c1..e682c6b93a0 100644
--- a/tests/integration/targets/ec2_ami_instance/tasks/main.yml
+++ b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
@@ -92,10 +92,10 @@
       amazon.aws.ec2_ami:
         instance_id: "{{ ec2_instance_id }}"
         state: present
-        name: "{{ ec2_ami_name }}_ami"
+        name: "{{ ec2_ami_name }}_ami_check"
         description: "{{ ec2_ami_description }}"
         tags:
-          Name: "{{ ec2_ami_name }}_ami"
+          Name: "{{ ec2_ami_name }}_ami_check"
         wait: true
         root_device_name: "{{ ec2_ami_root_disk }}"
       check_mode: true
@@ -110,10 +110,10 @@
       amazon.aws.ec2_ami:
         instance_id: "{{ ec2_instance_id }}"
         state: present
-        name: "{{ ec2_ami_name }}_ami"
+        name: "{{ ec2_ami_name }}_ami_instance"
         description: "{{ ec2_ami_description }}"
         tags:
-          Name: "{{ ec2_ami_name }}_ami"
+          Name: "{{ ec2_ami_name }}_ami_instance"
         wait: true
         root_device_name: "{{ ec2_ami_root_disk }}"
       register: result
@@ -127,7 +127,7 @@
         that:
           - result.changed
           - result.image_id.startswith('ami-')
-          - "'Name' in result.tags and result.tags.Name == ec2_ami_name + '_ami'"
+          - "'Name' in result.tags and result.tags.Name == ec2_ami_name + '_ami_instance'"
 
     - name: get related snapshot info and ensure the tags have been propagated
       amazon.aws.ec2_snapshot_info:
@@ -139,7 +139,7 @@
       ansible.builtin.assert:
         that:
           - "'tags' in snapshot_result.snapshots[0]"
-          - "'Name' in snapshot_result.snapshots[0].tags and snapshot_result.snapshots[0].tags.Name == ec2_ami_name + '_ami'"
+          - "'Name' in snapshot_result.snapshots[0].tags and snapshot_result.snapshots[0].tags.Name == ec2_ami_name + '_ami_instance'"
 
     # ============================================================
 
@@ -347,6 +347,29 @@
           - result.failed
           - "result.msg == 'state is absent but all of the following are missing: image_id'"
 
+    # ==============================================================
+
+    - name: Create image with imdsv2 enabled
+      amazon.aws.ec2_ami:
+        state: present
+        instance_id: "{{ ec2_instance_id }}"
+        name: "{{ ec2_ami_name }}_imdsv2"
+        imdsv2_enable: true
+        tags:
+          Name: "{{ ec2_ami_name }}_imdsv2"
+      register: imdsv2_image
+
+    - name: Get ami info
+      amazon.aws.ec2_ami_info:
+        image_ids: "{{ imdsv2_image.image_id }}"
+        describe_image_attributes: true
+      register: imdsv2_info
+
+    - name: Assert ami IMDSV2 is set
+      ansible.builtin.assert:
+        that:
+          - imdsv2_info.images[0].imdsv2_enabled
+
   always:
     # ============================================================
 

From 2bd6e1289d7930d2f42c317e09246e73b83a6f3c Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Wed, 9 Oct 2024 16:23:33 +0200
Subject: [PATCH 14/21] pass imdssupport parameters explicitly

---
 plugins/modules/ec2_ami.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index d2c7906714b..bbfa7ba7c2c 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -785,9 +785,10 @@ def set_imdsv2(connection, image, imdsv2_enable):
         if image["ImdsSupport"] == "v2.0":
             return False
 
-        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": {"Value": "v2.0"}}
         try:
-            modify_image_attribute(connection, image_id=image["imageId"], **imds_params)
+            modify_image_attribute(
+                connection, image_id=image["imageId"], Attribute="ImdsSupport", ImdsSupport={"Value": "v2.0"}
+            )
             return True
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image['imageId']}", e)
@@ -873,9 +874,10 @@ def set_launch_permissions(connection, launch_permissions, image_id):
 
     @staticmethod
     def set_imdsv2(connection, image_id):
-        imds_params = {"Attribute": "ImdsSupport", "ImdsSupport": {"Value": "v2.0"}}
         try:
-            modify_image_attribute(connection, image_id=image_id, **imds_params)
+            modify_image_attribute(
+                connection, image_id=image_id, Attribute="ImdsSupport", ImdsSupport={"Value": "v2.0"}
+            )
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image_id}", e)
 

From 0b8d21d8272838c09a41418128883d44814e3c78 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 07:00:53 +0200
Subject: [PATCH 15/21] wait for imdsv2 enabled image

---
 tests/integration/targets/ec2_ami_instance/tasks/main.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/targets/ec2_ami_instance/tasks/main.yml b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
index e682c6b93a0..0950a642670 100644
--- a/tests/integration/targets/ec2_ami_instance/tasks/main.yml
+++ b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
@@ -355,6 +355,7 @@
         instance_id: "{{ ec2_instance_id }}"
         name: "{{ ec2_ami_name }}_imdsv2"
         imdsv2_enable: true
+        wait: true
         tags:
           Name: "{{ ec2_ami_name }}_imdsv2"
       register: imdsv2_image

From 4a59e11606574e48be2948dd26785fb9b06ec5e6 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 09:50:02 +0200
Subject: [PATCH 16/21] add imdsv2 test image cleanup

---
 .../targets/ec2_ami_instance/tasks/main.yml          | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/integration/targets/ec2_ami_instance/tasks/main.yml b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
index 0950a642670..e9130bd735e 100644
--- a/tests/integration/targets/ec2_ami_instance/tasks/main.yml
+++ b/tests/integration/targets/ec2_ami_instance/tasks/main.yml
@@ -371,6 +371,10 @@
         that:
           - imdsv2_info.images[0].imdsv2_enabled
 
+    - name: set image id fact for deletion later
+      ansible.builtin.set_fact:
+        ec2_ami_image_id_imdsv2: "{{ imdsv2_image.image_id }}"
+
   always:
     # ============================================================
 
@@ -403,6 +407,14 @@
         wait: true
       ignore_errors: true
 
+    - name: delete imdsv2 ami
+      amazon.aws.ec2_ami:
+        state: absent
+        image_id: "{{ ec2_ami_image_id_imdsv2 }}"
+        name: "{{ ec2_ami_name }}_imdsv2"
+        wait: true
+      ignore_errors: true
+
     - name: delete ami
       amazon.aws.ec2_ami:
         state: absent

From 750eba3a900871b28905ff322aee54522b530f6c Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 12:25:42 +0200
Subject: [PATCH 17/21] imdssupport should be lower case

---
 plugins/modules/ec2_ami.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index bbfa7ba7c2c..ae13997cd35 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -787,7 +787,7 @@ def set_imdsv2(connection, image, imdsv2_enable):
 
         try:
             modify_image_attribute(
-                connection, image_id=image["imageId"], Attribute="ImdsSupport", ImdsSupport={"Value": "v2.0"}
+                connection, image_id=image["imageId"], Attribute="imdsSupport", ImdsSupport={"Value": "v2.0"}
             )
             return True
         except AnsibleEC2Error as e:
@@ -876,7 +876,7 @@ def set_launch_permissions(connection, launch_permissions, image_id):
     def set_imdsv2(connection, image_id):
         try:
             modify_image_attribute(
-                connection, image_id=image_id, Attribute="ImdsSupport", ImdsSupport={"Value": "v2.0"}
+                connection, image_id=image_id, Attribute="imdsSupport", ImdsSupport={"Value": "v2.0"}
             )
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image_id}", e)

From 83e47a1bddc2a4a0f8da36d7607253556380789c Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 14:34:46 +0200
Subject: [PATCH 18/21] only "value" needed when attr is imdssupport

---
 plugins/modules/ec2_ami.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/plugins/modules/ec2_ami.py b/plugins/modules/ec2_ami.py
index ae13997cd35..830172e2de3 100644
--- a/plugins/modules/ec2_ami.py
+++ b/plugins/modules/ec2_ami.py
@@ -786,9 +786,7 @@ def set_imdsv2(connection, image, imdsv2_enable):
             return False
 
         try:
-            modify_image_attribute(
-                connection, image_id=image["imageId"], Attribute="imdsSupport", ImdsSupport={"Value": "v2.0"}
-            )
+            modify_image_attribute(connection, image_id=image["imageId"], Attribute="imdsSupport", Value="v2.0")
             return True
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image['imageId']}", e)
@@ -875,9 +873,7 @@ def set_launch_permissions(connection, launch_permissions, image_id):
     @staticmethod
     def set_imdsv2(connection, image_id):
         try:
-            modify_image_attribute(
-                connection, image_id=image_id, Attribute="imdsSupport", ImdsSupport={"Value": "v2.0"}
-            )
+            modify_image_attribute(connection, image_id=image_id, Attribute="imdsSupport", Value="v2.0")
         except AnsibleEC2Error as e:
             raise Ec2AmiFailure(f"Error setting IMDS Support to v2 for image {image_id}", e)
 

From 656b96e483712d428f47653e29f10e012041b96e Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 16:03:19 +0200
Subject: [PATCH 19/21] fix attribute name when getting imds support

---
 plugins/modules/ec2_ami_info.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/modules/ec2_ami_info.py b/plugins/modules/ec2_ami_info.py
index d277bb74f51..7dd0642871a 100644
--- a/plugins/modules/ec2_ami_info.py
+++ b/plugins/modules/ec2_ami_info.py
@@ -282,9 +282,9 @@ def list_ec2_images(ec2_client, module, request_args):
                 ).get("LaunchPermissions", [])
                 image["launch_permissions"] = [camel_dict_to_snake_dict(perm) for perm in launch_permissions]
                 imdsv2_enabled = describe_image_attribute(
-                    ec2_client, attribute="ImdsSupport", image_id=image["image_id"]
+                    ec2_client, attribute="imdsSupport", image_id=image["image_id"]
                 ).get("ImdsSupport", "")
-                image["imdsv2_enabled"] = True if imdsv2_enabled == "v2.0" else False
+                image["imdsv2_enabled"] = True if imdsv2_enabled["Value"] == "v2.0" else False
         except is_ansible_aws_error_code("AuthFailure"):
             # describing launch permissions of images owned by others is not permitted, but shouldn't cause failures
             pass

From f8db7f7eef315e0d358e4a5c6bb83b6cf7f91518 Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 16:51:19 +0200
Subject: [PATCH 20/21] check if Value is returned by the API before checking
 value

---
 plugins/modules/ec2_ami_info.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/plugins/modules/ec2_ami_info.py b/plugins/modules/ec2_ami_info.py
index 7dd0642871a..99efba5ae10 100644
--- a/plugins/modules/ec2_ami_info.py
+++ b/plugins/modules/ec2_ami_info.py
@@ -283,8 +283,10 @@ def list_ec2_images(ec2_client, module, request_args):
                 image["launch_permissions"] = [camel_dict_to_snake_dict(perm) for perm in launch_permissions]
                 imdsv2_enabled = describe_image_attribute(
                     ec2_client, attribute="imdsSupport", image_id=image["image_id"]
-                ).get("ImdsSupport", "")
-                image["imdsv2_enabled"] = True if imdsv2_enabled["Value"] == "v2.0" else False
+                ).get("ImdsSupport", {})
+                image["imdsv2_enabled"] = (
+                    True if "Value" in imdsv2_enabled and imdsv2_enabled["Value"] == "v2.0" else False
+                )
         except is_ansible_aws_error_code("AuthFailure"):
             # describing launch permissions of images owned by others is not permitted, but shouldn't cause failures
             pass

From b3afb671aa65055ace378f98b13b5cb1ee8ee40b Mon Sep 17 00:00:00 2001
From: Andrea Tartaglia <me@andreatartaglia.com>
Date: Thu, 10 Oct 2024 16:51:55 +0200
Subject: [PATCH 21/21] update unit test to add imds support tests

---
 tests/unit/plugins/modules/test_ec2_ami_info.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/unit/plugins/modules/test_ec2_ami_info.py b/tests/unit/plugins/modules/test_ec2_ami_info.py
index 12c7e175acd..7482990895a 100644
--- a/tests/unit/plugins/modules/test_ec2_ami_info.py
+++ b/tests/unit/plugins/modules/test_ec2_ami_info.py
@@ -140,6 +140,7 @@ def test_list_ec2_images(m_get_images, m_describe_image_attribute):
     m_describe_image_attribute.return_value = {
         "ImageId": "ami-1234567890",
         "LaunchPermissions": [{"UserId": "1234567890"}, {"UserId": "0987654321"}],
+        "ImdsSupport": {"Value": "v2.0"},
     }
 
     images = m_get_images.return_value
@@ -166,13 +167,14 @@ def test_list_ec2_images(m_get_images, m_describe_image_attribute):
         [call(ec2_client, attribute="launchPermission", image_id=images[1]["image_id"])],
     )
     m_describe_image_attribute.assert_has_calls(
-        [call(ec2_client, attribute="ImdsSupport", image_id=images[0]["image_id"])],
-        [call(ec2_client, attribute="ImdsSupport", image_id=images[1]["image_id"])],
+        [call(ec2_client, attribute="imdsSupport", image_id=images[0]["image_id"])],
+        [call(ec2_client, attribute="imdsSupport", image_id=images[1]["image_id"])],
     )
 
     assert len(list_ec2_images_result) == 2
     assert list_ec2_images_result[0]["image_id"] == "ami-1234567890"
     assert list_ec2_images_result[1]["image_id"] == "ami-1523498760"
+    assert list_ec2_images_result[0]["imdsv2_enabled"]
 
 
 @patch(module_name + ".AnsibleAWSModule")