#!/usr/bin/env python

"""
Get/set AWS SSM parameters.
"""

import os
import sys
import boto3
from docopt import docopt
from termcolor import colored

COMMAND_VERSION = "1.0.6"
COMMAND_NAME = os.path.basename(sys.argv[0]).lower()
USAGE_TEXT = """command_basename command_version

Usage:
    command_basename get [-s] [-p AWS_PROFILE] [-r AWS_REGION]
    command_basename set FILE [-d] [-p AWS_PROFILE] [-r AWS_REGION]
    command_basename (-h | --help)
    command_basename (-v | --version)

Arguments:
    get           Get parameters.
    set           Create/update/delete parameters.
    FILE          Path to file that contains parameters.

Options:
    -d            Perform a dryrun.
    -h --help     Show this screen.
    -p PROFILE    AWS profile to use.
    -r REGION     AWS region to use.
    -s            Get secret value.
    -v --version  Show version.

FILE format example:
    #comment
    /create/string=foo
    */create/securestring=password
    !/create/string/but/do/not/update=bar
    !*/create/securestring/but/do/not/update=password
    -/delete/anything=true
"""

USAGE_TEXT = USAGE_TEXT.replace("command_basename", COMMAND_NAME)
USAGE_TEXT = USAGE_TEXT.replace("command_version", COMMAND_VERSION)

COLOR_INFO = "white"
COLOR_ERROR = "magenta"
COLOR_DRYRUN = "cyan"
COLOR_UPDATED = "yellow"
COLOR_CREATED = "green"
COLOR_DELETED = "red"


def do_get():
    """
    Get SSM parameters.
    """

    next_token = ""

    while next_token is not None:
        ssmargs = {
            "Path": "/",
            "Recursive": True,
            "WithDecryption": cmdargs["-s"],
            "MaxResults": 10,
        }

        if next_token != "":
            ssmargs["NextToken"] = next_token

        try:
            response = ssmclient.get_parameters_by_path(**ssmargs)
        except Exception as _e:
            print(colored(_e, COLOR_ERROR))
            sys.exit(1)

        next_token = response["NextToken"] if "NextToken" in response else None
        print_params(response["Parameters"])


def do_set():
    """
    Create/update/delete SSM parameters
    """

    parameters = read_params_from_file(cmdargs["FILE"])

    for key, value in parameters.items():
        if key.startswith("-"):
            if key[1:].startswith("!*"):
                delete_parameter(key[3:])
            elif key[1:].startswith("!") or key[1:].startswith("*"):
                delete_parameter(key[2:])
            else:
                delete_parameter(key[1:])
        elif key.startswith("!"):
            if key[1:].startswith("*"):
                create_parameter(key[2:], value, "SecureString")
            else:
                create_parameter(key[1:], value)
        elif key.startswith("*"):
            upsert_parameter(key[1:], value, "SecureString")
        else:
            upsert_parameter(key, value)


def delete_parameter(key):
    """
    Delete SSM parameter if it exists.
    """

    try:
        ssmclient.get_parameter(Name=key)
    except ssmclient.exceptions.ParameterNotFound as _e:
        print(colored("[-] Skipped (does not exist): %s" % key, COLOR_INFO))
    except Exception as _e:
        print(colored(_e, COLOR_ERROR))
        sys.exit(1)
    else:
        if not cmdargs["-d"]:
            try:
                ssmclient.delete_parameter(Name=key)
            except Exception as _e:
                print(colored(_e, COLOR_ERROR))
                sys.exit(1)
        print(colored("[-] Deleted: %s" % key, COLOR_DELETED))


def create_parameter(key, value, paramtype="String"):
    """
    Create, but do not update, SSM parameter.
    """

    decrypt = bool(paramtype == "SecureString")

    try:
        ssmclient.get_parameter(Name=key, WithDecryption=decrypt)
    except ssmclient.exceptions.ParameterNotFound as _e:
        if not cmdargs["-d"]:
            try:
                ssmclient.put_parameter(Name=key, Value=value, Type=paramtype, Overwrite=False)
            except Exception as _e:
                print(colored(_e, COLOR_ERROR))
                sys.exit(1)
        if paramtype == "SecureString":
            key = "*" + key
        print(colored("[!] Created: %s" % key, COLOR_CREATED))
    except Exception as _e:
        print(colored(_e, COLOR_ERROR))
        sys.exit(1)
    else:
        if paramtype == "SecureString":
            key = "*" + key
        print(colored("[!] Skipped (already exists): %s" % key, COLOR_INFO))


def upsert_parameter(key, value, paramtype="String"):
    """
    Create or update (if value is different) SSM parameter.
    """

    decrypt = bool(paramtype == "SecureString")

    try:
        response = ssmclient.get_parameter(Name=key, WithDecryption=decrypt)
    except ssmclient.exceptions.ParameterNotFound as _e:
        if not cmdargs["-d"]:
            try:
                ssmclient.put_parameter(Name=key, Value=value, Type=paramtype, Overwrite=False)
            except Exception as _e:
                print(colored(_e, COLOR_ERROR))
                sys.exit(1)
        if paramtype == "SecureString":
            key = "*" + key
        print(colored("[+] Created: %s" % key, COLOR_CREATED))
    except Exception as _e:
        print(colored(_e, COLOR_ERROR))
        sys.exit(1)
    else:
        if response.get("Parameter").get("Value") == value:
            if paramtype == "SecureString":
                key = "*" + key
            print(colored("[+] Skipped (identical): %s" % key, COLOR_INFO))
        else:
            if not cmdargs["-d"]:
                try:
                    ssmclient.put_parameter(Name=key, Value=value, Type=paramtype, Overwrite=True)
                except Exception as _e:
                    print(colored(_e, COLOR_ERROR))
                    sys.exit(1)
            if paramtype == "SecureString":
                key = "*" + key
            print(colored("[+] Updated: %s" % key, COLOR_UPDATED))


def print_params(params):
    """
    Print SSM parameters.

    Parameters:
        - params: the key value parameters to print.
    """

    for param in params:
        key = "*" + param["Name"] if param["Type"] == "SecureString" else param["Name"]
        if cmdargs["-s"] or param["Type"] == "String":
            print("%s=%s" % (key, param["Value"]))
        else:
            print("%s=%s" % (key, "<SECRET>"))


def read_params_from_file(file):
    """
    Parse SSM parameters from file.

    Parameters:
        - file: the file to parse for ssm parameters.

    Returns: a dictionary of the key/values.
    """

    result = {}

    try:
        with open(file) as _f:
            content = _f.readlines()
    except EnvironmentError as _e:
        print(colored(_e, COLOR_ERROR))
        sys.exit(1)

    tups = [
        line.strip().split("=", 1)
        for line in content
        if not (line.startswith("#") or len(line.strip()) == 0)
    ]

    for key, value in tups:
        result[key] = value

    return result


if __name__ == "__main__":
    cmdargs = docopt(USAGE_TEXT, version=COMMAND_VERSION)

    try:
        if cmdargs["-p"] or cmdargs["-r"]:
            profile = cmdargs["-p"] if cmdargs["-p"] else None
            region = cmdargs["-r"] if cmdargs["-r"] else None
            session = boto3.session.Session(profile_name=profile, region_name=region)
            ssmclient = session.client("ssm")
        else:
            ssmclient = boto3.client("ssm")
    except Exception as _e:
        print(colored(_e, COLOR_ERROR))
        sys.exit(1)

    if cmdargs["-d"]:
        print(colored("Dryrun enabled. No changes will be made.", COLOR_DRYRUN))

    if cmdargs["set"]:
        do_set()

    if cmdargs["get"]:
        do_get()

    if cmdargs["-d"]:
        print(colored("Dryrun enabled. No changes were made.", COLOR_DRYRUN))

    sys.exit(0)
