#! /usr/bin/python
################################################################################
## instance-action
## Written by N2W Software
## Date: May 2014
## License: You can use/modify/circulate or do whatever you want.
##          Just note that this script is given "As Is" without any warranty
##
## Usage: see README file
##
################################################################################
import sys
import time
import argparse
from boto import ec2
from boto.exception import EC2ResponseError

AWS_ACCESS_KEY = None
AWS_SECRET_KEY = None
AWS_REGION = u'us-east-1'

def confirm_action (action, instance, region):
    """
    Confirms action with user
    """
    
    if action=='start': msg = 'Confirm starting instance %s in region %s [yes/no] ' % (instance, region)
    elif action=='stop': msg = 'Confirm stopping instance %s in region %s [yes/no] ' % (instance, region)
    elif action=='terminate_with_vols': msg = 'Confirm terminating instance %s (region: %s) and deleting of all attached volumes [yes/no] ' % \
                                                (instance, region)
    elif action=='terminate': msg = 'Confirm terminating instance %s (region: %s)[yes/no] ' % (instance, region)
    else:
        sys.stderr.write ('Unknown action %s' % action)
        return False
    
    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 wait_for_state (instance, target_state):
    """
    Waits for instance to move to desired state
    """

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

        
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

    wait_for_state (instance, u'running')
        
    sys.stdout.write ('Successfully started 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

    wait_for_state (instance, u'stopped')
        
    sys.stdout.write ('Successfully stopped instance %s\n' % instance_id)
    return True
    
def terminate_with_vols_instance (conn, instance_id):
    """
    Terminates instance with all attached volumes
    """
    return terminate_instance (conn, instance_id, True)
    
def terminate_instance (conn, instance_id, delete_vols=False):
    """
    Gets a connection object and instance id
    Kills an instance, meaning terminating it and, if asked for, making sure all attached volumes
    are deleted
    """
    if delete_vols:
        sys.stdout.write ('Killing instance %s and all its attached volumes\n' % instance_id)
    else:
        sys.stdout.write ('Terminating instance %s\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
    if delete_vols:
        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
    
    wait_for_state (instance, u'terminated')
        
    # deletes extra volumes
    if delete_vols:
        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 (10)

        
    if delete_vols:
        sys.stdout.write ('Successfully terminated instance %s with all attached volumes\n' % instance_id)    
    else:
        sys.stdout.write ('Successfully terminated instance %s\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', 'stop', 'terminate', 'terminate_with_vols'])
    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 (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)
        # Calls the relevant function according to action
        retval = locals()[args.action+'_instance'](conn, args.instance)
    else:
        retval = True
        
    if retval:
        sys.exit (0)
    else:
        sys.exit (1)