added option --owner-id

.travis.yml

@@ -1,6 +1,9 @@
 env:
   global:
     - AWS_DEFAULT_REGION=eu-west-1
+    - AWS_ACCESS_KEY_ID=local
+    - AWS_SECRET_ACCESS_KEY=secret
+    - BOTO_CONFIG=/dev/null
 
 dist: trusty
 language: python

README.rst

@@ -135,6 +135,12 @@ Exclude amis based on tag values
 
     amicleaner --mapping-key tags --mapping-values role env -excluded-mapping-values prod
 
+Specify the aws account owner of the AMIs
+
+ .. code:: bash
+
+     amicleaner --owner-id 1234567
+
 Skip confirmation, can be useful for automation
 
 .. code:: bash

amicleaner/cli.py

@@ -10,7 +10,7 @@
 from amicleaner import __version__
 from .core import AMICleaner, OrphanSnapshotCleaner
 from .fetch import Fetcher
-from .resources.config import MAPPING_KEY, MAPPING_VALUES, EXCLUDED_MAPPING_VALUES
+from .resources.config import MAPPING_KEY, MAPPING_VALUES, EXCLUDED_MAPPING_VALUES, OWNER_ID
 from .resources.config import TERM
 from .utils import Printer, parse_args
 
@@ -29,6 +29,7 @@ def __init__(self, args):
         self.full_report = args.full_report
         self.force_delete = args.force_delete
         self.ami_min_days = args.ami_min_days
+        self.owner_id = args.owner_id or OWNER_ID
 
         self.mapping_strategy = {
             "key": self.mapping_key,
@@ -43,7 +44,7 @@ def fetch_candidates(self, available_amis=None, excluded_amis=None):
         AMIs from ec2 instances, launch configurations, autoscaling groups
         and returns unused AMIs.
         """
-        f = Fetcher()
+        f = Fetcher(owner_id=self.owner_id)
 
         available_amis = available_amis or f.fetch_available_amis()
         excluded_amis = excluded_amis or []
@@ -68,7 +69,7 @@ def prepare_candidates(self, candidates_amis=None):
         if not candidates_amis:
             return None
 
-        c = AMICleaner()
+        c = AMICleaner(owner_id=self.owner_id)
 
         mapped_amis = c.map_candidates(
             candidates_amis=candidates_amis,
@@ -119,7 +120,7 @@ def clean_orphans(self):
 
         """ Find and removes orphan snapshots """
 
-        cleaner = OrphanSnapshotCleaner()
+        cleaner = OrphanSnapshotCleaner(owner_id=self.owner_id)
         snaps = cleaner.fetch()
 
         if not snaps:
@@ -140,6 +141,7 @@ def clean_orphans(self):
     def print_defaults(self):
 
         print(TERM.bold("\nDefault values : ==>"))
+        print(TERM.green("owner_id : {0}".format(self.owner_id)))
         print(TERM.green("mapping_key : {0}".format(self.mapping_key)))
         print(TERM.green("mapping_values : {0}".format(self.mapping_values)))
         print(TERM.green("excluded_mapping_values : {0}".format(self.excluded_mapping_values)))

amicleaner/core.py

@@ -8,7 +8,7 @@
 from botocore.exceptions import ClientError
 from botocore.config import Config
 
-from .resources.config import BOTO3_RETRIES
+from .resources.config import BOTO3_RETRIES, OWNER_ID
 from .resources.models import AMI
 
 from datetime import datetime
@@ -18,8 +18,9 @@ class OrphanSnapshotCleaner(object):
 
     """ Finds and removes ebs snapshots left orphaned """
 
-    def __init__(self, ec2=None):
+    def __init__(self, ec2=None, owner_id=None):
         self.ec2 = ec2 or boto3.client('ec2', config=Config(retries={'max_attempts': BOTO3_RETRIES}))
+        self.owner_id = owner_id or OWNER_ID
 
     def get_snapshots_filter(self):
 
@@ -49,7 +50,7 @@ def fetch(self):
 
         """ retrieve orphan snapshots """
 
-        resp = self.ec2.describe_images(Owners=['self'])
+        resp = self.ec2.describe_images(Owners=[self.owner_id])
 
         used_snaps = [
             ebs.get("Ebs", {}).get("SnapshotId")
@@ -95,8 +96,9 @@ def log(self, msg):
 
 class AMICleaner(object):
 
-    def __init__(self, ec2=None):
+    def __init__(self, ec2=None, owner_id=None):
         self.ec2 = ec2 or boto3.client('ec2', config=Config(retries={'max_attempts': BOTO3_RETRIES}))
+        self.owner_id = owner_id or OWNER_ID
 
     @staticmethod
     def get_ami_sorting_key(ami):
@@ -141,7 +143,7 @@ def remove_amis_from_ids(self, ami_ids):
             return False
 
         my_custom_images = self.ec2.describe_images(
-            Owners=['self'],
+            Owners=[self.owner_id],
             ImageIds=ami_ids
         )
         amis = []

amicleaner/fetch.py

@@ -5,28 +5,29 @@
 from builtins import object
 import boto3
 from botocore.config import Config
-from .resources.config import BOTO3_RETRIES
+from .resources.config import BOTO3_RETRIES, OWNER_ID
 from .resources.models import AMI
 
 
 class Fetcher(object):
 
     """ Fetches function for AMI candidates to deletion """
 
-    def __init__(self, ec2=None, autoscaling=None):
+    def __init__(self, ec2=None, autoscaling=None, owner_id=None):
 
         """ Initializes aws sdk clients """
 
         self.ec2 = ec2 or boto3.client('ec2', config=Config(retries={'max_attempts': BOTO3_RETRIES}))
         self.asg = autoscaling or boto3.client('autoscaling')
+        self.owner_id = owner_id or OWNER_ID
 
     def fetch_available_amis(self):
 
         """ Retrieve from your aws account your custom AMIs"""
 
         available_amis = dict()
 
-        my_custom_images = self.ec2.describe_images(Owners=['self'])
+        my_custom_images = self.ec2.describe_images(Owners=[self.owner_id])
         for image_json in my_custom_images.get('Images'):
             ami = AMI.object_with_json(image_json)
             available_amis[ami.id] = ami

amicleaner/resources/config.py

@@ -29,6 +29,9 @@
 EXCLUDED_MAPPING_VALUES = []
 
 
+OWNER_ID = "self"
+
+
 # Number of days amis to keep based on creation date and grouping strategy
 # not including the ami currently running by an ec2 instance
 AMI_MIN_DAYS = -1

amicleaner/utils.py

@@ -8,7 +8,7 @@
 
 from prettytable import PrettyTable
 
-from .resources.config import KEEP_PREVIOUS, AMI_MIN_DAYS
+from .resources.config import KEEP_PREVIOUS, AMI_MIN_DAYS, OWNER_ID
 
 
 class Printer(object):
@@ -119,6 +119,12 @@ def parse_args(args):
                         help="Number of days AMI to keep excluding those "
                              "currently being running")
 
+    parser.add_argument("--owner-id",
+                        dest='owner_id',
+                        default=OWNER_ID,
+                        help="Owner of the images to clean: "
+                             "[self|amazon|aws-marketplace|123456]")
+
     parsed_args = parser.parse_args(args)
     if parsed_args.mapping_key and not parsed_args.mapping_values:
         print("missing mapping-values\n")

requirements_build.txt

@@ -3,6 +3,6 @@ blessings==1.6
 codecov
 moto
 pep8>=1.7.0
-pytest
+pytest>=4.1.1
 pytest-pep8
 pytest-cov

tests/test_cli.py

@@ -166,6 +166,7 @@ def test_parse_args_no_args():
     assert parser.mapping_values is None
     assert parser.keep_previous is 4
     assert parser.ami_min_days is -1
+    assert parser.owner_id is "self"
 
 
 def test_parse_args():
@@ -185,6 +186,9 @@ def test_parse_args():
     assert parser.ami_min_days == 10
     assert parser.full_report is True
 
+    parser = parse_args(['--owner-id', '123456'])
+    assert parser.owner_id == "123456"
+
 
 def test_print_report():
     assert Printer.print_report({}) is None