#
# AWS Lambda Python script to automatically create AMIs
#

import boto3
import botocore
import json
import datetime
import dateutil.parser

print('Loading function')

# Constants
RETENTION_DAYS = 7
TAG_KEY = 'automatic-backup'

def lambda_handler(event, context):
    # print("Received event: " + json.dumps(event, indent=2))
    print("\nAWS instance backups stated at %s...\n" % datetime.datetime.now())

    if 'retention_days' in event:
        retention_days = event['retention_days']

    ec2 = boto3.resource('ec2')

    # filter instances that are currently running and have a tag 'backup=true'
    instances = list(ec2.instances.filter(
        Filters = [
            {'Name': 'instance-state-name', 'Values': ['running']},
            {'Name': 'tag-key', 'Values': [TAG_KEY]}
        ]
    ))

    if len(instances) == 0:
        print "Did not find any instance matching the creteria"
        
    for i in instances:
        instance_name = filter(lambda tag: tag['Key'] == 'Name', i.tags)[0]['Value']
        print("Creating image for Instance %s: %s..." % (i.id, instance_name))

        try:
            # loop through instances and create AMIs
            image = i.create_image(
                Name = ("%s - generated on %s" % (instance_name, datetime.date.today())),
                NoReboot = True
            )
            
            # Add tag to identify AMI that were autogenerated
            image.create_tags(Tags = [{'Key': TAG_KEY,'Value': ''}])
            print("OK")
        except botocore.exceptions.ClientError as e:
            print("AMI creation failed: %s" % e.response['Error']['Message'])

    # Look for available AMIs owned by us and that were auto-generated 
    images = list(ec2.images.filter(
        Owners=['self'],
        Filters = [
            {'Name': 'tag-key', 'Values': [TAG_KEY]},
            {'Name': 'state', 'Values': ['available']}
        ]
    ))
    
    if len(instances) == 0:
        print "Did not find any AMI matching the creteria"
    
    # Check if AMI is older than retention_days
    for ami in images:
        ami_creation_date = dateutil.parser.parse(ami.creation_date).replace(tzinfo=None)
        limit_date = datetime.datetime.now() - datetime.timedelta(days=retention_days)
        
        if ami_creation_date <= limit_date:
            print("Deregistering AMI %s with name '%s'" % (ami.id, ami.name))
            # deregister AMI
            ami.deregister()
        
    print("\nAWS instance backups completed at %s\n" % datetime.datetime.now())
    
    return True
