Repository: cloud-custodian/cloud-custodian

Issue:
Ability to remove individual account IDs from the AMI launch permissions with aws.ami
**Is your feature request related to a problem? Please describe.**
I have a use case where I want Cloud Custodian to find AMIs that are shared with accounts not in a whitelist, but in many cases, the AMI might be shared with both authorised and unauthorised accounts. When I use Cloud Custodian to remediate this (i.e. `remove-launch-permissions`) the permissions for _all_ accounts are removed, including the accounts that do have valid access to the AMI.

**Describe the solution you'd like**
I'd like the remediation available in Cloud Custodian to remove launch permissions to the AMI for accounts that are _not_ defined in the whitelist; all other accounts (i.e. the ones that _are_ in the whitelist) should not have their launch permissions removed.

**Describe alternatives you've considered**
I haven't tried it yet, but a Lambda function may be useful here to remove individual account IDs from the launch permissions. I'm not sure how much I like the idea of maintaining both a Cloud Custodian policy _and_ a Lambda function just to remove individual account IDs from the launch permissions rather than all of them.


Discussion:

Resolving patch:
diff --git a/c7n/resources/ami.py b/c7n/resources/ami.py
--- a/c7n/resources/ami.py
+++ b/c7n/resources/ami.py
@@ -18,7 +18,7 @@
 import jmespath
 
 from c7n.actions import BaseAction
-from c7n.exceptions import ClientError
+from c7n.exceptions import ClientError, PolicyValidationError
 from c7n.filters import (
     AgeFilter, Filter, CrossAccountAccessFilter)
 from c7n.manager import resources
@@ -162,8 +162,25 @@ class RemoveLaunchPermissions(BaseAction):
 
     """
 
-    schema = type_schema('remove-launch-permissions')
-    permissions = ('ec2:ResetImageAttribute',)
+    schema = type_schema(
+        'remove-launch-permissions',
+        accounts={'oneOf': [
+            {'enum': ['matched']},
+            {'type': 'string', 'minLength': 12, 'maxLength': 12}]})
+
+    permissions = ('ec2:ResetImageAttribute', 'ec2:ModifyImageAttribute',)
+
+    def validate(self):
+        if 'accounts' in self.data and self.data['accounts'] == 'matched':
+            found = False
+            for f in self.manager.iter_filters():
+                if isinstance(f, AmiCrossAccountFilter):
+                    found = True
+                    break
+            if not found:
+                raise PolicyValidationError(
+                    "policy:%s filter:%s with matched requires cross-account filter" % (
+                        self.manager.ctx.policy.name, self.type))
 
     def process(self, images):
         client = local_session(self.manager.session_factory).client('ec2')
@@ -171,8 +188,25 @@ def process(self, images):
             self.process_image(client, i)
 
     def process_image(self, client, image):
-        client.reset_image_attribute(
-            ImageId=image['ImageId'], Attribute="launchPermission")
+        accounts = self.data.get('accounts')
+        if not accounts:
+            return client.reset_image_attribute(
+                ImageId=image['ImageId'], Attribute="launchPermission")
+        if accounts == 'matched':
+            accounts = image.get(AmiCrossAccountFilter.annotation_key)
+        if not accounts:
+            return
+        remove = []
+        if 'all' in accounts:
+            remove.append({'Group': 'all'})
+            accounts.remove('all')
+        remove.extend([{'UserId': a} for a in accounts])
+        if not remove:
+            return
+        client.modify_image_attribute(
+            ImageId=image['ImageId'],
+            LaunchPermission={'Remove': remove},
+            OperationType='remove')
 
 
 @AMI.action_registry.register('copy')
@@ -318,6 +352,7 @@ class AmiCrossAccountFilter(CrossAccountAccessFilter):
         whitelist={'type': 'array', 'items': {'type': 'string'}})
 
     permissions = ('ec2:DescribeImageAttribute',)
+    annotation_key = 'c7n:CrossAccountViolations'
 
     def process_resource_set(self, client, accounts, resource_set):
         results = []
@@ -326,10 +361,11 @@ def process_resource_set(self, client, accounts, resource_set):
                 client.describe_image_attribute,
                 ImageId=r['ImageId'],
                 Attribute='launchPermission')['LaunchPermissions']
+            r['c7n:LaunchPermissions'] = attrs
             image_accounts = {a.get('Group') or a.get('UserId') for a in attrs}
             delta_accounts = image_accounts.difference(accounts)
             if delta_accounts:
-                r['c7n:CrossAccountViolations'] = list(delta_accounts)
+                r[self.annotation_key] = list(delta_accounts)
                 results.append(r)
         return results
 


Test patch:
diff --git a/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_1.json b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_1.json
@@ -0,0 +1,15 @@
+{
+    "status_code": 200,
+    "data": {
+        "ImageId": "ami-03f9a84f3963d5028",
+        "LaunchPermissions": [
+            {
+                "UserId": "112233445566"
+            },
+            {
+                "UserId": "665544332211"
+            }
+        ],
+        "ResponseMetadata": {}
+    }
+}
diff --git a/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_2.json b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_2.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImageAttribute_2.json
@@ -0,0 +1,8 @@
+{
+    "status_code": 200,
+    "data": {
+        "ImageId": "ami-03f9a84f3963d5028",
+        "LaunchPermissions": [],
+        "ResponseMetadata": {}
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImages_1.json b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImages_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_ami_remove_perms/ec2.DescribeImages_1.json
@@ -0,0 +1,40 @@
+{
+    "status_code": 200,
+    "data": {
+        "Images": [
+            {
+                "Architecture": "arm64",
+                "CreationDate": "2020-10-08T13:31:46.000Z",
+                "ImageId": "ami-03f9a84f3963d5028",
+                "ImageLocation": "644160558196/my-image",
+                "ImageType": "machine",
+                "Public": false,
+                "OwnerId": "644160558196",
+                "PlatformDetails": "Linux/UNIX",
+                "UsageOperation": "RunInstances",
+                "State": "available",
+                "BlockDeviceMappings": [
+                    {
+                        "DeviceName": "/dev/sda1",
+                        "Ebs": {
+                            "DeleteOnTermination": true,
+                            "SnapshotId": "snap-03a81fee8252db12c",
+                            "VolumeSize": 8,
+                            "VolumeType": "gp2",
+                            "Encrypted": false
+                        }
+                    }
+                ],
+                "Description": "",
+                "EnaSupport": true,
+                "Hypervisor": "xen",
+                "Name": "my-image",
+                "RootDeviceName": "/dev/sda1",
+                "RootDeviceType": "ebs",
+                "SriovNetSupport": "simple",
+                "VirtualizationType": "hvm"
+            }
+        ],
+        "ResponseMetadata": {}
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_ami_remove_perms/ec2.ModifyImageAttribute_1.json b/tests/data/placebo/test_ami_remove_perms/ec2.ModifyImageAttribute_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_ami_remove_perms/ec2.ModifyImageAttribute_1.json
@@ -0,0 +1,6 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {}
+    }
+}
\ No newline at end of file
diff --git a/tests/test_ami.py b/tests/test_ami.py
--- a/tests/test_ami.py
+++ b/tests/test_ami.py
@@ -37,7 +37,29 @@ def test_query(self):
         resources = p.run()
         self.assertEqual(len(resources), 1)
 
-    def test_ami_ssq(self):
+    def test_ami_remove_launch_permissions(self):
+        factory = self.replay_flight_data('test_ami_remove_perms')
+        p = self.load_policy({
+            'name': 'ami-check',
+            'resource': 'aws.ami',
+            'filters': ['cross-account'],
+            'actions': [{
+                'type': 'remove-launch-permissions',
+                'accounts': 'matched'}]},
+            session_factory=factory)
+        resources = p.run()
+        self.assertEqual(len(resources), 1)
+        self.assertEqual(
+            sorted(resources[0]['c7n:CrossAccountViolations']),
+            ['112233445566', '665544332211'])
+
+        client = factory().client('ec2')
+        perms = client.describe_image_attribute(
+            ImageId=resources[0]['ImageId'],
+            Attribute='launchPermission')['LaunchPermissions']
+        assert perms == []
+
+    def test_ami_sse(self):
         factory = self.replay_flight_data('test_ami_sse')
         p = self.load_policy({
             'name': 'ubuntu-bionic',