#!/usr/bin/env python3
#
# Startup/Shutdown Groups of our systems
#
# Common usage:
# 
# View the current state of the things on a schedule:
#   xdrtest status
# 
# At the start of your shift:
#   xdrtest start --profile mdr-test-c2-gov
#
# At the end of your shift:
#   xdrtest stop
#
# You can also start/stop individual instances by ID (provided the tag _also_ matches):
#   xdrtest start openvpn
import argparse
import boto3
import fnmatch
import json
import logging
import os
import re
import sys

DEBUG=False # Turn on with --debug

logger = logging.getLogger()

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
            It must be "yes" (the default), "no" or None (meaning
            an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".
    """
    valid = {"yes": True, "y": True, "ye": True, "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        sys.stdout.write(question + prompt)
        choice = input().lower()
        if default is not None and choice == "":
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' " "(or 'y' or 'n').\n")


def gather_profiles(file='~/.aws/config', commercial=False):
    file=os.path.expanduser(file)
    profiles=list()
    try:
        with open(os.path.expanduser(file), "r") as configfile:
            for line in configfile:
                match = re.search('^\[profile (.*)\]', line)
                if match:
                    if 'test' in match.group(1):
                        if 'gov' in match.group(1) or commercial:
                            profiles.append(match.group(1))
    except Exception as e:
        logger.error(f'Could not open file {file} for reading: {e}')
        sys.exit(1)
    return profiles


def gather_instances(ec2, tagname, tagvalue, exclude_state):
    ''' 
    Gathers the instance information from the provided ec2 resource
    '''
    desired_instances = ec2.instances.filter(
            Filters=[
                { 
                  'Name': f'tag:{tagname}',
                  'Values': [ f'{tagvalue}' ]
                }
              ]
    )
    for i in desired_instances:
        name=""
        for t in i.tags:
            if t['Key']=='Name':
                name = t['Value']
        logger.debug(f'DEBUG: Discovered instance {i.instance_id}, Name: {name}, State: {i.state["Name"]}')
        if i.state['Name'] != exclude_state:
            # Why an iterator? Why not?
            yield {
                    'instance_id': i.instance_id,
                    'name': name,
                    'state': i.state['Name']
                  }


if __name__ == "__main__":
    handler = logging.StreamHandler()
    logger.addHandler(handler)

    # Turn off debugging for AWS
    for module in [ 'boto3', 'botocore', 'nose', 's3transfer', 'urllib3', 'urllib3.connectionpool' ]:
            l = logging.getLogger(module)
            l.setLevel(logging.WARNING) # bump this up if you think you have a boto3 bug

    parser = argparse.ArgumentParser(description='Startup/Shutdown Groups of Instances')
    parser.add_argument('action', help='Action to Take on Instances', choices=['stop', 'start', 'status'])
    parser.add_argument('--config', help='AWS Config File', default='~/.aws/config')
    parser.add_argument('--profile', help='Profile to Use', default=None)
    parser.add_argument('--tagname', help='Tag Name to Filter On', default='Schedule')
    parser.add_argument('--tagvalue', help='Tag Value to Filter On', default='MSOC')
    parser.add_argument('instancename', help='One or more instance names or IDs to action on. Wildcards are allowed.', nargs='*', default=None)
    parser.add_argument('--commercial', help='Include Commercial Accounts', action='store_true')
    parser.add_argument('--dry-run', help='Dry Run', action='store_true')
    parser.add_argument('--debug', help='Output Debug Information', action='store_true')

    args = parser.parse_args()
    if args.debug:
        DEBUG=True
        logger.setLevel(logging.DEBUG) # Debug
    else:
        logger.setLevel(logging.INFO) # Info

    if args.profile is not None:
        profiles=[ args.profile ]
    else:
        profiles=gather_profiles(file=args.config, commercial=args.commercial)

    logger.debug(f'Profiles to gather: {json.dumps(profiles)}')

    # Connect to our profiles:
    session = dict()
    ec2 = dict()
    instances = dict()
    instance_count = 0
    print(f'Instances to { args.action }:')
    for p in profiles:
        logger.debug(f'Gathering from profile {p}...')
        try:
            session[p] = boto3.session.Session(profile_name=p)
            ec2[p] = session[p].resource('ec2')
        except Exception as e:
            logger.error(f'Could not connect to profile "{p}": {e}')

        exclude_state = ""
        if args.action=="start":
            exclude_state="running"
        elif args.action=="stop":
            exclude_state="stopped"
       
        instances[p] = list()
        try:
            for instance in gather_instances(ec2[p], args.tagname, args.tagvalue, exclude_state):
                if args.instancename:
                    # At least one instance name was specified, so go through and match
                    for i in args.instancename:
                        if fnmatch.fnmatch(instance['name'], i):
                            print(f'{p:<30}\t{ instance["instance_id"] }\t{instance["state"]:<10}\t{instance["name"]}')
                            instances[p].append(instance['instance_id'])
                            instance_count += 1
                else:
                    # No instance name specified, we add everything
                    print(f'{p:<30}\t{ instance["instance_id"] }\t{instance["state"]:<10}\t{instance["name"]}')
                    instances[p].append(instance['instance_id'])
                    instance_count += 1
        except Exception as e:
            logger.error(f'Could not enumerate instances in profile "{p}": {e}')
            profiles.remove(p)

    if instance_count == 0:
        print('')
        print('No applicable instances discovered.')
        sys.exit(0)

    logger.debug(f'Num instances: {len(instances)}: {json.dumps(instances)}')
    if args.action=="status":
        logger.debug('Exiting after status check')
        sys.exit(0)

    print('')
    if query_yes_no(question='Continue?', default='no') is False:
        logger.debug('Exiting on user request')
        sys.exit(2)

    for p in profiles:
        if len(instances[p]) > 0:
            for i in ec2[p].instances.filter(InstanceIds=instances[p]):
                if args.action == 'start':
                    print(f'Starting instances in profile { p }: { i.instance_id }')
                    try:
                        i.start(DryRun=args.dry_run)
                    except Exception as e:
                        print(f'An error occured while starting instance {i.id}. Error: {e}. Skipping.')
                elif args.action == 'stop':
                    print(f'Stopping instances in profile { p }: { i.instance_id }')
                    try:
                        i.stop(DryRun=args.dry_run)
                    except Exception as e:
                        print(f'An error occured while stopping instance {i.id}. Error: {e}. Skipping.')
                else:
                    print(f'Unsupported action: {args.action}')
                    sys.exit(3)