provision.py 4.08 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
import boto3
from botocore.exceptions import ClientError
import time

'''
Original incantation for starting and stopping EC2 was from the documentation:
https://aws.amazon.com/premiumsupport/knowledge-center/start-stop-lambda-cloudwatch/

Starting and stopping RDS is from:
https://stackoverflow.com/questions/44738071/boto3-start-stop-rds-instance-with-aws-lambda

Allocating/releasing elastic IP addresses is from:
https://boto3.readthedocs.io/en/latest/guide/ec2-example-elastic-ip-addresses.html
'''


def lambda_handler(event, context):
    '''

    :param event: must contain 'mode' (start|stop), 'region' ('us-east-1'), 'ec2_instance'
    :param context:
    :return:
    '''
    mode = event['mode']
    # the region instances are in without specifying Availability Zone; e.g., 'us-east-1'
    region = event['region']

    #ec2 instance ex. ['X-XXXXXXXX', 'X-XXXXXXXX']
    ec2_instance = event['ec2_instance']
    ec2 = boto3.client('ec2', region_name=region)

    rds_instance = event['rds_instance']
    rds = boto3.client('rds')

    if mode == 'start':
        start(ec2, ec2_instance, rds, rds_instance)
    elif mode == 'stop':
        stop(ec2, ec2_instance, rds, rds_instance)
    else:
        print("Must specify a mode.")
        exit(1)

    return 0

def start(ec2, ec2_instance, rds, rds_instance):
    print("Starting up soup")

    print(f"Starting ec2 instance {ec2_instance}")
    ec2.start_instances(InstanceIds=[ec2_instance])
    print(f"Started ec2 instance {ec2_instance}")

    print(f"Starting rds instance {rds_instance}")
    rds.start_db_instance(DBInstanceIdentifier=rds_instance)
    print(f"Starting rds instance {rds_instance}")

    print(f"Allocating EIP address to EC2 instance {ec2_instance}")
    try:
        wait_ec2_target_state(ec2, ec2_instance, "running")

        allocation = ec2.allocate_address(Domain='vpc')
        response = ec2.associate_address(AllocationId=allocation['AllocationId'],
                                         InstanceId=ec2_instance)
        print(response)
        print(f"Allocated EIP address {allocation} to EC2 instance {ec2_instance}")
    except ClientError as e:
        print(e)
        exit(1)

    print("successfully finished starting up soup")

def stop(ec2, ec2_instance, rds, rds_instance):
    print("stopping soup")

    print(f"Stopping ec2 instance {ec2_instance}")
    ec2.stop_instances(InstanceIds=[ec2_instance])
    print(f"Stopped ec2 instance {ec2_instance}")

    wait_ec2_target_state(ec2, ec2_instance, "stopped")

    print(f"Releasing EIP address from EC2 instance {ec2_instance}")
    try:
        allocations = ec2.describe_addresses()
        if len(allocations['Addresses']) != 1:
            print("wrong number of ip allocations returned")
            exit(1)
        association_id = allocations['Addresses'][0]['AssociationId']
        ec2.disassociate_address(AssociationId=association_id)
        allocation_id = allocations['Addresses'][0]['AllocationId']
        ec2.release_address(AllocationId=allocation_id)
        print(f"Released EIP address {allocation_id} from EC2 instance {ec2_instance}")
    except ClientError as e:
        print(e)
        exit(1)

    print(f"Stopping rds instance {rds_instance}")
    rds.stop_db_instance(DBInstanceIdentifier=rds_instance)
    print(f"Stopped rds instance {rds_instance}")

    print("successfully finished stopping soup")

def wait_ec2_target_state(ec2, ec2_instance, target_state):
    for i in range(1, 11):
        print(f"Attempt {i} checking ec2 state")
        response = ec2.describe_instance_status(InstanceIds=[ec2_instance], IncludeAllInstances=True)
        instance_states = response['InstanceStatuses']
        if len(instance_states) != 1:
            print("wrong number of instances returned")
            exit(1)
        state = instance_states[0]['InstanceState']['Name']
        id = instance_states[0]['InstanceId']
        print(f"Instance {id} detected as {state}")
        if id == ec2_instance and state == target_state:
            break
        time.sleep(5.0)
    else:
        print(f"ec2 instance {ec2_instance} did not reach target state {target_state}")
        exit(1)