Update required ami test logic (#239)

* Update required ami test logic

* style: ran black

* make options in new ami test configurable

* rename test file based on new test name

* small style fix

aws/ec2/resources.py

@@ -135,3 +135,18 @@ def ec2_security_groups_with_in_use_flag():
             sec_group["InUse"] = False
 
     return sec_groups
+
+
+def ec2_images_owned_by(account_ids):
+    "Returns a list of EC2 images owned by a list of provided account ids"
+    return (
+        botocore_client.get(
+            "ec2",
+            "describe_images",
+            [],
+            {"Filters": [{"Name": "owner-id", "Values": account_ids}]},
+        )
+        .extract_key("Images")
+        .flatten()
+        .values()
+    )

aws/ec2/test_ec2_instance_on_acceptable_ami.py

@@ -0,0 +1,46 @@
+import pytest
+
+from aws.ec2.helpers import ec2_instance_test_id
+from aws.ec2.resources import ec2_instances, ec2_images_owned_by
+
+from datetime import datetime, timedelta
+
+
+@pytest.fixture
+def owned_amis(pytestconfig):
+    return ec2_images_owned_by(pytestconfig.custom_config.aws.owned_ami_account_ids)
+
+
+@pytest.fixture
+def max_ami_age(pytestconfig):
+    return pytestconfig.custom_config.aws.max_ami_age_in_days
+
+
+@pytest.mark.ec2
+@pytest.mark.parametrize("ec2_instance", ec2_instances(), ids=ec2_instance_test_id)
+def test_ec2_instance_on_acceptable_ami(ec2_instance, owned_amis, max_ami_age):
+    """
+    Checks that all EC2 instances are running on acceptable AMIs, meaning
+    an AMI that is not older than X days and is owned by us.
+    Default is 180 days.
+    """
+    for tag in ec2_instance["Tags"]:
+        if tag["Key"] == "Name":
+            instanceName = tag["Value"]
+
+    minAge = datetime.now() - timedelta(days=max_ami_age)
+    foundAmi = False
+    for ami in owned_amis:
+        if ami["ImageId"] == ec2_instance["ImageId"]:
+            assert (
+                ami["CreationDate"] > minAge
+            ), "Instance {} {} is running on an AMI created on {} that's older than 180 days".format(
+                instanceName, ec2_instance["InstanceId"], ami["CreationDate"]
+            )
+            foundAmi = True
+            break
+
+    if not foundAmi:
+        assert False, "Instance {} {} uses AMI {} not owned by us".format(
+            instanceName, ec2_instance["InstanceId"], ec2_instance["ImageId"]
+        )

config.yaml.example

@@ -40,16 +40,16 @@ aws:
     - Type
     - App
     - Env
-  required_amis:
-    - ami-00000000000000000
-    - ami-55555555555555555
   whitelisted_ports_global:
     - 25
   whitelisted_ports:
     - test_param_id: '*bastion'
       ports:
         - 22
         - 2222
+  max_ami_age_in_days: 90
+  owned_ami_account_ids:
+    - 1234567890
 gsuite:
   domain: 'example.com'
   min_number_of_owners: 2

custom_config.py

@@ -67,6 +67,10 @@ def __init__(self, config):
         self.access_key_expires_after = config.get("access_key_expires_after", None)
         self.admin_policies = frozenset(config.get("admin_policies", []))
         self.admin_groups = frozenset(config.get("admin_groups", []))
+        self.owned_ami_account_ids = [
+            str(x) for x in config.get("owned_ami_account_ids", [])
+        ]
+        self.max_ami_age_in_days = config.get("max_ami_age_in_days", 180)
         super().__init__(config)
 
     def get_whitelisted_ports(self, test_id):