check also lambda for running images awslabs#40

Author: idan-miara

README.md

@@ -1,5 +1,11 @@
 # Automated Image Cleanup for Amazon ECR
-The Python script and Lambda function described here help clean up images in [Amazon ECR](https://aws.amazon.com/ecr). The script looks for images that are not used in running [Amazon ECS](https://aws.amazon.com/ecs) tasks that can be deleted. You can configure the script to print the image list first to confirm deletions, specify a region, or specify a number of images to keep for potential rollbacks.
+The Python script and Lambda function described here help clean up images in [Amazon ECR](https://aws.amazon.com/ecr).
+The script looks for images that are not used in running 
+[Amazon ECS](https://aws.amazon.com/ecs) tasks, 
+[Amazon EKS](https://aws.amazon.com/eks) tasks and 
+[Amazon Lambda](https://aws.amazon.com/lambda) container images that can be deleted.
+You can configure the script to print the image list first to confirm deletions, specify a region, 
+or specify a number of images to keep for potential rollbacks.
 
 ## Authenticate with AWS
 [Configuring the AWS Command Line Interface.](http://docs.aws.amazon.com/cli/latest/userguide/cli-chap-getting-started.html)

main.py

@@ -43,6 +43,7 @@ class Inputs:
     protect_tags_regex: re.Pattern | None = None
     protect_latest: bool | None = None
     check_ecs: bool | None = None
+    check_lambda: bool | None = None
     check_eks: bool | None = None
     connect_timeout: int | None = None
     read_timeout: int | None = None
@@ -105,6 +106,12 @@ def __post_init__(self):
             else:
                 raise Exception(f"Unexpected {type(self.check_ecs)=}")
 
+        if not isinstance(self.check_lambda, bool):
+            if isinstance(self.check_lambda, str):
+                self.check_lambda = (self.check_lambda or "false").lower() != 'false'
+            else:
+                raise Exception(f"Unexpected {type(self.check_lambda)=}")
+
         if not isinstance(self.check_eks, bool):
             if isinstance(self.check_eks, str):
                 self.check_eks = (self.dryrun or "false").lower() != 'false'
@@ -126,6 +133,7 @@ def from_env(cls):
             protect_tags_regex=os.environ.get('PROTECT_TAGS_REGEX'),
             protect_latest=os.environ.get('PROTECT_LATEST'),
             check_ecs=os.environ.get('CHECK_ECS'),
+            check_lambda=os.environ.get('CHECK_LAMBDA'),
             check_eks=os.environ.get('CHECK_EKS'),
             connect_timeout=os.environ.get('CONNECT_TIMEOUT'),
             read_timeout=os.environ.get('READ_TIMEOUT'),
@@ -165,6 +173,8 @@ def from_parser(cls):
                             dest="protect_latest", action='store_false')
         parser.add_argument('-no-ecs', help="Don't search ECS for running images",
                             dest="check_ecs", action='store_false')
+        parser.add_argument('-no-lambda', help="Don't search Lambda for running images",
+                            dest="check_lambda", action='store_false')
         parser.add_argument('-no-eks', help="Don't search EKS for running images",
                             dest="check_eks", action='store_false')
 
@@ -185,6 +195,7 @@ def from_parser(cls):
             protect_tags_regex=args.protect_tags_regex,
             protect_latest=args.protect_latest,
             check_ecs=args.check_ecs,
+            check_lambda=args.check_lambda,
             check_eks=args.check_eks,
             connect_timeout=args.connect_timeout,
             read_timeout=args.read_timeout,
@@ -432,6 +443,14 @@ def get_running_containers(ecs_client, inputs: Inputs):
         running_containers_ecs = []
         logger.info('Skip checking running containers on ECS...')
 
+    if inputs.check_lambda:
+        running_containers_lambda = get_running_containers_from_lambda(inputs.region, logger)
+        logger.info(f"{len(running_containers_lambda)} running containers "
+                    f"from Lambda: {sorted(running_containers_lambda)}...")
+    else:
+        running_containers_lambda = []
+        logger.info('Skip checking running containers on Lambda...')
+
     if inputs.check_eks:
         running_containers_eks = get_running_containers_from_eks(logger)
         logger.info(f"{len(running_containers_eks)} running containers "
@@ -440,11 +459,25 @@ def get_running_containers(ecs_client, inputs: Inputs):
         running_containers_eks = []
         logger.info('Skip checking running containers on EKS...')
 
-    running_containers_ecs = sorted(running_containers_ecs | running_containers_eks)
+    running_containers_ecs = sorted(running_containers_ecs | running_containers_eks | running_containers_lambda)
     logger.info(f"{len(running_containers_ecs)} total running containers: {running_containers_ecs}...")
     return running_containers_ecs
 
 
+def get_running_containers_from_lambda(region_name: str, logger: logging.Logger):
+    running_containers = set()
+    lambda_client = boto3.client('lambda', region_name=region_name)
+    listlambdas_paginator = lambda_client.get_paginator('list_functions')
+    for response_listlambdapaginator in listlambdas_paginator.paginate():
+        for function in response_listlambdapaginator['Functions']:
+            if function["PackageType"] == "Image":
+                function_config = lambda_client.get_function(FunctionName=function["FunctionName"])
+                if function_config["Code"]["ImageUri"] not in running_containers:
+                    running_containers.add(function_config["Code"]["ImageUri"])
+
+    return running_containers
+
+
 def get_running_containers_from_ecs(ecs_client, logger: logging.Logger):
     running_containers = set()
     list_clusters_paginator = ecs_client.get_paginator('list_clusters')