Merge pull request #279 from mdr-engineering/feature/ftd_na_AddRemoveOptionToAMIScript

Adds optional `--remove` flag to update_ami_accounts

bin/update-ami-accounts

@@ -130,7 +130,7 @@ def find_amis_matching_filter(ami_filter,region=None):
         newentry = { k: entry[k] for k in entry.keys() & fields_of_interest }
         yield newentry
 
-def share_ami(ami,region,accounts):
+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
@@ -139,8 +139,12 @@ def share_ami(ami,region,accounts):
     launchparam = { }
     launchparam['Add'] = []
 
-    for account in accounts:
-        launchparam['Add'].extend([{'UserId': account}])
+    if 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))
 
@@ -148,7 +152,7 @@ def share_ami(ami,region,accounts):
     image = ec2.Image(ami)
     image.modify_attribute(Attribute='launchPermission', LaunchPermission=launchparam)
 
-def runmain(ami_filter,accounts,region_filters):
+def runmain(ami_filter,accounts,region_filters,remove=False):
     """
     main
     """
@@ -190,7 +194,7 @@ def runmain(ami_filter,accounts,region_filters):
         for ami in find_amis_matching_filter(ami_filter,region):
             if len(accounts) > 0:
                 try:
-                    share_ami(ami.get('ImageId'),region,accounts)
+                    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"))
@@ -203,6 +207,8 @@ 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()
@@ -224,7 +230,7 @@ def cli():
             print(message)
         sys.exit(1)
 
-    runmain(args.ami_filter,args.accounts,args.region)
+    runmain(args.ami_filter,args.accounts,args.region,args.remove)
 
 if __name__ == "__main__":
     cli()