Repository: getmoto/moto (python)

Issue:
When modifying EC2 image attribute to add organization ARN to launch permissions, moto silently ignores it.
**How to reproduce**

```
image = boto3.resource("ec2").Image(ami_id)
arn = "someOrganizationArn"
image.modify_attribute(
    Attribute="launchPermission",
    OperationType="add",
    OrganizationArns=[arn],
)

launch_permissions = image.describe_attribute(Attribute="launchPermission")["LaunchPermissions"]
assert launch_permissions == [{"OrganizationArn": arn}]
```

I expected that the organization arn that was added to the launch permission was visible afterwards via describe_attribute however it was not. Moto API calls are successful, but what happened was there, there was no launch permissions at all. Moto silently ignored adding an organization ARN to launch permission of an image.

Using boto3==1.21.9, botocore==1.24.46 and moto==5.0.0

Pull request:
EC2: modify_image_attributes() now supports all LaunchPermissions
Fixes #7318 

`modify_image_attributes()` now supports the `LaunchPermission`, `OrganizationArns` and `OrganizationalUnitArns` attributes

Internally, the launch_permissions are now stored as a dict, more closely resembling the actual data.

Patch:
diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md
index d7cb91469..cb973e06d 100644
--- a/IMPLEMENTATION_COVERAGE.md
+++ b/IMPLEMENTATION_COVERAGE.md
@@ -2523,7 +2523,7 @@
 - [X] modify_hosts
 - [ ] modify_id_format
 - [ ] modify_identity_id_format
-- [ ] modify_image_attribute
+- [X] modify_image_attribute
 - [X] modify_instance_attribute
 - [ ] modify_instance_capacity_reservation_attributes
 - [ ] modify_instance_credit_specification
diff --git a/docs/docs/services/ec2.rst b/docs/docs/services/ec2.rst
index 4398c475c..2750ebd76 100644
--- a/docs/docs/services/ec2.rst
+++ b/docs/docs/services/ec2.rst
@@ -545,7 +545,7 @@ ec2
 - [X] modify_hosts
 - [ ] modify_id_format
 - [ ] modify_identity_id_format
-- [ ] modify_image_attribute
+- [X] modify_image_attribute
 - [X] modify_instance_attribute
 - [ ] modify_instance_capacity_reservation_attributes
 - [ ] modify_instance_credit_specification
diff --git a/moto/ec2/models/amis.py b/moto/ec2/models/amis.py
index 12028f2cd..862b97536 100644
--- a/moto/ec2/models/amis.py
+++ b/moto/ec2/models/amis.py
@@ -101,11 +101,10 @@ class Ami(TaggedEC2Resource):
             if not description:
                 self.description = source_ami.description
 
-        self.launch_permission_groups: Set[str] = set()
-        self.launch_permission_users: Set[str] = set()
+        self.launch_permissions: List[Dict[str, str]] = []
 
         if public:
-            self.launch_permission_groups.add("all")
+            self.launch_permissions.append({"Group": "all"})
 
         # AWS auto-creates these, we should reflect the same.
         volume = self.ec2_backend.create_volume(size=15, zone_name=region_name)
@@ -119,7 +118,7 @@ class Ami(TaggedEC2Resource):
 
     @property
     def is_public(self) -> bool:
-        return "all" in self.launch_permission_groups
+        return {"Group": "all"} in self.launch_permissions
 
     @property
     def is_public_string(self) -> str:
@@ -273,8 +272,9 @@ class AmiBackend:
                 tmp_images = []
                 for ami in images:
                     for user_id in exec_users:
-                        if user_id in ami.launch_permission_users:
-                            tmp_images.append(ami)
+                        for lp in ami.launch_permissions:
+                            if lp.get("UserId") == user_id:
+                                tmp_images.append(ami)
                 images = tmp_images
 
             # Limit by owner ids
@@ -305,38 +305,39 @@ class AmiBackend:
         else:
             raise InvalidAMIIdError(ami_id)
 
-    def validate_permission_targets(
-        self, user_ids: Optional[List[str]] = None, group: Optional[str] = None
-    ) -> None:
-        # If anything is invalid, nothing is added. (No partial success.)
-        if user_ids:
-            """
-            AWS docs:
-              "The AWS account ID is a 12-digit number, such as 123456789012, that you use to construct Amazon Resource Names (ARNs)."
-              http://docs.aws.amazon.com/general/latest/gr/acct-identifiers.html
-            """
-            for user_id in user_ids:
-                if len(user_id) != 12 or not user_id.isdigit():
-                    raise InvalidAMIAttributeItemValueError("userId", user_id)
+    def validate_permission_targets(self, permissions: List[Dict[str, str]]) -> None:
+        for perm in permissions:
+            # If anything is invalid, nothing is added. (No partial success.)
+            # AWS docs:
+            # The AWS account ID is a 12-digit number, such as 123456789012, that you use to construct Amazon Resource Names (ARNs)."
+            # http://docs.aws.amazon.com/general/latest/gr/acct-identifiers.html
+
+            if "UserId" in perm and (
+                len(perm["UserId"]) != 12 or not perm["UserId"].isdigit()
+            ):
+                raise InvalidAMIAttributeItemValueError("userId", perm["UserId"])
 
-        if group and group != "all":
-            raise InvalidAMIAttributeItemValueError("UserGroup", group)
+            if "Group" in perm and perm["Group"] != "all":
+                raise InvalidAMIAttributeItemValueError("UserGroup", perm["Group"])
 
-    def add_launch_permission(
+    def modify_image_attribute(
         self,
         ami_id: str,
-        user_ids: Optional[List[str]] = None,
-        group: Optional[str] = None,
+        launch_permissions_to_add: List[Dict[str, str]],
+        launch_permissions_to_remove: List[Dict[str, str]],
     ) -> None:
         ami = self.describe_images(ami_ids=[ami_id])[0]
-        self.validate_permission_targets(user_ids=user_ids, group=group)
-
-        if user_ids:
-            for user_id in user_ids:
-                ami.launch_permission_users.add(user_id)
-
-        if group:
-            ami.launch_permission_groups.add(group)
+        self.validate_permission_targets(launch_permissions_to_add)
+        self.validate_permission_targets(launch_permissions_to_remove)
+        for lp in launch_permissions_to_add:
+            if lp not in ami.launch_permissions:
+                ami.launch_permissions.append(lp)
+        for lp in launch_permissions_to_remove:
+            try:
+                ami.launch_permissions.remove(lp)
+            except ValueError:
+                # The LaunchPermission may not exist
+                pass
 
     def register_image(
         self, name: Optional[str] = None, description: Optional[str] = None
@@ -353,21 +354,5 @@ class AmiBackend:
         self.amis[ami_id] = ami
         return ami
 
-    def remove_launch_permission(
-        self,
-        ami_id: str,
-        user_ids: Optional[List[str]] = None,
-        group: Optional[str] = None,
-    ) -> None:
-        ami = self.describe_images(ami_ids=[ami_id])[0]
-        self.validate_permission_targets(user_ids=user_ids, group=group)
-
-        if user_ids:
-            for user_id in user_ids:
-                ami.launch_permission_users.discard(user_id)
-
-        if group:
-            ami.launch_permission_groups.discard(group)
-
     def describe_image_attribute(self, ami_id: str, attribute_name: str) -> Any:
         return self.amis[ami_id].__getattribute__(attribute_name)
diff --git a/moto/ec2/responses/amis.py b/moto/ec2/responses/amis.py
index c9e36607d..7238c32e5 100644
--- a/moto/ec2/responses/amis.py
+++ b/moto/ec2/responses/amis.py
@@ -65,10 +65,7 @@ class AmisResponse(EC2BaseResponse):
             "description": "description",
             "kernel": "kernel_id",
             "ramdisk": "ramdisk",
-            "launchPermission": {
-                "groups": "launch_permission_groups",
-                "users": "launch_permission_users",
-            },
+            "launchPermission": "launch_permissions",
             "productCodes": "product_codes",
             "blockDeviceMapping": "bdm",
             "sriovNetSupport": "sriov",
@@ -85,45 +82,63 @@ class AmisResponse(EC2BaseResponse):
             # https://github.com/aws/aws-cli/issues/1083
             raise AuthFailureRestricted
 
-        groups = None
-        users = None
         attribute_value = None
+        launch_permissions = None
         if attribute_name == "launchPermission":
-            groups = self.ec2_backend.describe_image_attribute(
-                ami_id, valid_atrributes_list[attribute_name]["groups"]  # type: ignore[index]
-            )
-            users = self.ec2_backend.describe_image_attribute(
-                ami_id, valid_atrributes_list[attribute_name]["users"]  # type: ignore[index]
+            launch_permissions = self.ec2_backend.describe_image_attribute(
+                ami_id, valid_atrributes_list[attribute_name]
             )
         else:
             attribute_value = self.ec2_backend.describe_image_attribute(
                 ami_id, valid_atrributes_list[attribute_name]
             )
+
         template = self.response_template(DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE)
         return template.render(
             ami_id=ami_id,
-            users=users,
-            groups=groups,
+            launch_permissions=launch_permissions,
             attribute_name=attribute_name,
             attribute_value=attribute_value,
         )
 
     def modify_image_attribute(self) -> str:
         ami_id = self._get_param("ImageId")
+        launch_permissions_to_add = list(
+            self._get_params().get("LaunchPermission", {}).get("Add", {}).values()
+        )
+        launch_permissions_to_remove = list(
+            self._get_params().get("LaunchPermission", {}).get("Remove", {}).values()
+        )
+        # If only one OperationType is added, the other attributes are submitted as different variables
         operation_type = self._get_param("OperationType")
-        group = self._get_param("UserGroup.1")
-        user_ids = self._get_multi_param("UserId")
+        if operation_type in ["add", "remove"]:
+            group = self._get_param("UserGroup.1")
+            lp = (
+                launch_permissions_to_add
+                if operation_type == "add"
+                else launch_permissions_to_remove
+            )
+            if group:
+                lp.append({"Group": group})
+
+            for user_id in self._get_multi_param("UserId"):
+                lp.append({"UserId": user_id})
+
+            org_arn = self._get_param("OrganizationArn.1")
+            if org_arn:
+                lp.append({"OrganizationArn": org_arn})
+
+            ou_arn = self._get_param("OrganizationalUnitArn.1")
+            if ou_arn:
+                lp.append({"OrganizationalUnitArn": ou_arn})
 
         self.error_on_dryrun()
 
-        if operation_type == "add":
-            self.ec2_backend.add_launch_permission(
-                ami_id, user_ids=user_ids, group=group
-            )
-        elif operation_type == "remove":
-            self.ec2_backend.remove_launch_permission(
-                ami_id, user_ids=user_ids, group=group
-            )
+        self.ec2_backend.modify_image_attribute(
+            ami_id=ami_id,
+            launch_permissions_to_add=launch_permissions_to_add,
+            launch_permissions_to_remove=launch_permissions_to_remove,
+        )
         return MODIFY_IMAGE_ATTRIBUTE_RESPONSE
 
     def register_image(self) -> str:
@@ -232,17 +247,13 @@ DESCRIBE_IMAGE_ATTRIBUTES_RESPONSE = """
         {% endfor %}
     {% endif %}
     {% if attribute_name == 'launchPermission' %}
-         {% if groups %}
-            {% for group in groups %}
-               <item>
-                  <group>{{ group }}</group>
-               </item>
-            {% endfor %}
-         {% endif %}
-         {% if users %}
-            {% for user in users %}
+         {% if launch_permissions %}
+            {% for lp in launch_permissions %}
                <item>
-                  <userId>{{ user }}</userId>
+                  {% if lp['UserId'] %}<userId>{{ lp['UserId'] }}</userId>{% endif %}
+                  {% if lp['Group'] %}<group>{{ lp['Group'] }}</group>{% endif %}
+                  {% if lp['OrganizationArn'] %}<organizationArn>{{ lp['OrganizationArn'] }}</organizationArn>{% endif %}
+                  {% if lp['OrganizationalUnitArn'] %}<organizationalUnitArn>{{ lp['OrganizationalUnitArn'] }}</organizationalUnitArn>{% endif %}
                </item>
             {% endfor %}
          {% endif %}


Test patch:
diff --git a/tests/test_ec2/test_amis.py b/tests/test_ec2/test_amis.py
index e8bff43a0..87905a3ce 100644
--- a/tests/test_ec2/test_amis.py
+++ b/tests/test_ec2/test_amis.py
@@ -426,6 +426,9 @@ def test_ami_filters():
     ]
     imageB = boto3.resource("ec2", region_name="us-east-1").Image(imageB_id)
     imageB.modify_attribute(LaunchPermission={"Add": [{"Group": "all"}]})
+    assert imageB.describe_attribute(Attribute="launchPermission")[
+        "LaunchPermissions"
+    ] == [{"Group": "all"}]
 
     amis_by_architecture = ec2.describe_images(
         Filters=[{"Name": "architecture", "Values": ["x86_64"]}]
@@ -479,7 +482,6 @@ def test_ami_filters():
         Filters=[{"Name": "is-public", "Values": ["false"]}]
     )["Images"]
     assert imageA.id in [ami["ImageId"] for ami in amis_by_nonpublic]
-    assert len(amis_by_nonpublic) >= 2, "Should have at least 2 non-public images"
 
 
 @mock_aws
@@ -684,6 +686,53 @@ def test_ami_attribute_user_permissions():
     image.modify_attribute(**REMOVE_USERS_ARGS)
 
 
+@mock_aws
+def test_ami_attribute_organizations():
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+    reservation = ec2.run_instances(ImageId=EXAMPLE_AMI_ID, MinCount=1, MaxCount=1)
+    instance_id = reservation["Instances"][0]["InstanceId"]
+    image_id = ec2.create_image(InstanceId=instance_id, Name="test-ami-A")["ImageId"]
+    image = boto3.resource("ec2", "us-east-1").Image(image_id)
+    arn = "someOrganizationArn"
+    image.modify_attribute(
+        Attribute="launchPermission",
+        OperationType="add",
+        OrganizationArns=[arn],
+    )
+    image.modify_attribute(
+        Attribute="launchPermission",
+        OperationType="add",
+        OrganizationalUnitArns=["ou1"],
+    )
+
+    ec2.modify_image_attribute(
+        Attribute="launchPermission",
+        ImageId=image_id,
+        LaunchPermission={
+            "Add": [
+                {"UserId": "111122223333"},
+                {"UserId": "555566667777"},
+                {"Group": "all"},
+                {"OrganizationArn": "orgarn"},
+                {"OrganizationalUnitArn": "ou2"},
+            ]
+        },
+    )
+
+    launch_permissions = image.describe_attribute(Attribute="launchPermission")[
+        "LaunchPermissions"
+    ]
+    assert launch_permissions == [
+        {"OrganizationArn": "someOrganizationArn"},
+        {"OrganizationalUnitArn": "ou1"},
+        {"UserId": "111122223333"},
+        {"UserId": "555566667777"},
+        {"Group": "all"},
+        {"OrganizationArn": "orgarn"},
+        {"OrganizationalUnitArn": "ou2"},
+    ]
+
+
 @mock_aws
 def test_ami_describe_executable_users():
     conn = boto3.client("ec2", region_name="us-east-1")