Merge pull request #78 from jantman/pr64

Fixup of #64 - region and STS support

.travis.yml

@@ -6,14 +6,12 @@ cache:
     - $HOME/.pip-cache/
 
 env:
-  - TOXENV=py26-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py27-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py32-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py33-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py34-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=pypy-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=pypy3-unit PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
-  - TOXENV=py26-versioncheck PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py27-versioncheck PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py32-versioncheck PIP_DOWNLOAD_CACHE=$HOME/.pip-cache
   - TOXENV=py33-versioncheck PIP_DOWNLOAD_CACHE=$HOME/.pip-cache

CHANGES.rst

@@ -6,6 +6,12 @@ Pre-release (develop branch)
 
 * Update trove classifier Development Status in setup.py to Beta
 * Fix markup formatting issue in ``docs/source/getting_started.rst``
+* temporarily disable py26 testenv in Travis; failing due to upstream bug https://github.com/pytest-dev/pytest/issues/1035
+* `PR #64 <https://github.com/jantman/awslimitchecker/pull/64>`_ and `#68 <https://github.com/jantman/awslimitchecker/issues/68>`_ -
+  support [STS](http://docs.aws.amazon.com/STS/latest/APIReference/Welcome.html) and regions
+  * Add support for passing in a region to connect to via ``-r`` / ``--region``
+  * Add support for using STS to check resources in another account, including support for ``external_id``
+  * Major refactor of how service classes connect to AWS API
 
 0.1.2 (2015-08-13)
 ------------------

README.rst

@@ -63,9 +63,8 @@ Full project documentation is available at `http://awslimitchecker.readthedocs.o
 Status
 ------
 
-This project is currently in very early development. At this time please consider it alpha code and not reliable;
-furthermore its API may be changing rapidly. I hope to have this stabilized soon. I wouldn't call it ready for
-use, but contributions are certainly welcome.
+This project is currently in very early development. At this time please consider it beta code and not fully tested in all situations;
+furthermore its API may be changing rapidly. I hope to have this stabilized soon.
 
 What It Does
 ------------
@@ -77,6 +76,8 @@ What It Does
   exceed thresholds, and (CLI wrapper) exit non-0 if thresholds are exceeded
 - Define custom thresholds per-limit
 - where possible, pull current limits from Trusted Advisor API
+- Supports explicitly setting the AWS region
+- Supports using `STS <http://docs.aws.amazon.com/STS/latest/APIReference/Welcome.html>`_ to assume roles in other accounts, including using ``external_id``.
 
 Requirements
 ------------
@@ -101,14 +102,18 @@ system-wide, you can (using sudo).
 Credentials
 -----------
 
-awslimitchecker does nothing with AWS credentials, it leaves that to boto itself.
+Aside from STS, awslimitchecker does nothing with AWS credentials, it leaves that to boto itself.
 You must either have your credentials configured in one of boto's supported config
 files, or set as environment variables. See
 `boto config <http://docs.pythonboto.org/en/latest/boto_config_tut.html>`_
 and
 `this project's documentation <http://awslimitchecker.readthedocs.org/en/latest/getting_started.html#credentials>`_
 for further information.
 
+When using STS, you will need to specify the ``-r`` / ``--region`` option as well as the ``-A`` / ``--sts-account-id``
+and ``-R`` / ``--sts-account-role`` options to specify the Account ID that you want to assume a role in, and the
+name of the role you want to assume. If an external ID is required, you can specify it with ``-E`` / ``--external-id``.
+
 Usage
 -----
 

awslimitchecker/checker.py

@@ -48,7 +48,9 @@
 
 class AwsLimitChecker(object):
 
-    def __init__(self, warning_threshold=80, critical_threshold=99):
+    def __init__(self, warning_threshold=80, critical_threshold=99,
+                 account_id=None, account_role=None, region=None,
+                 external_id=None):
         """
         Main AwsLimitChecker class - this should be the only externally-used
         portion of awslimitchecker.
@@ -65,6 +67,22 @@ def __init__(self, warning_threshold=80, critical_threshold=99):
           integer percentage, for any limits without a specifically-set
           threshold.
         :type critical_threshold: int
+        :param account_id: `AWS Account ID <http://docs.aws.amazon.com/general/
+          latest/gr/acct-identifiers.html>`_
+          (12-digit string, currently numeric) for the account to connect to
+          (destination) via STS
+        :type account_id: str
+        :param account_role: the name of an
+          `IAM Role <http://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles.
+          html>`_
+          (in the destination account) to assume
+        :param region: AWS region name to connect to
+        :type region: str
+        :type account_role: str
+        :param external_id: (optional) the `External ID <http://docs.aws.amazon.
+          com/IAM/latest/UserGuide/id_roles_create_for-user_externalid.html>`_
+          string to use when assuming a role via STS.
+        :type external_id: str
         """
         # ###### IMPORTANT license notice ##########
         # Pursuant to Sections 5(b) and 13 of the GNU Affero General Public
@@ -91,10 +109,21 @@ def __init__(self, warning_threshold=80, critical_threshold=99):
         )
         self.warning_threshold = warning_threshold
         self.critical_threshold = critical_threshold
+        self.account_id = account_id
+        self.account_role = account_role
+        self.external_id = external_id
+        self.region = region
         self.services = {}
-        self.ta = TrustedAdvisor()
+        self.ta = TrustedAdvisor(
+            account_id=account_id,
+            account_role=account_role,
+            region=region,
+            external_id=external_id
+        )
         for sname, cls in _services.items():
-            self.services[sname] = cls(warning_threshold, critical_threshold)
+            self.services[sname] = cls(warning_threshold, critical_threshold,
+                                       account_id, account_role, region,
+                                       external_id)
 
     def get_version(self):
         """

awslimitchecker/connectable.py

@@ -0,0 +1,103 @@
+"""
+awslimitchecker/connectable.py
+
+The latest version of this package is available at:
+<https://github.com/jantman/awslimitchecker>
+
+################################################################################
+Copyright 2015 Jason Antman <jason@jasonantman.com> <http://www.jasonantman.com>
+
+    This file is part of awslimitchecker, also known as awslimitchecker.
+
+    awslimitchecker is free software: you can redistribute it and/or modify
+    it under the terms of the GNU Affero General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    awslimitchecker is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU Affero General Public License for more details.
+
+    You should have received a copy of the GNU Affero General Public License
+    along with awslimitchecker.  If not, see <http://www.gnu.org/licenses/>.
+
+The Copyright and Authors attributions contained herein may not be removed or
+otherwise altered, except to add the Author attribution of a contributor to
+this work. (Additional Terms pursuant to Section 7b of the AGPL v3)
+################################################################################
+While not legally required, I sincerely request that anyone who finds
+bugs please submit them at <https://github.com/jantman/awslimitchecker> or
+to me via email, and that you send any contributions or improvements
+either as a pull request on GitHub, or to me via email.
+################################################################################
+
+AUTHORS:
+Jason Antman <jason@jasonantman.com> <http://www.jasonantman.com>
+################################################################################
+"""
+
+import logging
+import boto.sts
+
+logger = logging.getLogger(__name__)
+
+
+class Connectable(object):
+
+    """
+    Mix-in helper class for connecting to AWS APIs. Centralizes logic of
+    connecting via regions and/or STS.
+    """
+
+    def connect_via(self, driver):
+        """
+        Connect to an AWS API and return the connection object. If
+        ``self.account_id`` is None, call ``driver(self.region)``. Otherwise,
+        call :py:meth:`~._get_sts_token` to get STS token credentials using
+        :py:meth:`boto.sts.STSConnection.assume_role` and call ``driver()`` with
+        those credentials to use an assumed role.
+
+        :param driver: the connect_to_region() function of the boto
+          submodule to use to create this connection
+        :type driver: :py:obj:`function`
+        :returns: connected boto service class instance
+        """
+        if self.account_id is not None:
+            logger.debug("Connecting to %s for account %s (STS; %s)",
+                         self.service_name, self.account_id, self.region)
+            self.credentials = self._get_sts_token()
+            conn = driver(
+                self.region,
+                aws_access_key_id=self.credentials.access_key,
+                aws_secret_access_key=self.credentials.secret_key,
+                security_token=self.credentials.session_token)
+        else:
+            logger.debug("Connecting to %s (%s)",
+                         self.service_name, self.region)
+            conn = driver(self.region)
+        logger.info("Connected to %s", self.service_name)
+        return conn
+
+    def _get_sts_token(self):
+        """
+        Assume a role via STS and return the credentials.
+
+        First connect to STS via :py:func:`boto.sts.connect_to_region`, then
+        assume a role using :py:meth:`boto.sts.STSConnection.assume_role`
+        using ``self.account_id`` and ``self.account_role`` (and optionally
+        ``self.external_id``). Return the resulting
+        :py:class:`boto.sts.credentials.Credentials` object.
+
+        :returns: STS assumed role credentials
+        :rtype: :py:class:`boto.sts.credentials.Credentials`
+        """
+        logger.debug("Connecting to STS in region %s", self.region)
+        sts = boto.sts.connect_to_region(self.region)
+        arn = "arn:aws:iam::%s:role/%s" % (self.account_id, self.account_role)
+        logger.debug("STS assume role for %s", arn)
+        role = sts.assume_role(arn, "awslimitchecker",
+                               external_id=self.external_id)
+        logger.debug("Got STS credentials for role; access_key_id=%s",
+                     role.credentials.access_key)
+        return role.credentials

awslimitchecker/runner.py

@@ -128,6 +128,20 @@ def parse_args(self, argv):
                        type=int, default=99,
                        help='default critical threshold (percentage of '
                        'limit); default: 99')
+        p.add_argument('-A', '--sts-account-id', action='store',
+                       type=str, default=None,
+                       help='for use with STS, the Account ID of the '
+                       'destination account (account to assume a role in)')
+        p.add_argument('-R', '--sts-account-role', action='store',
+                       type=str, default=None,
+                       help='for use with STS, the name of the IAM role to '
+                       'assume')
+        p.add_argument('-E', '--external-id', action='store', type=str,
+                       default=None, help='External ID to use when assuming '
+                       'a role via STS')
+        p.add_argument('-r', '--region', action='store',
+                       type=str, default=None,
+                       help='AWS region name to connect to; required for STS')
         p.add_argument('--skip-ta', action='store_true', default=False,
                        help='do not attempt to pull *any* information on limits'
                        ' from Trusted Advisor')
@@ -281,7 +295,11 @@ def console_entry_point(self):
         # the rest of these actually use the checker
         self.checker = AwsLimitChecker(
             warning_threshold=args.warning_threshold,
-            critical_threshold=args.critical_threshold
+            critical_threshold=args.critical_threshold,
+            account_id=args.sts_account_id,
+            account_role=args.sts_account_role,
+            region=args.region,
+            external_id=args.external_id
         )
 
         if args.version:

awslimitchecker/services/autoscaling.py

@@ -39,6 +39,7 @@
 
 import abc  # noqa
 import boto
+import boto.ec2.autoscale
 import logging
 
 from .base import _AwsService
@@ -52,11 +53,13 @@ class _AutoscalingService(_AwsService):
     service_name = 'AutoScaling'
 
     def connect(self):
-        """connect to API if not already connected; set self.conn"""
-        if self.conn is None:
-            logger.debug("Connecting to %s", self.service_name)
+        """Connect to API if not already connected; set self.conn."""
+        if self.conn is not None:
+            return
+        elif self.region:
+            self.conn = self.connect_via(boto.ec2.autoscale.connect_to_region)
+        else:
             self.conn = boto.connect_autoscale()
-            logger.info("Connected to %s", self.service_name)
 
     def find_usage(self):
         """

awslimitchecker/services/base.py

@@ -39,15 +39,18 @@
 
 import abc
 import logging
+from awslimitchecker.connectable import Connectable
+
 logger = logging.getLogger(__name__)
 
 
-class _AwsService(object):
+class _AwsService(Connectable):
     __metaclass__ = abc.ABCMeta
 
     service_name = 'baseclass'
 
-    def __init__(self, warning_threshold, critical_threshold):
+    def __init__(self, warning_threshold, critical_threshold, account_id=None,
+                 account_role=None, region=None, external_id=None):
         """
         Describes an AWS service and its limits, and provides methods to
         query current utilization.
@@ -65,9 +68,30 @@ def __init__(self, warning_threshold, critical_threshold):
           integer percentage, for any limits without a specifically-set
           threshold.
         :type critical_threshold: int
+        :param account_id: `AWS Account ID <http://docs.aws.amazon.com/general/
+          latest/gr/acct-identifiers.html>`_
+          (12-digit string, currently numeric) for the account to connect to
+          (destination) via STS
+        :type account_id: str
+        :param account_role: the name of an
+          `IAM Role <http://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles.
+          html>`_
+          (in the destination account) to assume
+        :param region: AWS region name to connect to
+        :type region: str
+        :type account_role: str
+        :param external_id: (optional) the `External ID <http://docs.aws.amazon.
+          com/IAM/latest/UserGuide/id_roles_create_for-user_externalid.html>`_
+          string to use when assuming a role via STS.
+        :type external_id: str
         """
         self.warning_threshold = warning_threshold
         self.critical_threshold = critical_threshold
+        self.account_id = account_id
+        self.account_role = account_role
+        self.region = region
+        self.external_id = external_id
+
         self.limits = {}
         self.limits = self.get_limits()
         self.conn = None
@@ -77,13 +101,18 @@ def __init__(self, warning_threshold, critical_threshold):
     def connect(self):
         """
         If not already done, establish a connection to the relevant AWS service
-        and save as ``self.conn``.
+        and save as ``self.conn``. If ``self.region`` is defined, call
+        ``self.connect_via()`` (:py:meth:`~.Connectable.connect_via`)
+        passing the appripriate boto ``connect_to_region()`` function as the
+        argument, else call the boto.connect_SERVICE_NAME() method directly.
         """
         """
-        if self.conn is None:
-            logger.debug("Connecting to %s", self.service_name)
-            # self.conn = boto.<connect to something>
-            logger.info("Connected to %s", self.service_name)
+        if self.conn is not None:
+            return
+        elif self.region:
+            self.conn = self.connect_via(boto.ec2.connect_to_region)
+        else:
+            self.conn = boto.connect_ec2()
         """
         raise NotImplementedError('abstract base class')
 

awslimitchecker/services/ebs.py

@@ -39,6 +39,7 @@
 
 import abc  # noqa
 import boto
+import boto.ec2
 import logging
 from .base import _AwsService
 from ..limit import AwsLimit
@@ -50,11 +51,13 @@ class _EbsService(_AwsService):
     service_name = 'EBS'
 
     def connect(self):
-        """connect to API if not already connected; set self.conn"""
-        if self.conn is None:
-            logger.debug("Connecting to %s", self.service_name)
+        """Connect to API if not already connected; set self.conn."""
+        if self.conn is not None:
+            return
+        elif self.region:
+            self.conn = self.connect_via(boto.ec2.connect_to_region)
+        else:
             self.conn = boto.connect_ec2()
-            logger.info("Connected to %s", self.service_name)
 
     def find_usage(self):
         """