style: minor changes

lambda/ec2-scheduler/lambda_function.py

@@ -4,13 +4,14 @@ import json
 
 # reference: https://aws.amazon.com/premiumsupport/knowledge-center/start-stop-lambda-eventbridge/
 
-ec2 = boto3.client('ec2', region_name=os.environ['region_name'])
-
 def lambda_handler(event, context):
-    if (event['action'] == 'start'):
+    ec2 = boto3.client('ec2', region_name=os.environ['region_name'])
-        resp = ec2.start_instances(InstanceIds=json.loads(os.environ['instances']))
+    instances = json.loads(os.environ['instances'])
-    elif (event['action'] == 'stop'):
+    if event['action'] == 'start':
-        resp = ec2.stop_instances(InstanceIds=json.loads(os.environ['instances']))
+        resp = ec2.start_instances(InstanceIds=instances)
+    elif event['action'] == 'stop':
+        resp = ec2.stop_instances(InstanceIds=instances)
     else:
-        resp = "Event action not provided"
+        raise ValueError("Invalid event action")
     return resp
+

lambda/lambda-cf-wafipset/main.py

@@ -0,0 +1,71 @@
+#!/usr/bin/python3
+r"""
+Documentation
+
+License: This program is released under the MIT License
+"""
+
+# Imports
+import boto3
+
+
+class AwsPrefixList:
+    def __init__(self):
+        ec2 = boto3.client('ec2')
+        response = ec2.describe_managed_prefix_lists(
+            Filters=[{'Name': 'prefix-list-name', 'Values': ['com.amazonaws.global.cloudfront.origin-facing']}]
+        )
+        prefix_lists = response.get('PrefixLists', [])
+        prefix_list_id = prefix_lists[0]['PrefixListId']
+        entries = []
+        paginator = ec2.get_paginator('get_managed_prefix_list_entries')
+        for page in paginator.paginate(PrefixListId=prefix_list_id):
+            entries.extend(page.get('Entries', []))
+
+        self.pl = [entry['Cidr'] for entry in entries]
+        self.pl.sort()
+
+    def getHash(self):
+        return hash(tuple(self.pl))
+
+    def getTuple(self):
+        return tuple(self.pl)
+
+    def getLength(self):
+        return len(self.pl)
+
+class WafIpSet:
+    def __init__(self, name: str, id: str):
+        waf_client = boto3.client('wafv2')
+        temp = waf_client.get_ip_set(
+            Name=name,
+            Scope='REGIONAL',
+            Id=id)
+        self.ip_set = temp["IPSet"]["Addresses"]
+        self.ip_set.sort()
+        self.lock_token = temp['LockToken'] # need this to update ipset
+
+    def getHash(self):
+        return hash(tuple(self.ip_set))
+
+    def getTuple(self):
+        return tuple(self.ip_set)
+
+    def getLength(self):
+        return len(self.ip_set)
+
+# Main function
+def main() -> None:
+    pl = AwsPrefixList()
+    ipset = WafIpSet(name="cloudfront_ip_ipset", id="951120be-31d7-415f-9aa3-5ad9e56b6195")
+    print(f"PrefixList length: {pl.getLength()}")
+    print(f"IpSet length: {ipset.getLength()}")
+    # missing = set(pl.getTuple()) - set(ipset.getTuple())
+    # notInPl = set(ipset.getTuple()) - set(pl.getTuple())
+    # print(f"Missing in WAF ipset: {len(missing)}")
+    # print(f"Not in PL: {len(notInPl)}")
+
+
+# Call main function
+if __name__ == '__main__':
+    main()

lambda/secret-rotation-reminder/SecretRotationReminder.py

@@ -56,7 +56,7 @@ def lambda_handler(event, context):
     step = event['Step']
 
     # Secretsmanager sends 4 rotation events, but we will only use the createSecret event
-    # and send reminder out
+    # and the finishSecret event
     if step == "createSecret":
         # send notification and create a new secret from existing secret
         send_notification(secret_id, token)
@@ -64,11 +64,11 @@ def lambda_handler(event, context):
         # set new secret with version AWSCURRENT
         swap_current_version(secret_id, token)
     else:
-        print("Ignore step " + step)
+        print(f"Steps other than createSecret and finishSecret will be ignored: {step}")
     return True
 
-def send_notification(secret_id, token):
+def send_notification(secret_id: str, token: str) -> None:
-    print("Clone secret and send notification for", secret_id)
+    print(f"Clone secret and send notification for {secret_id}")
     sm_client = boto3.client('secretsmanager')
     """
     A new secret version is required by rotation workflow
@@ -90,11 +90,11 @@ def send_notification(secret_id, token):
     sns_client = boto3.client('sns')
     sns_client.publish(
         TopicArn=SNS_TOPIC_ARN,
-        Message='Please rotate the secret ' + secret_id + '\n\nThis message is generated by lambda function SecretRotationReminder',
+        Message=f'Your secret {secret_id} is due for update. Please change it on secretsmanager and on your applications.',
         Subject='Secret rotation reminder for ' + secret_id.split(":")[6]
     )
 
-def swap_current_version(secret_id, token):
+def swap_current_version(secret_id: str, token: str) -> None:
     print("Point AWSCURRENT to new secret version")
     sm_client = boto3.client('secretsmanager')
     metadata = sm_client.describe_secret(SecretId=secret_id)
@@ -116,4 +116,3 @@ def swap_current_version(secret_id, token):
         VersionStage='AWSPENDING',
         RemoveFromVersionId=token
     )
-    return True