spotinst-cli

#!/usr/bin/env python

# This tool allows the user to interact with spotinst API
# run the script with the -h flag for help

from json import loads, dumps
from collections import OrderedDict
from optparse import OptionParser
from os import system, environ
from sys import argv, stdout
from os import path
from base64 import b64encode, b64decode
from commands import getoutput, getstatusoutput


# This function prints error and exits
def errorAndExit(exception_message):
    print "\nERROR - %s" %(exception_message)
    exit()

# This function checks if a package is installed and prints a message if necessary
def check_package_is_installed(package_name):
    if not package_name in packages:
        errorAndExit("{} is not installed, please run 'sudo pip install {}' and try again.".format(package_name,package_name))


# This function changes the color of the text in the terminal
def change_prompt_color(which_color):
    if which_color == 'red':
        stdout.write("\033[1;31m")
    elif which_color == 'green':
        stdout.write("\033[0;32m")
    elif which_color == 'normal':
        stdout.write("\033[0;0m")


# This function prints a message in green/red and returns back to white text
def print_in_color(color, msg):
    change_prompt_color(color)
    print msg
    change_prompt_color('normal')


# This function prints a group name surrounded with =
def print_header(text):
    text_len = len(text)
    print "\n" + ("=" * text_len)
    print text
    print ("=" * text_len)


# This function queries spotinst API endpoint
def query_api(api_url, spotinst_token, req_method='GET', req_payload=None):
    try:
        headers = {
            'Content-Type': 'application/json',
            'Authorization': 'Bearer %s' % spotinst_token,
        }
        if req_method == 'GET':
            return get(api_url, headers=headers)
        elif req_method == 'PUT':
            return put(api_url, headers=headers, data=req_payload)
        elif req_method == 'POST':
            return post(api_url, headers=headers, data=req_payload)
        elif req_method == 'DELETE':
            return delete(api_url, headers=headers, data=req_payload)
    except Exception as e:
        errorAndExit(e.message)


# This function returns a groups dict, filtered by name (contains) or all if parameter is empty
def get_groups(grep_list, ungrep_list):
    print "\nQuerying spotinst API, hold on..."
    result  = query_api(base_path, token)
    data = loads(result.content)
    groups_dict = {}
    all_groups = data['response']['items']
    for group in all_groups:
        # check if the spot_group_name parameter is not empty
        # is so, need to return only the groups that matches the search
        group_name = group['name'].lower()
        if [n for n in grep_list if n in group_name] and not [n for n in ungrep_list if n in group_name]:
            groups_dict[group['name']] = [group['id'], group['capacity']['minimum'], group['capacity']['target'], group['capacity']['maximum'], group['compute']['launchSpecification']['imageId']]
    # sort the list by name
    sorted_inst = OrderedDict(sorted(groups_dict.items()))
    return sorted_inst

#exact matching
def get_groups_exact(grep_list, ungrep_list):
    print "\nQuerying spotinst API, hold on..."
    result = query_api(base_path, token)
    data = loads(result.content)
    groups_dict = {}
    all_groups = data['response']['items']
    for group in all_groups:
        # check if the spot_group_name parameter is not empty
        # is so, need to return only the groups that matches the search
        group_name = group['name'].lower()
        group_name_filter = group_name.split(".")[0]
        for grep_name in grep_list:
            if (group_name_filter == grep_name) and not [n for n in ungrep_list if n in group_name]:
                groups_dict[group['name']] = [group['id'], group['capacity']['minimum'], group['capacity']['target'], group['capacity']['maximum']]
    # sort the list by name
    sorted_inst = OrderedDict(sorted(groups_dict.items()))
    return sorted_inst
# This function prints groups as a table
def print_all_groups(groups_to_print):
    table = PrettyTable(['#' ,'Group name', 'ID', 'Min', 'Target', 'Max', 'AMI ID'])
    table.align = "l"
    table.border = True
    counter = 1
    for key, value in groups_to_print.iteritems():
        table.add_row([counter, key, value[0], value[1], value[2], value[3], value[4]])
        counter += 1
    print_message("Found the following groups: ")

    print table


# This function asks the user to choose a specific group from a list
def get_specific_group(grep_term, ungrep_term):
    groups_to_update = get_groups(grep_term, ungrep_term)
    # check if filter has more than one list and ask the user to choose a group
    if len(groups_to_update) > 1 and options.list is None:
        print_all_groups(groups_to_update)
        ans = raw_input("\nWhich group would you like to update?\nGroup number: ")
        try:
            group_name = list(groups_to_update)[int(ans) - 1]
            group_values = groups_to_update[group_name]
        except:
            errorAndExit("no such group exists!")
    else:
        group_name = list(groups_to_update)[0]
        group_values = groups_to_update[group_name]
    return [group_name, group_values]


# This function calculates the percentage change between two given numbers
def calc_percentage_change(old_val, new_val):
    percentage_change = 0
    try:
        if old_val != 0:
            percentage_change = ((int(new_val) - int(old_val)) / float(old_val)) * 100
        return abs(percentage_change)
    except Exception as e:
        return percent_alert_threshold

# This function checks for changes in the group target and alerts the user
def check_change_in_target(group_name, current_group_target, desired_group_target, alert_threshold):
    target_change = calc_percentage_change(current_group_target, desired_group_target)
    ans_change = "y"
    if int(target_change) >= alert_threshold and not options.skip_validation:
        print "{} target value is: {}, requested target is: {}.".format(group_name, current_group_target, desired_group_target)
        ans_change = raw_input("This represent a change of over {}%, are you sure you want to proceed? [y/n] ".format(alert_threshold))
    return ans_change


# This function prints a message surrounded by #
def print_message(msg):
    all_lines = msg.splitlines()
    longest_line = 0
    for line in all_lines:
        if len(line.strip()) > longest_line:
            longest_line = len(line.strip())
    padding = 8
    print "\n" + ("#" * (longest_line + padding))
    for line in msg.splitlines():
        print (" " * (padding / 2)) + line
    print ("#" * (longest_line + padding)) + "\n"


# This function makes sure the user wants to process with his choice
def user_make_sure(question):
    ans_from_user = raw_input(question)
    if ans_from_user !="y":
        print "\nExiting.."
        exit()


# function to get y or n from the user and return True/ False accordingly
def get_yes_no_from_user(question):
    ans_from_user = ""
    while ans_from_user != "y" and ans_from_user != "n":
        ans_from_user = raw_input(question)
        if ans_from_user != "y" and ans_from_user != "n":
            print "please answer 'y' or 'n' only!"
    if ans_from_user == "y":
        return True
    else:
        return False

# function to split parsed option to list by comma
def parser_split(option, opt, value, parser):
    setattr(parser.values, option.dest, value.split(','))

##### MAIN ####
# check required pip packages
pip_installed = False
packages = [pkg.split('==')[0] for pkg in getoutput('pip freeze').split('\n')]
required_packages = ['prettytable', 'requests']
pip_version = getstatusoutput('pip -V')
if pip_version[0] != 0:
    errorAndExit("python-pip is not installed, please install it before using this tool!\n"
        "On MacOS, you can install it using one of the following:\n"
        "'brew install python' or 'sudo easy_install pip")
if 'python 2.7' not in pip_version[1]:
    errorAndExit("This script currently supports python 2.7 only!")

check_package_is_installed("prettytable")
check_package_is_installed("requests")
from prettytable import PrettyTable
from requests import put, get, post, delete

# define help and options
usage = "Usage: %prog [options]"
parser = OptionParser(usage)
parser.add_option("-g", "--grep", type="string", action="callback", dest="grep", callback=parser_split, default="", help="text to filter groups by")
parser.add_option("-d", "--get-data", action="store_true", dest="data", default="", help="get groups data")
parser.add_option("-s", "--get-status", action="store_true", dest="status", default="", help="get groups status")
parser.add_option("--suspension-status", action="store_true", dest="suspension", default="", help="get groups suspension status")
parser.add_option("-u", "--ungrep", type="string", action="callback", default="", callback=parser_split, dest="ungrep", help="text to exclude groups")
parser.add_option("-l", "--list", action="store_true", dest="list", help="show group list and exit")
parser.add_option("--min", action="store", dest="min", type=int, help="update group minimum capacity, must supply with max and target")
parser.add_option("--target", action="store", dest="target", type=int, help="update group target capacity, must supply with min and max")
parser.add_option("--max", action="store", dest="max", type=int, help="update group maximum capacity, must supply with min and target")
parser.add_option("--scale-up", action="store", dest="scale_up", type=int, help="scale up group by X number of instances")
parser.add_option("--scale-down", action="store", dest="scale_down", type=int, help="scale down group by X number of instances")
parser.add_option("--suspend", action="store", dest="suspend", choices=['AUTO_SCALE', 'AUTO_HEALING', 'OUT_OF_STRATEGY', 'PREVENTIVE_REPLACEMENT', 'REVERT_PREFERRED', 'SCHEDULING', 'BOTH'], help="suspend scaling or healing for a group. use the BOTH option to suspend scaling & healing (AUTO_SCALE/AUTO_HEALING/BOTH/REVERT_PREFERRED/OUT_OF_STRATEGY/PREVENTIVE_REPLACEMENT)")
parser.add_option("--unsuspend", action="store", dest="unsuspend", choices=['AUTO_SCALE', 'AUTO_HEALING', 'OUT_OF_STRATEGY', 'PREVENTIVE_REPLACEMENT', 'REVERT_PREFERRED', 'SCHEDULING', 'BOTH'], help="unsuspend scaling or healing for a group. (AUTO_SCALE/AUTO_HEALING/BOTH/REVERT_PREFERRED/OUT_OF_STRATEGY/PREVENTIVE_REPLACEMENT)")
parser.add_option("--ttl", action="store", dest="ttl", type=int, help="set an expiration for your suspension request - in minutes ( 10/20/60/300/600) ")
parser.add_option("--roll", action="store_true", dest="roll", help="roll a group, must supply batch-size, and grace-period")
parser.add_option("--batch-size", action="store", dest="batch", type=int, help="roll batch size - must supply with the roll flag")
parser.add_option("--grace-period", action="store", dest="grace", type=int, help="roll grace period - must supply with the roll flag")
parser.add_option("--replace-ami", action="store", dest="ami", help="replace AMI for group")
parser.add_option("--replace-health", action="store", dest="health", choices=['EC2', 'ELB', "HCS"], help="replace health check type for a group (EC2/ELB/HCS)")
parser.add_option("--user-data", action="store", dest="user_data", help="updated user data - supply a file path which contains the user data script (cloud init)")
parser.add_option("--detach-batch", action="store_true", dest="detach_batch", help="detach all instances for specific batch - choose from a list of batches")
parser.add_option("-y", "--skip-validation", action="store_true", dest="skip_validation", help="skip prompt validation for non-interactive mode")
parser.add_option("-e", "--exact", type="string", action="callback", dest="exact", callback=parser_split, help="grep exactly")
parser.add_option("--show-response", action="store", dest="show_response", type="string", help="Put True here is you want to see the request")


(options, args) = parser.parse_args()

# show help as default and exit
if len(argv) == 1:
    system("python " + argv[0] + " -h")
    exit()


# global variables
percent_alert_threshold = 30

# validate input
if options.max is not None and (options.min is None or options.target is None):
    errorAndExit("you must supply max with min and target!")
if options.min is not None and (options.max is None or options.target is None):
    errorAndExit("you must supply min with max and target!")
if options.target is not None and (options.max is None or options.min is None):
    errorAndExit("you must supply target with max and min!")
if options.roll is not None and (options.batch is None or options.grace is None):
    errorAndExit("you must supply batch size and grace period with the roll flag!")
if options.max is not None and options.min is not None and options.target is not None:
    if options.target < options.min or options.max < options.min or options.target > options.max:
        errorAndExit("check input - the number makes no sense!")

# set spotinst main path and check spotinst token
base_path = "https://api.spotinst.io/aws/ec2/group"
token = environ.get('spotinst_token')
if not token: errorAndExit("you must define an environment variable called 'spotinst_token' with a valid token (use export or .bashrc file)!")

# handle the kill batch option
if options.detach_batch:
    if options.skip_validation:
        errorAndExit("the skip validation option is not valid with the --detach-batch option!")
    # first get a specific group
    kill_in_group = get_specific_group(options.grep, options.ungrep)
    group_id = kill_in_group[1][0]
    group_name = kill_in_group[0]
    req_path = "{}/{}/status".format(base_path, group_id)
    res = query_api(req_path, token, 'GET', '')
    # get the group instances
    if res.status_code != 200:
        data = loads(res.content)
        print_in_color('red', "Error while getting the group status!\nstatus code: {}, reason: {}\n".format(res.status_code, data['response']['errors'][0]['message']))
    else:
        # find which instances has the create date supplied by the user
        batches = []
        data = loads(res.content)
        for instance in data['response']['items']:
            if not instance['createdAt'] in batches:
                batches.append(instance['createdAt'])
        batches = sorted(batches)
        print "\nFound the following batches for the group: {}\n".format(group_name)
        counter = 1
        for batch in batches:
            print "[{}] {}".format(counter, batch)
            counter += 1
        batch_number =  raw_input("\nWhich batch would you like to kill? ")
        try:
            batch_to_detach = batches[int(batch_number) - 1]
        except Exception as e:
            errorAndExit("no such batch exists!")

        instances_to_kill = []
        for instance in data['response']['items']:
            if batch_to_detach in instance['createdAt']:
                instances_to_kill.append(instance['instanceId'])

        # finally send a request to detach instances
        print "\nFound {} relevant instances for the group: {}\nInstances ID's:".format(len(instances_to_kill), group_name)
        for i in instances_to_kill:
            print "{} ".format(i),
        user_make_sure ("\n\nAre you sure you want to detach the above instances? [y/n] ")
        decrement = get_yes_no_from_user("\nDo you want to decrement the capacity of the group? [y/n] ")
        terminate = get_yes_no_from_user("\nDo you want to terminate the above instances (or just remove from group)? [y/n] ")
        req_path = "{}/{}/detachInstances".format(base_path, group_id)
        payload = dumps ({ "instancesToDetach": instances_to_kill, "shouldTerminateInstances" : terminate, "shouldDecrementTargetCapacity" : decrement })
        print_header(group_name)
        print "\nGoing to detach (terminate) the above instances..."
        res = query_api(req_path, token, 'PUT', payload)
        if res.status_code != 200:
            data = loads(res.content)
            print_in_color('red', "Error while getting the group status!\nstatus code: {}, reason: {}\n".format(res.status_code, data['response']['errors'][0]['message']))
        else:
            print_in_color('green', "\nSuccessfully terminated the required batch!")
    exit()

if options.exact is None:
    groups_to_update = get_groups(options.grep, options.ungrep)
    if len(groups_to_update) == 0:
        errorAndExit("could not find any groups for the search term: {}".format(options.grep))
    print_all_groups(groups_to_update)
else:
    groups_to_update = get_groups_exact(options.exact, options.ungrep)
    if len(groups_to_update) == 0:
        errorAndExit("could not find any groups for the search term: {}".format(options.exact))
    print_all_groups(groups_to_update)

# if the list flag was triggered, need to exit
if options.list is not None:
    exit()

# loop through each group and update
payload = None
change_validation = False
for group, group_values in groups_to_update.iteritems():
    group_id = group_values[0]

    # handle data request
    if options.data:
        req_type = 'GET'
        message = "retrieving data for group: {}".format(group)
        req_path = "{}/{}".format(base_path, group_id)

    # handle status request
    elif options.status:
        req_type = 'GET'
        message = "retrieving status for group: {}".format(group)
        req_path = "{}/{}/status".format(base_path, group_id)

    # handle suspension data request
    elif options.suspension:
        req_type = 'GET'
        message = "retrieving suspenstion status for group: {}".format(group)
        req_path = "{}/{}/suspension".format(base_path, group_id)

    # handle suspension request
    elif options.suspend is not None or options.unsuspend is not None:
        req_path = "{}/{}/suspension".format(base_path, group_id)

        if options.suspend is not None:
            req_type = 'POST'
            command = "suspend"
            if options.ttl is None:
                if options.suspend != "BOTH":
                    payload = dumps({"suspensions": [{"name": options.suspend}]})
                else:
                    payload = dumps({"suspensions": [{"name": "AUTO_HEALING"}, {"name": "AUTO_SCALE"}]})
            else:
                if options.suspend != "BOTH":
                    payload = dumps({"suspensions": [{"name": options.suspend, "ttlInMinutes": options.ttl}]})
                else:
                    payload = dumps({"suspensions": [{"name": "AUTO_HEALING", "ttlInMinutes": options.ttl}, {"name": "AUTO_SCALE", "ttlInMinutes": options.ttl}]})


        elif options.unsuspend is not None:
            if options.unsuspend != 'BOTH':
                req_type = 'DELETE'
                command = "unsuspend"
                payload = '{"processes":["%s"]}' % options.unsuspend
            else:
                req_type = 'DELETE'
                command = "unsuspend"
                payload = '{"processes":["AUTO_SCALE","AUTO_HEALING"]}'


        message = "sending request to {} the group".format(command)
    # handle scale request
    elif options.scale_down is not None or options.scale_up is not None:
        req_type = 'PUT'
        if options.scale_up is not None:
            command = "up"
            adjustment = options.scale_up
            target = int(group_values[2]) + int(adjustment)
        elif options.scale_down is not None:
            command = "down"
            adjustment = options.scale_down
            target = int(group_values[2]) - int(adjustment)
        # check that this change does not pass the percent_alert_threshold
        ans = check_change_in_target(group, group_values[2], target, percent_alert_threshold)
        if ans != "y":
            print "\nSkipping update for this group..."
            continue
        req_path = "{}/{}/scale/{}\?adjustment={}".format(base_path, group_id, command, adjustment)
        message = "sending request to scale {} {} instances...".format(command, adjustment)

    # handle group update request
    elif options.min is not None and options.target is not None and options.max is not None:
        # check that this change does not pass the percent_alert_threshold
        change_ans = check_change_in_target(group, group_values[2], options.target, percent_alert_threshold)
        if change_ans != "y":
            print "\nSkipping update for this group..."
            continue
        payload = dumps({"group": {"capacity": {"target": options.target, "minimum": options.min, "maximum": options.max}}})
        req_type = 'PUT'
        message = "sending the following payload: {}".format(payload)
        req_path = "{}/{}".format(base_path, group_id)

    # handle roll request
    elif options.roll is not None and options.batch is not None and options.grace is not None:
        req_type = 'PUT'
        req_path = "{}/{}/roll".format(base_path, group_id)
        payload = dumps({"batchSizePercentage": options.batch, "gracePeriod" : options.grace})
        message = "rolling group: {} with batch-size: {} and a grace period: {}".format(group, options.batch, options.grace)

    # handle replace ami
    elif options.ami is not None or options.health is not None:
        req_type = 'PUT'
        req_path = "{}/{}".format(base_path, group_id)
        if options.ami is not None:
            payload = dumps({ "group": { "compute": { "launchSpecification": { "imageId": options.ami  }}}})
            message = "replacing AMI for {}".format(group)
        else:
            payload = dumps({"group": {"compute": {"launchSpecification": {"healthCheckType": options.health}}}})
            message = "replacing health check type for {}".format(group)

    # handle replace of user data script
    elif options.user_data is not None:
        if not path.isfile(options.user_data):
            errorAndExit("{} is not a valid file name!".format(options.user_data))
        f = open(options.user_data, "r")
        user_data_script = f.read()
        f.close()
        user_data_script = b64encode(user_data_script)
        req_path = "{}/{}".format(base_path, group_id)
        req_type = 'PUT'
        payload = dumps({ "group": { "compute": { "launchSpecification": { "userData": user_data_script  }}}})
        message = "sending a request to update the user data script..."

    else:
        exit()

    # make sure the user wants to continue with the current filter (unless the req type is GET)
    if not change_validation and req_type != 'GET':
        if not options.skip_validation:
            user_make_sure("\nAre you sure you want to updated these groups[y/n]? ")
        else:
            print "Found skip validation option, skipping user prompt..."
        change_validation = True

    # send request
    print_header(group)
    print message
    res = query_api(req_path, token, req_type, payload)
    if res.status_code == 200:
        if req_type == 'GET' and options.show_response == 'True':
            print res.content
        elif options.suspension:
            json_response = loads(res.content)
            if not json_response['response']['items']:
                print "No suspensions in this Elastigroup"
            else:
                for suspension_type in json_response['response']['items']:
                    print suspension_type['processes']

        elif req_type == 'POST' and options.show_response == 'True':
            print res.content
        else:
            print_in_color('green', "Update was successful!")
    else:
        data = loads(res.content)
        print_in_color('red', "Error while updating group!\nstatus code: {}, reason: {}\n".format(res.status_code, data['response']['errors'][0]['message']))

print_message("Done running on all groups!")