#!/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 ./duane.py '*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 ./duane.py --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 ./duane.py --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 ./duane.py --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 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):
    """
    Share a specific AMI (by id) with a list of AWS account IDs
    within a specific region
    """

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

    for account in accounts:
        launchparam['Add'].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):
    """
    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)
                    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')))


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--region',action='append',required=False,
            help='Region to add sharing in (can specify multiple)')
    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()

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