Fix cloudfront lambda cleanup

RyanJarv · Aug 20, 2021 · 8cd7f48 · 8cd7f48
1 parent 4e17985
commit 8cd7f48
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 16 deletions.
diff --git a/cdn_bypass/cdn.py b/cdn_bypass/cdn.py
diff --git a/cdn_bypass/cloudfront/cli.py b/cdn_bypass/cloudfront/cli.py
@@ -31,7 +31,7 @@ def create(target: str = typer.Argument(..., help='Optional device name to limit
     """
 
     cloudfront = CloudFront(sess, target)
-    with typer.progressbar(cloudfront.create(), length=18, label=f"Creating {target}") as progress:
+    with typer.progressbar(cloudfront.create(), length=20, label=f"Creating {target}") as progress:
         for update in progress:
             progress.update(1)
             progress.label = typer.style(update, fg=typer.colors.CYAN)
@@ -47,7 +47,7 @@ def delete(target: str = typer.Argument(..., help='Optional device name to limit
     """
 
     cloudfront = CloudFront(sess, target)
-    with typer.progressbar(cloudfront.delete(), length=12, label=f"Destroying {id}") as progress:
+    with typer.progressbar(cloudfront.delete(), length=20, label=f"Destroying {id}") as progress:
         for update in progress:
             progress.update(1)
             progress.label = typer.style(update, fg=typer.colors.CYAN)

diff --git a/cdn_bypass/cloudfront/main.py b/cdn_bypass/cloudfront/main.py
@@ -3,7 +3,7 @@
 import re
 import zipfile
 from pathlib import Path
-from time import time
+from time import time, sleep
 
 import botocore.exceptions
 
@@ -140,10 +140,10 @@ def create_function(self, lambda_role_arn):
             source_code = f.read()
 
         # Lambda@Edge functions can't use environment variables so set this as a global.
-        source_code = re.sub(r'HOST\w+=.*$', 'HOST = "{}"'.format(self.target), source_code)
+        source_code = re.sub(r'HOST\s+=.*$', 'HOST = "{}"'.format(self.target), source_code)
 
         if self.x_forwarded_for:
-            source_code = re.sub(r'X_FORWARDED_FOR\w+=.*$', 'X_FORWARDED_FOR = "{}"'.format(self.x_forwarded_for),
+            source_code = re.sub(r'X_FORWARDED_FOR\s+=.*$', 'X_FORWARDED_FOR = "{}"'.format(self.x_forwarded_for),
                                  source_code)
 
         zip = io.BytesIO()
@@ -163,7 +163,6 @@ def create_function(self, lambda_role_arn):
                 raise e
 
         yield f'Lambda@Edge Function {self.lambda_function_name} -- Creating function'
-        import pdb; pdb.set_trace()
         resp = lamb.create_function(
             FunctionName=self.lambda_function_name,
             Runtime='python3.8',
@@ -196,11 +195,26 @@ def delete_function(self):
         lamb = self.sess.client('lambda', region_name='us-east-1')
         yield f'Lambda@Edge Function {self.lambda_function_name} -- Deleting'
 
-        try:
-            lamb.delete_function(FunctionName=self.lambda_function_name)
-        except botocore.exceptions.ClientError as e:
-            if e.response['Error']['Code'] != 'ResourceNotFoundException':
-                raise e
+        # It takes some time after you disassociate a lambda function from a CloudFront distribution before you can
+        # delete it successfully.
+        while True:
+            i = 1
+            try:
+                yield f'Lambda@Edge Function {self.lambda_function_name} -- Deleting (attempt {i}, this may take a ' \
+                      'while)'
+                lamb.delete_function(FunctionName=self.lambda_function_name)
+                break
+            except botocore.exceptions.ClientError as e:
+                if e.response['Error']['Code'] == 'InvalidParameterValueException':
+                    pass
+                elif e.response['Error']['Code'] == 'ResourceNotFoundException':
+                    print("\n[ERROR] Lambda function {} not found.".format(self.lambda_function_name))
+                    return
+                else:
+                    raise e
+
+            i += 1
+            sleep(10)
 
         yield f'Lambda@Edge Function {self.lambda_function_name} -- Deleted'
 
@@ -283,7 +297,7 @@ def create_distribution(self, lambda_arn):
                         'CacheBehaviors': {
                             'Quantity': 0,
                         },
-                        'Comment': 'cdn-bypass todo update description',
+                        'Comment': 'cdn-proxy distribution for {}'.format(self.target),
                         'PriceClass': 'PriceClass_100',
                         'Enabled': True,
                         'ViewerCertificate': {
@@ -331,7 +345,9 @@ def delete_distribution(self, dist_id):
                                         'if so.')
         config = resp['DistributionConfig']
         config['Enabled'] = False
-        config['DefaultCacheBehavior']['LambdaFunctionAssociations'] = {}
+        config['DefaultCacheBehavior']['LambdaFunctionAssociations'] = {
+            "Quantity": 0,
+        }
         yield f'CloudFront Distribution {dist_id} -- Disabling'
         client.update_distribution(Id=dist_id, DistributionConfig=config, IfMatch=resp['ETag'])
         yield f'CloudFront Distribution {dist_id} -- Waiting for propagation (this may take a while)'
@@ -348,6 +364,4 @@ def get_latest_lambda_version(self, lambda_name):
         resp = lamb.list_versions_by_function(FunctionName=lambda_name)
         versions = [v['Version'] for v in resp['Versions']]
         versions.remove('$LATEST')
-        return max(versions)
-
-
+        return max(versions)