Update Launch Templates with new AMI

# Patrick Domack
# Base taken from Orion Anderson https://misterorion.com/lambda-update-ami/
# https://github.com/ranman/awesome-sns/blob/master/Events.md#windows-ami-update
# https://docs.aws.amazon.com/AWSEC2/latest/WindowsGuide/aws-windows-ami.html#subscribe-notifications

import boto3, os, json
from datetime import datetime, timezone, timedelta


def lambda_handler(event, context):
    # Get values from Lambda environment variables.
    launch_templates_env = os.environ.get("launch_template_id")
    if launch_templates_env:
        launch_templates = launch_templates_env.split(" ")
    else:
        launch_templates = None
    sns_arn = os.environ.get("sns_arn")
    asg_name = os.environ.get("asg_name")
    ami_name = os.environ.get("ami_name")
    region = os.environ.get("region")

    sns_message = None
    new_ami = None

    def detect_running_region():
        """Dynamically determine the region from a running Glue job (or anything on EC2 for
        that matter)."""
        easy_checks = [
            # check if set through ENV vars
            os.environ.get('AWS_REGION'),
            os.environ.get('AWS_DEFAULT_REGION'),
            # else check if set in config or in boto already
            boto3.DEFAULT_SESSION.region_name if boto3.DEFAULT_SESSION else None,
            boto3.Session().region_name,
        ]
        for region in easy_checks:
            if region:
                return region
        # else query an external service
        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html
        r = requests.get("http://169.254.169.254/latest/dynamic/instance-identity/document")
        response_json = r.json()
        return response_json.get('region')
    
    if region is None:
        region = detect_running_region()
    
    # Create boto3 clients
    ec2 = boto3.client("ec2", region_name=region)
    asg = boto3.client("autoscaling", region_name=region)
    sns = boto3.client("sns", region_name=region)
    ssm = boto3.client("ssm", region_name=region)

    if event:
        if "Records" in event:
            sns_message = json.loads(event["Records"][0]["Sns"]["Message"])

    if sns_message:
        # Parse the SNS message and get the new image id (ECS)
        if "ECSAmis" in sns_message:
            new_ami = sns_message["ECSAmis"][0]["Regions"][region]["ImageId"]
    
        # Parse the SNS message and get the new image id (Linux)
        if "v1" in sns_message:
            for x in sns_message["v1"]["Regions"][region]:
                if x.name == ami_name:
                    new_ami = x.ImageId
    
    # Parse the SSM Latest Windows AMI
    if new_ami is None:
        response = ssm.get_parameter( Name='/aws/service/ami-windows-latest/' + ami_name )
        if "Parameter" in response:
            new_ami = response["Parameter"]["Value"]

    def update_current_launch_template_ami(launch_template_id, ami):
        response = ec2.create_launch_template_version(
            LaunchTemplateId=launch_template_id,
            SourceVersion="$Latest",
            VersionDescription="Latest-AMI",
            LaunchTemplateData={
                "ImageId": ami
            }
        )
        print(f"New launch template created with AMI {ami}")

    def set_launch_template_default_version(launch_template_id):
        response = ec2.modify_launch_template(
            LaunchTemplateId=launch_template_id,
            DefaultVersion="$Latest"
        )
        print("Default launch template set to $Latest.")
        previous_version = str(
            int(response["LaunchTemplate"]["LatestVersionNumber"]) - 2)
        response = ec2.delete_launch_template_versions(
            LaunchTemplateId=launch_template_id,
            Versions=[
                previous_version,
            ]
        )
        print(f"Old launch template {previous_version} deleted.")

    def create_asg_scheduled_action(start_time, desired_capacity):
        response = asg.put_scheduled_update_group_action(
            AutoScalingGroupName=asg_name,
            ScheduledActionName=f"Desire {desired_capacity}",
            StartTime=start_time,
            DesiredCapacity=desired_capacity
        )
        print(f"""
            ASG action created
            Start time: {start_time}"
            Desired capacity: {desired_capacity}
            """)

    def send_sns_notification(subject, message):
        if sns_arn:
            response = sns.publish(
                TargetArn=sns_arn,
                Message=message,
                Subject=subject,
            )
        print(f"""
            Notification email sent.
            Subject: {subject}
            Message: {message}
            """)

    def update_launch_template_and_asg(launch_template_id):
        # Update template AMI and set as default
        if new_ami:
            update_current_launch_template_ami(launch_template_id,new_ami)
            set_launch_template_default_version(launch_template_id)

            # Create future ASG actions to roll out the new AMI
            if asg_name:
                now_utc = datetime.now(timezone.utc)
                in_01_min = now_utc + timedelta(minutes=1)
                in_15_min = now_utc + timedelta(minutes=15)
                create_asg_scheduled_action(in_01_min, 2)
                create_asg_scheduled_action(in_15_min, 1)

                # Send a notification that the update succeeded.
                subject = "AMI updated, ASG rotated!"
                message = f"AMI updated! New AMI is {new_ami}."
            else:
                subject = "AMI updated!"
                message = f"AMI updated! New AMI is {new_ami}."
        else:
            subject = "AMI not updated!"
            message = f"AMI not defined."
        send_sns_notification(subject, message)
        print(message)
        return message

    ami_status = f"AMI found {new_ami}\n"
    if launch_templates:
        for x in launch_templates:
            if x:
                ami_status += update_launch_template_and_asg(x)
                ami_status += "\n"

    # Show if AMI was updated in Lambda console.
    return ami_status