## Usage: see README file
##
#####################################################################
import sys
import time
import argparse
from boto import ec2
from boto.exception import EC2ResponseError

AWS_ACCESS_KEY = ''
AWS_SECRET_KEY = ''
AWS_REGION = u'us-west-1'

def confirm_action (msg):
    """
    Confirms action with user
    """

    valid = {'yes':True, 'y':True,
             'no':False, 'n':False}

    while True:
        sys.stdout.write (msg+' ')
        choice = raw_input().lower()
        if choice in valid.keys ():
            return valid[choice]
        sys.stdout.write ('Please respond with \'yes\' or \'no\' (or \'y\' or \'n\').\n')

def ec2_connect (access_key, secret_key, region):
    """
    Connects to EC2, returns a connection object
    """

    try:
        conn = ec2.connect_to_region (region,
                                      aws_access_key_id=access_key,
                                      aws_secret_access_key=secret_key)
    except Exception, e:
        sys.stderr.write ('Could not connect to region: %s. Exception: %s\n' % (region, e))
        conn = None

    return conn


def start_instance (conn, instance_id):
    """
    Gets a connection object and instance id
    Starts the instance and waits until it's in 'running' state
    """

    sys.stdout.write ('Starting instance %s\n' % instance_id)
    try:
        instance = conn.start_instances (instance_ids=[instance_id])[0]
    except EC2ResponseError, e:
        sys.stderr.write ('Could not start instance %s. Error: %s\n' % (instance_id, e.error_message))
        return False

    status = instance.update ()
    while status != u'running':
        sys.stdout.write ('Waiting for instance %s to be in \'running\' state\n' % instance_id)
        time.sleep (5)
        status = instance.update ()

    sys.stdout.write ('Successfully started instance %s\n' % instance_id)
    return True

def reboot_instance (conn, instance_id):
    """
    Gets a connection object and instance id
    Reboots the instance and waits until it's in 'running' state
    """

    sys.stdout.write ('Rebooting instance %s\n' % instance_id)
    try:
        instance = conn.reboot_instances (instance_ids=[instance_id])[0]
    except EC2ResponseError, e:
        sys.stderr.write ('Could not re-reboot instance %s. Error: %s\n' % (instance_id, e.error_message))
        return False

    status = instance.update ()
    while status != u'running':
        sys.stdout.write ('Waiting for instance %s to be in \'running\' state\n' % instance_id)
        time.sleep (5)
        status = instance.update ()

    sys.stdout.write ('Successfully rebooted instance %s\n' % instance_id)
    return True

def stop_instance (conn, instance_id):
    """
    Gets a connection object and instance id
    Stops the instance and waits until it's in 'stopped' state
    """

    sys.stdout.write ('Stopping instance %s\n' % instance_id)
    try:
        instance = conn.stop_instances (instance_ids=[instance_id])[0]
    except EC2ResponseError, e:
        sys.stderr.write ('Could not stop instance %s. Error: %s\n' % (instance_id, e.error_message))
        return False

    status = instance.update ()
    while status != u'stopped':
        sys.stdout.write ('Waiting for instance %s to be in \'stopped\' state\n' % instance_id)
        time.sleep (5)
        status = instance.update ()

    sys.stdout.write ('Successfully stopped instance %s\n' % instance_id)
    return True

def kill_instance (conn, instance_id):
    """
    Gets a connection object and instance id
    Kills an instancemeaning terminating it and making sure all attached volumes
    are deleted
    """

    sys.stdout.write ('Killing instance %s and all its attached volumes\n' % instance_id)
    # gets instance
    try:
        instance = conn.get_all_instances (instance_ids=[instance_id])[0].instances[0]
    except EC2ResponseError, e:
        sys.stderr.write ('Could not kill instance %s. Error: %s\n' % (instance_id, e.error_message))
        return False

    # find block devices
    vols_to_delete = []
    for bd in instance.block_device_mapping.values():
        if bd.delete_on_termination == False:
            vols_to_delete.append (bd.volume_id)

    # terminate instance
    try:
        conn.terminate_instances (instance_ids=[instance_id])
    except EC2ResponseError, e:
        sys.stderr.write ('Could not kill instance %s. Error: %s\n' % (instance_id, e.error_message))
        return False

    # waits for termination
    status = instance.update ()
    while status != u'terminated':
        sys.stdout.write ('Waiting for instance %s to be in \'terminated\' state\n' % instance_id)
        time.sleep (5)
        status = instance.update ()

    # deletes extra volumes
    first_vol_delete = True
    for vol in vols_to_delete:
        for try_num in 1,2:
            sys.stdout.write ('Deleting attached volume %s, try number %d\n' % (vol, try_num))
            try:
                conn.delete_volume (vol)
                break
            except EC2ResponseError, e:
                sys.stderr.write ('Could not delete attached volume %s. Error: %s\n' % (vol, e.error_message))
                if try_num == 1:
                    time.sleep (15)


    sys.stdout.write ('Successfully killed instance %s with all attached volumes\n' % instance_id)
    return True


if __name__ == '__main__':

    # Define command line argument parser
    parser = argparse.ArgumentParser(description='Performs actions on an EC2 instance.')
    parser.add_argument('action', choices=['start', 'reboot', 'stop', 'kill'])
    parser.add_argument('--instance', required=True, help='EC2 instance ID to perform the action on')
    parser.add_argument('--region', default = AWS_REGION, help='AWS region the instance is in. If missing default is used')
    parser.add_argument('--access_key', default = AWS_ACCESS_KEY, help='AWS API access key.  If missing default is used')
    parser.add_argument('--secret_key', default = AWS_SECRET_KEY, help='AWS API secret key.  If missing default is used')
    parser.add_argument('--force', action='store_true', help='If set, this flag will prevent confirmation for the action')

    args = parser.parse_args ()

    if args.force:
        perform = True
    else:
        perform = confirm_action ('Confirm %s of instance %s in region %s [yes/no]' % \
                                    (args.action, args.instance, args.region))

    if perform:
        conn = ec2_connect (args.access_key, args.secret_key, args.region)
        if conn == None:
            sys.exit (1)
        retval = locals()[args.action+'_instance'](conn, args.instance)
    else:
        retval = True

    if retval:
        sys.exit (0)
    else:
        sys.exit (1)