Add tags option for select AMI and other changes

README.rst

@@ -1,15 +1,15 @@
 DistAMI
 =======
 
-.. image:: https://travis-ci.org/Answers4AWS/distami.png?branch=master
+.. image:: https://img.shields.io/travis/Answers4AWS/distami/master.svg
    :target: https://travis-ci.org/Answers4AWS/distami
    :alt: Build Status
 
-.. image:: https://pypip.in/d/distami/badge.png
+.. image:: https://img.shields.io/pypi/dm/distami.svg
    :target: https://pypi.python.org/pypi/distami
    :alt: PyPI Downloads
 
-.. image:: https://pypip.in/v/distami/badge.png
+.. image:: https://img.shields.io/pypi/v/distami.svg
    :target: https://pypi.python.org/pypi/distami
    :alt: PyPI Version
 
@@ -20,18 +20,21 @@ Usage
 
 ::
 
-    usage: distami [-h] [--region REGION] [--to REGIONS] [--non-public]
+    usage: distami [-h] [--ami_tags AMI_TAGs] [--region REGION] [--to REGIONS] [--non-public]
                    [--accounts AWS_ACCOUNT_IDs] [-p] [-v] [--version]
                    AMI_ID
 
     Distributes an AMI by copying it to one, many, or all AWS regions, and by
     optionally making the AMIs and Snapshots public.
 
     positional arguments:
-      AMI_ID                the source AMI ID to distribute. E.g. ami-1234abcd
+      AMI_ID                the source AMI ID to distribute. E.g. ami-1234abcd.
+                            When be set to "-", mean to use --ami-tags
 
     optional arguments:
       -h, --help            show this help message and exit
+      --ami-tags AMI_TAGs   the source AMI Tags to distribute. E.g.Name:Linux,Version:1.
+                            If use this argument, the AMI_ID must be set to "-"
       --region REGION       the region the AMI is in (default is current region of
                             EC2 instance this is running on). E.g. us-east-1
       --to REGIONS          comma-separated list of regions to copy the AMI to.
@@ -74,19 +77,25 @@ Copy AMI in ``us-east-1`` to ``us-west-1``
 
 ::
 
-    distami --region us-east-1 ami-abcd1234 --to us-west-1
+    distami --region us-east-1 --to us-west-1 ami-abcd1234
+
+Copy AMI with tags in ``us-east-1`` to ``us-west-1``
+
+::
+
+    distami --region us-east-1 --to us-west-1 --ami-tags Name:Linux,Version:3 -
 
 Copy an AMI in ``eu-west-1`` to ``us-west-1`` and ``us-west-2``, but do not make the AMI or its copies public
 
 ::
 
-    distami --region eu-west-1 ami-abcd1234 --to us-west-1,us-west-2 --non-public
+    distami --region eu-west-1 --to us-west-1,us-west-2 --non-public ami-abcd1234
 
 Share an AMI in ``us-east-1`` with the AWS account IDs 123412341234 and 987698769876. Do not copy to other regions and do not make public.
 
 ::
 
-    distami --region=us-east-1 ami-abcd1234 --to=none --accounts=123412341234,987698769876
+    distami --region=us-east-1 --to=none --accounts=123412341234,987698769876 ami-abcd1234
 
 
 Installation

distami/cli.py

@@ -41,7 +41,7 @@ def copy(param_array):
     to_region = param_array[1]
     args = param_array[2]
     copied_ami_id = distami.copy_to_region(to_region)
-    ami_cp = Distami(copied_ami_id, to_region)
+    ami_cp = Distami(copied_ami_id, to_region, args.ami_tags)
 
     if args.non_public:
         distami.make_ami_non_public()
@@ -58,12 +58,14 @@ def copy(param_array):
 def run():
     parser = argparse.ArgumentParser(description='Distributes an AMI by copying it to one, many, or all AWS regions, and by optionally making the AMIs and Snapshots public or shared with specific AWS Accounts.')
     parser.add_argument('ami_id', metavar='AMI_ID', 
-                        help='the source AMI ID to distribute. E.g. ami-1234abcd')
+                        help='the source AMI ID to distribute. E.g. ami-1234abcd. When be set to "-", mean to use --ami-tags')
+    parser.add_argument('--ami-tags', metavar='AMI_TAGs', default=None,
+                        help='the source AMI Tags to distribute. E.g. Name:Linux,Version:1. If use this argument, the ami_id must be set to "-"')
     parser.add_argument('--region', metavar='REGION', 
                         help='the region the AMI is in (default is current region of EC2 instance this is running on). E.g. us-east-1')
     parser.add_argument('--to', metavar='REGIONS', 
                         help='comma-separated list of regions to copy the AMI to. The default is all regions. Specify "none" to prevent copying to other regions. E.g. us-east-1,us-west-1,us-west-2')
-    parser.add_argument('--non-public', action='store_true', default=False, 
+    parser.add_argument('--non-public', action='store_true', default=True, 
                         help='Copies the AMIs to other regions, but does not make the AMIs or snapshots public. Bad karma, but good for AMIs that need to be private/internal only')
     parser.add_argument('--accounts', metavar='AWS_ACCOUNT_IDs', 
                         help='comma-separated list of AWS Account IDs to share an AMI with. Assumes --non-public. Specify --to=none to share without copying.')
@@ -98,7 +100,8 @@ def run():
         log.debug("Running in region: %s", ami_region)
 
     try:
-        distami = Distami(args.ami_id, ami_region)
+        args.ami_id = None if args.ami_id == "-" else args.ami_id
+        distami = Distami(args.ami_id, ami_region, args.ami_tags)
         if not args.non_public:
             distami.make_ami_public()
             distami.make_snapshot_public()

distami/core.py

@@ -25,8 +25,12 @@
 
 
 class Distami(object):
-    def __init__(self, ami_id, ami_region):
+    def __init__(self, ami_id, ami_region, ami_tags = None):
         self._ami_id = ami_id
+        if ami_tags is not None:
+            self._ami_tags = ami_tags.split(",")
+        else:
+            self._ami_tags = None
         self._ami_region = ami_region
 
         log.info("Looking for AMI %s in region %s", self._ami_id, self._ami_region)
@@ -36,7 +40,8 @@ def __init__(self, ami_id, ami_region):
             log.error('Could not connect to region %s' % self._ami_region)
             log.critical('No AWS credentials found. To configure Boto, please read: http://boto.readthedocs.org/en/latest/boto_config_tut.html')
             raise DistamiException('No AWS credentials found.')            
-        self._image = utils.wait_for_ami_to_be_available(self._conn, self._ami_id)
+        self._image = utils.wait_for_ami_to_be_available(self._conn, self._ami_id, self._ami_tags)
+        self._ami_id = self._image.id
         log.debug('AMI details: %s', vars(self._image))
 
         # Get current launch permissions
@@ -151,7 +156,6 @@ def copy_to_region(self, region):
         return copied_ami_id
 
 
-
 class Logging(object):
     # Logging formats
     _log_simple_format = '%(asctime)s [%(levelname)s] %(message)s'

distami/utils.py

@@ -22,17 +22,26 @@
 log = logging.getLogger(__name__)
 
 
-def get_ami(conn, ami_id):
+def get_ami(conn, ami_id, ami_tags = None):
     ''' Gets a single AMI as a boto.ec2.image.Image object '''
 
     attempts = 0
     max_attempts = 5
+    image_filters = {}
+    if ami_tags is not None:
+        for tag in ami_tags:
+            kv = tag.split(":")
+            if len(kv) == 2:
+                image_filters["tag:" + kv[0]] = kv[1]
+    else:
+        image_filters = None
+
     while (attempts < max_attempts):
         try:
             attempts += 1
-            images = conn.get_all_images(ami_id)
+            images = conn.get_all_images(ami_id, filters=image_filters)
         except boto.exception.EC2ResponseError:
-            msg = "Could not find AMI '%s' in region '%s'" % (ami_id, conn.region.name)
+            msg = "Could not find AMI '%s' with tags %s in region '%s'" % (ami_id, ami_tags, conn.region.name)
             if attempts < max_attempts:
                 # The API call to initiate an AMI copy is not blocking, so the
                 # copied AMI may not be available right away
@@ -84,19 +93,19 @@ def get_regions_to_copy_to(source_region):
     return regions
 
 
-def wait_for_ami_to_be_available(conn, ami_id):
+def wait_for_ami_to_be_available(conn, ami_id, ami_tags = None):
     ''' Blocking wait until the AMI is available '''
 
-    ami = get_ami(conn, ami_id)
+    ami = get_ami(conn, ami_id, ami_tags)
     log.debug('AMI details: %s', vars(ami))
 
     while ami.state != 'available':
         log.info("%s in %s not available, waiting...", ami_id, conn.region.name)
         time.sleep(30)
-        ami = get_ami(conn, ami_id)
+        ami = get_ami(conn, ami_id, ami_tags)
 
         if ami.state == 'failed':
             msg = "AMI '%s' is in a failed state and will never be available" % ami_id
             raise DistamiException(msg)
 
-    return ami
+    return ami