#!/usr/bin/env python3
"""
CLI tool to help with AMI sharing.  What I was doing before in bash
hit its reasonable limit of complexity, and was getting hard to read
because it was in bash.

Some notes:

    [1]  Specifying an AMI filter is mandatory, but regions and accounts are not.
    [2]  Standard AWS_PROFILE environment variables are supported (because boto3 supports)
    [3]  AWS_REGION environment variable (as used by the AWS CLI) is not supported because
         boto3 does not use it.  It uses AWS_DEFAULT_REGION instead.  If your
         $HOME/.aws/config lists a region for a given profile this is not needed.


Example 1: Let's just run a report of all AMIs matching '*Duane*' in all regions that the
profile has access to.  Notice the wildcards in quotes so bash won't try to expand them
out to filenames.

[duane.e.waddle@DPS0591 bin]$ AWS_PROFILE=gov-common-services-terraformer ./update-ami-accounts '*Duane*'
Looking for AMIs matching "*Duane*" in the following regions:
    us-gov-east-1
    us-gov-west-1

AMIs matching the filter:
region         |ami id                |ami name
===============|======================|========================================
us-gov-east-1  |ami-069f3e239427365b6 |Duane_Testing_20201124233617
us-gov-west-1  |ami-0ee37a86b09aefad0 |Duane_Testing_20201124233617


Example 2: Regions can be specified with a list or wildcard.  This is just a report too:

[duane.e.waddle@DPS0591 bin]$ AWS_PROFILE=gov-common-services-terraformer ./update-ami-accounts --region us-gov-east-1 --region '*west*' '*Duane*'
Looking for AMIs matching "*Duane*" in the following regions:
    us-gov-east-1
    us-gov-west-1

AMIs matching the filter:
region         |ami id                |ami name
===============|======================|========================================
us-gov-east-1  |ami-069f3e239427365b6 |Duane_Testing_20201124233617
us-gov-west-1  |ami-0ee37a86b09aefad0 |Duane_Testing_20201124233617

Example 3: If we list one or more accounts then sharing is updated

[duane.e.waddle@DPS0591 bin]$ AWS_PROFILE=gov-common-services-terraformer ./update-ami-accounts --region '*1' '*Duane*' 738800754746 721817724804
Looking for AMIs matching "*Duane*" in the following regions:
 us-gov-east-1
 us-gov-west-1

Sharing AMIs with these accounts:
 738800754746
 721817724804

AMIs matching the filter:
region         |ami id                |ami name                                |status
===============|======================|========================================|==========
us-gov-east-1  |ami-069f3e239427365b6 |Duane_Testing_20201124233617            |success
us-gov-west-1  |ami-0ee37a86b09aefad0 |Duane_Testing_20201124233617            |success

Example 4: Sharing updates are atomic so if you could get a failure because
one of several accounts you listed does not exist:

[duane.e.waddle@DPS0591 bin]$ AWS_PROFILE=gov-common-services-terraformer ./update-ami-accounts --region '*1' '*Duane*' 738800754746 72181772480
Looking for AMIs matching "*Duane*" in the following regions:
 us-gov-east-1
 us-gov-west-1

Sharing AMIs with these accounts:
 738800754746
 72181772480

AMIs matching the filter:
region         |ami id                |ami name                                |status
===============|======================|========================================|==========
us-gov-east-1  |ami-069f3e239427365b6 |Duane_Testing_20201124233617            |error
us-gov-west-1  |ami-0ee37a86b09aefad0 |Duane_Testing_20201124233617            |error

Notice one of the account numbers is missing a digit.  You don't get a clear
message as to which one caused the error.  Maybe one day I'll improve that...

"""

import argparse
import re
import sys
import boto3
import botocore
from botocore.config import Config


def list_matching_regions(region_filter):
    """
    Return the list of regions matching a wildcard.
    """
    ec2 = boto3.client('ec2')
    regions = ec2.describe_regions(
                Filters=[
                    {
                        'Name': 'region-name',
                        'Values': [ region_filter ]
                    }
                ],
                AllRegions = False
            )
    return [ x.get('RegionName') for x in regions.get('Regions',[]) ]

def find_amis_matching_filter(ami_filter,region=None):
    """
    Return a list of AMIs matching a given wildcard by name
    """

    if region is not None:
        ec2 = boto3.client('ec2',config=Config(region_name=region))
    else:
        ec2 = boto3.client('ec2')
    images = ec2.describe_images(
                Owners=['self'],
                Filters=[
                    {
                        'Name': 'name',
                        'Values': [ ami_filter ]
                    }
                ]
            )

    fields_of_interest = { 'ImageId', 'Name' }
    for entry in images.get('Images',[]):
        newentry = { k: entry[k] for k in entry.keys() & fields_of_interest }
        yield newentry

def share_ami(ami,region,accounts,remove=False):
    """
    Share a specific AMI (by id) with a list of AWS account IDs
    within a specific region
    """

    launchparam = { }
    launchparam['Add'] = []

    if not remove:
        for account in accounts:
            launchparam['Add'].extend([{'UserId': account}])
    else:
        for account in accounts:
            launchparam['Remove'].extend([{'UserId': account}])

    ec2 = boto3.resource('ec2',config=Config(region_name=region))

    # expecting to possibly raise something here
    image = ec2.Image(ami)
    image.modify_attribute(Attribute='launchPermission', LaunchPermission=launchparam)

def runmain(ami_filter,accounts,region_filters,remove=False):
    """
    main
    """

    region_list = []
    if region_filters is None:
        region_list.extend(list_matching_regions('*'))
    else:
        for filt in region_filters:
            region_list.extend(list_matching_regions(filt))

    region_list=list(sorted(set(region_list)))

    print('Looking for AMIs matching "{0}" in the following regions:'.format(ami_filter))
    for region in region_list:
        print(" {0}".format(region))
    print("")

    # Shorter format for when we're just doing a report of matching AMIs
    # No accounts means just report mode
    if len(accounts) > 0:

        print("Sharing AMIs with these accounts:")
        for account in accounts:
            print(" {0}".format(account))
        print("")

        print("AMIs matching the filter:")
        report_format="{0:<15}|{1:<22}|{2:<40}|{3:<10}"
        print(report_format.format('region','ami id','ami name','status'))
        print(report_format.format('='*15,'='*22,'='*40,'='*10))
    else:
        print("AMIs matching the filter:")
        report_format="{0:<15}|{1:<22}|{2:<40}"
        print(report_format.format('region','ami id','ami name'))
        print(report_format.format('='*15,'='*22,'='*40))

    for region in region_list:
        for ami in find_amis_matching_filter(ami_filter,region):
            if len(accounts) > 0:
                try:
                    share_ami(ami.get('ImageId'),region,accounts,remove)
                    print(report_format.format(region,ami.get('ImageId'),ami.get('Name'),"success"))
                except botocore.exceptions.ClientError:
                    print(report_format.format(region,ami.get('ImageId'),ami.get('Name'),"error"))
            # No accounts we're just making a report ...
            else:
                print(report_format.format(region,ami.get('ImageId'),ami.get('Name')))

def cli():

    parser = argparse.ArgumentParser()
    parser.add_argument('--region',action='append',required=False,
            help='Region to add sharing in (can specify multiple)')
    parser.add_argument('--remove',action='store_true',required=False,
            help='Set this flag to remove instead of add sharing.')
    parser.add_argument('ami_filter',help='AMI Filter to apply')
    parser.add_argument('accounts',nargs='*',help='list of AWS accounts to add AMIs to')
    args = parser.parse_args()

    # Remove dashes from account IDs as a nice thing for user
    args.accounts = [ f.replace("-","") for f in args.accounts ]

    # AWS Accounts are 12 digits, all digits
    invalid_accounts = []
    digit_check = re.compile(r"^\d+$")
    for account in args.accounts:
        if len(account) != 12:
            invalid_accounts.append("{0}: {1}".format(account,"Account length not 12"))
        elif digit_check.fullmatch(account) is None:
            invalid_accounts.append("{0}: {1}".format(account,"Account contains non-digit chars"))

    if len(invalid_accounts) > 0:
        for message in invalid_accounts:
            print(message)
        sys.exit(1)

    runmain(args.ami_filter,args.accounts,args.region,args.remove)

if __name__ == "__main__":
    cli()