Add support for sharing with Organizations and Organizational Units

This PR extends "share:" schema to support sharing images with
organizations & organization units by ARN.

Fixes: #105

Author: xnox

awspub/common.py

@@ -15,13 +15,20 @@ def _split_partition(val: str) -> Tuple[str, str]:
     :return: the partition and the resource
     :rtype: Tuple[str, str]
     """
-    if ":" in val:
-        partition, resource = val.split(":")
-    else:
-        # if no partition is given, assume default commercial partition "aws"
-        partition = "aws"
-        resource = val
-    return partition, resource
+
+    # ARNs encode partition https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html
+    if val.startswith("arn:"):
+        arn, partition, resource = val.split(":", maxsplit=2)
+        # Return extracted partition, but keep full ARN intact
+        return partition, val
+
+    # Partition prefix
+    if ":" in val and val.startswith("aws"):
+        partition, resource = val.split(":", maxsplit=1)
+        return partition, resource
+
+    # if no partition is given, assume default commercial partition "aws"
+    return "aws", val
 
 
 def _get_regions(region_to_query: str, regions_allowlist: List[str]) -> List[str]:

awspub/configmodels.py

@@ -1,4 +1,5 @@
 import pathlib
+import re
 from enum import Enum
 from typing import Dict, List, Literal, Optional
 
@@ -161,8 +162,8 @@ class ConfigImageModel(BaseModel):
     )
     imds_support: Optional[Literal["v2.0"]] = Field(description="Optional IMDS support", default=None)
     share: Optional[List[str]] = Field(
-        description="Optional list of account IDs the image and snapshot will be shared with. The account"
-        "ID can be prefixed with the partition and separated by ':'. Eg 'aws-cn:123456789123'",
+        description="Optional list of account IDs, organization ARN, OU ARN the image and snapshot will be shared with."
+        " The account ID can be prefixed with the partition and separated by ':'. Eg 'aws-cn:123456789123'",
         default=None,
     )
     temporary: Optional[bool] = Field(
@@ -193,11 +194,25 @@ def check_share(cls, v: Optional[List[str]]) -> Optional[List[str]]:
         """
         Make sure the account IDs are valid and if given the partition is correct
         """
+        patterns = [
+            # https://docs.aws.amazon.com/organizations/latest/APIReference/API_Account.html
+            r"\d{12}",
+            # Adjusted for partitions
+            # https://docs.aws.amazon.com/organizations/latest/APIReference/API_Organization.html
+            r"arn:aws(?:-cn)?(?:-us-gov)?:organizations::\d{12}:organization\/o-[a-z0-9]{10,32}",
+            # https://docs.aws.amazon.com/organizations/latest/APIReference/API_OrganizationalUnit.html
+            r"arn:aws(?:-cn)?(?:-us-gov)?:organizations::\d{12}:ou\/o-[a-z0-9]{10,32}\/ou-[0-9a-z]{4,32}-[0-9a-z]{8,32}",  # noqa:E501
+        ]
         if v is not None:
             for val in v:
-                partition, account_id = _split_partition(val)
-                if len(account_id) != 12:
-                    raise ValueError("Account ID must be 12 characters long")
+                partition, account_id_or_arn = _split_partition(val)
+                valid = False
+                for pattern in patterns:
+                    if re.fullmatch(pattern, account_id_or_arn):
+                        valid = True
+                        break
+                if not valid:
+                    raise ValueError("Account ID must be 12 digits long or an ARN for Organization or OU")
                 if partition not in ["aws", "aws-cn", "aws-us-gov"]:
                     raise ValueError("Partition must be one of 'aws', 'aws-cn', 'aws-us-gov'")
         return v

awspub/image.py

@@ -2,7 +2,7 @@
 import logging
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import boto3
 import botocore.exceptions
@@ -145,23 +145,30 @@ def _tags(self):
             tags.append({"Key": name, "Value": value})
         return tags
 
-    def _share_list_filtered(self, share_conf: List[str]) -> List[Dict[str, str]]:
+    def _share_list_filtered(self, share_conf: List[str]) -> Tuple[List[Dict[str, str]], List[Dict[str, str]]]:
         """
         Get a filtered list of share configurations based on the current partition
         :param share_conf: the share configuration
         :type share_conf: List[str]
         :return: a List of share configurations that is usable by modify_image_attribute()
-        :rtype: List[Dict[str, str]]
+        :rtype: Tuple[List[Dict[str, str]], List[Dict[str, str]]]
         """
         # the current partition
         partition_current = boto3.client("ec2").meta.partition
 
         share_list: List[Dict[str, str]] = []
+        volume_list: List[Dict[str, str]] = []
         for share in share_conf:
-            partition, account_id = _split_partition(share)
+            partition, account_id_or_arn = _split_partition(share)
             if partition == partition_current:
-                share_list.append({"UserId": account_id})
-        return share_list
+                if ":organization/o-" in account_id_or_arn:
+                    share_list.append({"OrganizationArn": account_id_or_arn})
+                elif ":ou/o-" in account_id_or_arn:
+                    share_list.append({"OrganizationalUnitArn": account_id_or_arn})
+                else:
+                    share_list.append({"UserId": account_id_or_arn})
+                    volume_list.append({"UserId": account_id_or_arn})
+        return share_list, volume_list
 
     def _share(self, share_conf: List[str], images: Dict[str, _ImageInfo]):
         """
@@ -172,7 +179,7 @@ def _share(self, share_conf: List[str], images: Dict[str, _ImageInfo]):
         :param images: a Dict with region names as keys and _ImageInfo objects as values
         :type images: Dict[str, _ImageInfo]
         """
-        share_list = self._share_list_filtered(share_conf)
+        share_list, volume_list = self._share_list_filtered(share_conf)
 
         if not share_list:
             logger.info("no valid accounts found for sharing in this partition, skipping")
@@ -188,11 +195,11 @@ def _share(self, share_conf: List[str], images: Dict[str, _ImageInfo]):
             )
 
             # modify snapshot permissions
-            if image_info.snapshot_id:
+            if image_info.snapshot_id and volume_list:
                 ec2client.modify_snapshot_attribute(
                     Attribute="createVolumePermission",
                     SnapshotId=image_info.snapshot_id,
-                    CreateVolumePermission={"Add": share_list},  # type: ignore
+                    CreateVolumePermission={"Add": volume_list},  # type: ignore
                 )
 
         logger.info(f"shared images & snapshots with '{share_conf}'")

awspub/tests/fixtures/config1.yaml

@@ -83,6 +83,8 @@ awspub:
         - "221020170000"
         - "aws:290620200000"
         - "aws-cn:334455667788"
+        - "arn:aws:organizations::123456789012:organization/o-123example"
+        - "arn:aws-cn:organizations::334455667788:ou/o-123example/ou-1234-5example"
       marketplace:
         entity_id: "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
         access_role_arn: "arn:aws:iam::xxxxxxxxxxxx:role/AWSMarketplaceAccess"

awspub/tests/test_common.py

@@ -12,6 +12,30 @@
         ("aws:123456789123", ("aws", "123456789123")),
         ("aws-cn:123456789123", ("aws-cn", "123456789123")),
         ("aws-us-gov:123456789123", ("aws-us-gov", "123456789123")),
+        (
+            "arn:aws:organizations::123456789012:organization/o-123example",
+            ("aws", "arn:aws:organizations::123456789012:organization/o-123example"),
+        ),
+        (
+            "arn:aws:organizations::123456789012:ou/o-123example/ou-1234-5example",
+            ("aws", "arn:aws:organizations::123456789012:ou/o-123example/ou-1234-5example"),
+        ),
+        (
+            "arn:aws-cn:organizations::123456789012:organization/o-123example",
+            ("aws-cn", "arn:aws-cn:organizations::123456789012:organization/o-123example"),
+        ),
+        (
+            "arn:aws-cn:organizations::123456789012:ou/o-123example/ou-1234-5example",
+            ("aws-cn", "arn:aws-cn:organizations::123456789012:ou/o-123example/ou-1234-5example"),
+        ),
+        (
+            "arn:aws-us-gov:organizations::123456789012:organization/o-123example",
+            ("aws-us-gov", "arn:aws-us-gov:organizations::123456789012:organization/o-123example"),
+        ),
+        (
+            "arn:aws-us-gov:organizations::123456789012:ou/o-123example/ou-1234-5example",
+            ("aws-us-gov", "arn:aws-us-gov:organizations::123456789012:ou/o-123example/ou-1234-5example"),
+        ),
     ],
 )
 def test_common__split_partition(input, expected_output):

awspub/tests/test_image.py

@@ -478,14 +478,38 @@ def test_image__verify(image_found, config, config_image_name, expected_problems
 
 
 @pytest.mark.parametrize(
-    "partition,imagename,share_list_expected",
+    "partition,imagename,share_list_expected,volume_list_expected",
     [
-        ("aws", "test-image-8", [{"UserId": "123456789123"}, {"UserId": "221020170000"}, {"UserId": "290620200000"}]),
-        ("aws-cn", "test-image-8", [{"UserId": "334455667788"}]),
-        ("aws-us-gov", "test-image-8", []),
+        (
+            "aws",
+            "test-image-8",
+            [
+                {"UserId": "123456789123"},
+                {"UserId": "221020170000"},
+                {"UserId": "290620200000"},
+                {"OrganizationArn": "arn:aws:organizations::123456789012:organization/o-123example"},
+            ],
+            [
+                {"UserId": "123456789123"},
+                {"UserId": "221020170000"},
+                {"UserId": "290620200000"},
+            ],
+        ),
+        (
+            "aws-cn",
+            "test-image-8",
+            [
+                {"UserId": "334455667788"},
+                {"OrganizationalUnitArn": "arn:aws-cn:organizations::334455667788:ou/o-123example/ou-1234-5example"},
+            ],
+            [
+                {"UserId": "334455667788"},
+            ],
+        ),
+        ("aws-us-gov", "test-image-8", [], []),
     ],
 )
-def test_image__share_list_filtered(partition, imagename, share_list_expected):
+def test_image__share_list_filtered(partition, imagename, share_list_expected, volume_list_expected):
     """
     Test _share_list_filtered() for a given image
     """
@@ -494,7 +518,7 @@ def test_image__share_list_filtered(partition, imagename, share_list_expected):
         instance.meta.partition = partition
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
         img = image.Image(ctx, imagename)
-        assert img._share_list_filtered(img.conf["share"]) == share_list_expected
+        assert img._share_list_filtered(img.conf["share"]) == (share_list_expected, volume_list_expected)
 
 
 @patch("awspub.s3.S3.bucket_region", return_value="region1")

docs/config-samples/config-minimal-share.yaml

@@ -10,3 +10,5 @@ awspub:
       share:
         - "123456789123"
         - "aws-cn:456789012345"
+        - "arn:aws:organizations::123456789012:organization/o-123example"
+        - "arn:aws-cn:organizations::334455667788:ou/o-123example/ou-1234-5example"

docs/how_to/publish.rst

@@ -199,6 +199,9 @@ In the above example, the image `my-custom-image` will be shared with the accoun
 when `awspub` runs in the commercial partition (``aws``, the default). It'll be shared
 with the account `456789012345` when `awspub` runs in the the china partition (``aws-cn``).
 
+Also sharing with an organization and organisation units is available by organization or organisational
+unit ARN (which already encodes partition).
+
 AWS Marketplace
 ~~~~~~~~~~~~~~~