Merge pull request #67 from toabctl/main-shared

Restrict share by partition

Author: toabctl

.github/workflows/pr.yaml

@@ -19,9 +19,7 @@ jobs:
           python-version: ${{ matrix.python }}
       - name: Install tox and any other packages
         run: |
-          # FIXME: not pining to an exact version results in getting poetry 1.0.10 which doesn't work
-          pip3 install tox poetry==1.7.1
-          poetry --version
+          pip3 install tox
       - uses: actions/checkout@v3
       - name: Run tox
         # Run tox using the version of Python in `PATH`

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

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

@@ -0,0 +1,12 @@
+awspub:
+  source:
+    path: "image.vmdk"
+    architecture: "x86_64"
+  s3:
+    bucket_name: "awspub-toabctl"
+  images:
+    "my-custom-image":
+      boot_mode: "uefi-preferred"
+      share:
+        - "123456789123"
+        - "aws-cn:456789012345"

docs/how_to/publish.rst

@@ -187,6 +187,18 @@ The image needs to be created and then published:
   awspub create config.yaml
   awspub publish config.yaml
 
+Sharing images
+~~~~~~~~~~~~~~
+
+Images can be shared with other AWS accounts. For that, the account IDs of the other accounts are needed.
+
+.. literalinclude:: ../config-samples/config-minimal-share.yaml
+   :language: yaml
+
+In the above example, the image `my-custom-image` will be shared with the account `1234567890123`
+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``).
+
 AWS Marketplace
 ~~~~~~~~~~~~~~~