feat: allow account sharing restrictions by partition

Currently the list of given account IDs in the "share" config section
for each image will be applied no matter on which AWS partition. But
account IDs are different on the different partitions so account
123456789123 is something different on "aws", "aws-cn" and
"aws-us-gov".
Being able to restrict account sharing by partition makes it possible
to reuse the same awspub.yaml configuration across different
partitions.

Author: toabctl

awspub/common.py

@@ -0,0 +1,18 @@
+from typing import Tuple
+
+
+def _split_partition(val: str) -> Tuple[str, str]:
+    """
+    Split a string into partition and resource, separated by a colon. If no partition is given, assume "aws"
+    :param val: the string to split
+    :type val: 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

awspub/configmodels.py

@@ -1,7 +1,9 @@
 import pathlib
 from typing import Dict, List, Literal, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from awspub.common import _split_partition
 
 
 class ConfigS3Model(BaseModel):
@@ -124,7 +126,9 @@ 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", default=None
+        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'",
+        default=None,
     )
     temporary: Optional[bool] = Field(
         description="Optional boolean field indicates that a image is only temporary", default=False
@@ -145,6 +149,21 @@ class ConfigImageModel(BaseModel):
     groups: Optional[List[str]] = Field(description="Optional list of groups this image is part of", default=[])
     tags: Optional[Dict[str, str]] = Field(description="Optional Tags to apply to this image only", default={})
 
+    @field_validator("share")
+    @classmethod
+    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
+        """
+        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")
+                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
+
 
 class ConfigModel(BaseModel):
     """

awspub/image.py

@@ -9,6 +9,7 @@
 from mypy_boto3_ssm import SSMClient
 
 from awspub import exceptions
+from awspub.common import _split_partition
 from awspub.context import Context
 from awspub.image_marketplace import ImageMarketplace
 from awspub.s3 import S3
@@ -159,16 +160,34 @@ def _tags(self):
             tags.append({"Key": name, "Value": value})
         return tags
 
+    def _share_list_filtered(self, share_conf: List[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]]
+        """
+        # the current partition
+        partition_current = boto3.client("ec2").meta.partition
+
+        share_list: List[Dict[str, str]] = []
+        for share in share_conf:
+            partition, account_id = _split_partition(share)
+            if partition == partition_current:
+                share_list.append({"UserId": account_id})
+        return share_list
+
     def _share(self, share_conf: List[str], images: Dict[str, _ImageInfo]):
         """
         Share images with accounts
 
-        :param share_conf: the share configuration. eg. self.conf["share_create"]
+        :param share_conf: the share configuration containing list
         :type share_conf: List[str]
         :param images: a Dict with region names as keys and _ImageInfo objects as values
         :type images: Dict[str, _ImageInfo]
         """
-        share_list: List[Dict[str, str]] = [{"UserId": user_id} for user_id in share_conf]
+        share_list = self._share_list_filtered(share_conf)
 
         for region, image_info in images.items():
             ec2client: EC2Client = boto3.client("ec2", region_name=region)

awspub/tests/fixtures/config1.yaml

@@ -78,6 +78,11 @@ awspub:
       public: true
       tags:
         key1: value1
+      share:
+        - "123456789123"
+        - "221020170000"
+        - "aws:290620200000"
+        - "aws-cn:334455667788"
       marketplace:
         entity_id: "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
         access_role_arn: "arn:aws:iam::xxxxxxxxxxxx:role/AWSMarketplaceAccess"

awspub/tests/test_common.py

@@ -0,0 +1,16 @@
+import pytest
+
+from awspub.common import _split_partition
+
+
+@pytest.mark.parametrize(
+    "input,expected_output",
+    [
+        ("123456789123", ("aws", "123456789123")),
+        ("aws:123456789123", ("aws", "123456789123")),
+        ("aws-cn:123456789123", ("aws-cn", "123456789123")),
+        ("aws-us-gov:123456789123", ("aws-us-gov", "123456789123")),
+    ],
+)
+def test_common__split_partition(input, expected_output):
+    assert _split_partition(input) == expected_output

awspub/tests/test_image.py

@@ -456,3 +456,23 @@ def test_image__verify(image_found, config, config_image_name, expected_problems
         img = image.Image(ctx, config_image_name)
         problems = img._verify("eu-central-1")
         assert problems == expected_problems
+
+
+@pytest.mark.parametrize(
+    "partition,imagename,share_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", []),
+    ],
+)
+def test_image__share_list_filtered(partition, imagename, share_list_expected):
+    """
+    Test _share_list_filtered() for a given image
+    """
+    with patch("boto3.client") as bclient_mock:
+        instance = bclient_mock.return_value
+        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