Provide STS role like backup_monkey

.gitignore

@@ -33,3 +33,4 @@ nosetests.xml
 .mr.developer.cfg
 .project
 .pydevproject
+.idea

README.rst

@@ -29,6 +29,7 @@ Usage
 	                       EC2 instance this is running on). E.g. us-east-1
 	  --profile PROFILE    the profile to use to connect to EC2 (default is 'default',
 	                       see Boto docs for profile credential options)
+	  --role ROLE          the role (sts role) to use to connect to EC2')
 	  --verbose, -v        enable verbose output (-vvv for more)
 	  --version            display version number and exit
 	  --config CONFIG.YML  read a yaml configuration file.  specify tags to propagate without changing code.

graffiti_monkey/cli.py

@@ -31,6 +31,7 @@ class GraffitiMonkeyCli(object):
     def __init__(self):
         self.region = None
         self.profile = None
+        self.role = None
         self.monkey = None
         self.args = None
         self.config = {"_instance_tags_to_propagate": ['Name'],
@@ -65,6 +66,8 @@ def set_cli_args(self):
                             help='the region to tag things in (default is current region of EC2 instance this is running on). E.g. us-east-1')
         parser.add_argument('--profile', metavar='PROFILE',
                             help='the profile (credentials) to use to connect to EC2')
+        parser.add_argument('--role', metavar='ROLE',
+                            help='the role (sts role) to use to connect to EC2')
         parser.add_argument('--verbose', '-v', action='count',
                             help='enable verbose output (-vvv for more)')
         parser.add_argument('--version', action='version', version='%(prog)s ' + __version__,
@@ -136,6 +139,15 @@ def set_profile(self):
             self.profile = 'default'
         log.debug("Using profile: %s", self.profile)
 
+    def set_role(self):
+        if self.args.role:
+            self.role = self.args.role
+        elif "role" in self.config.keys():
+            self.role = self.config["role"]
+            log.debug("Using role: %s", self.role)
+        else:
+            log.debug("Not using role: %s", self.role)
+
     def set_dryrun(self):
         self.dryrun = self.args.dryrun
 
@@ -172,6 +184,7 @@ def config_default(self, key):
     def initialize_monkey(self):
         self.monkey = GraffitiMonkey(self.region,
                                      self.profile,
+                                     self.role,
                                      self.config["_instance_tags_to_propagate"],
                                      self.config["_volume_tags_to_propagate"],
                                      self.config_default("_volume_tags_to_be_set"),
@@ -201,6 +214,7 @@ def run(self):
         self.set_config()
         self.set_region()
         self.set_profile()
+        self.set_role()
         self.set_dryrun()
         self.set_append()
         self.set_volumes()

graffiti_monkey/core.py

@@ -26,7 +26,7 @@
 
 
 class GraffitiMonkey(object):
-    def __init__(self, region, profile, instance_tags_to_propagate, volume_tags_to_propagate, volume_tags_to_be_set, snapshot_tags_to_be_set, dryrun, append, volumes_to_tag, snapshots_to_tag, instance_filter, novolumes, nosnapshots):
+    def __init__(self, region, profile, role, instance_tags_to_propagate, volume_tags_to_propagate, volume_tags_to_be_set, snapshot_tags_to_be_set, dryrun, append, volumes_to_tag, snapshots_to_tag, instance_filter, novolumes, nosnapshots):
         # This list of tags associated with an EC2 instance to propagate to
         # attached EBS volumes
         self._instance_tags_to_propagate = instance_tags_to_propagate
@@ -47,6 +47,9 @@ def __init__(self, region, profile, instance_tags_to_propagate, volume_tags_to_p
         # The profile to use
         self._profile = profile
 
+        # The role to use if set
+        self._role = role
+
         # Whether this is a dryrun
         self._dryrun = dryrun
 
@@ -70,18 +73,41 @@ def __init__(self, region, profile, instance_tags_to_propagate, volume_tags_to_p
 
         log.info("Starting Graffiti Monkey")
         log.info("Options: dryrun %s, append %s, novolumes %s, nosnapshots %s", self._dryrun, self._append, self._novolumes, self._nosnapshots)
-        log.info("Connecting to region %s using profile %s", self._region, self._profile)
-        try:
-            self._conn = ec2.connect_to_region(self._region, profile_name=self._profile)
-        except boto.exception.NoAuthHandlerFound:
-            raise GraffitiMonkeyException('No AWS credentials found - check your credentials')
-        except boto.provider.ProfileNotFoundError:
-            log.info("Connecting to region %s using default credentials", self._region)
+        self._conn = self.get_connection()
+
+    def get_connection(self):
+        import boto
+        ret = None
+        if self._role:
+            from boto.sts import STSConnection
+            import boto
+            try:
+                sts = STSConnection()
+                assumed_role = sts.assume_role(role_arn=self._role, role_session_name='AssumeRoleSession')
+                ret = ec2.connect_to_region(
+                    self._region,
+                    aws_access_key_id=assumed_role.credentials.access_key,
+                    aws_secret_access_key=assumed_role.credentials.secret_key,
+                    security_token=assumed_role.credentials.session_token
+                )
+            except Exception,e:
+                print e
+                raise GraffitiMonkeyException('Cannot complete cross account access')
+        else:
+            log.info("Connecting to region %s using profile %s", self._region, self._profile)
             try:
-                self._conn = ec2.connect_to_region(self._region)
+                ret = ec2.connect_to_region(self._region, profile_name=self._profile)
             except boto.exception.NoAuthHandlerFound:
                 raise GraffitiMonkeyException('No AWS credentials found - check your credentials')
-
+            except boto.provider.ProfileNotFoundError:
+                log.info("Connecting to region %s using default credentials", self._region)
+                try:
+                    ret = ec2.connect_to_region(self._region)
+                except boto.exception.NoAuthHandlerFound:
+                    raise GraffitiMonkeyException('No AWS credentials found - check your credentials')
+        if not ret:
+            raise GraffitiMonkeyException('AWS connection failed - check your region and credentials')
+        return ret
 
     def propagate_tags(self):
         ''' Propagates tags by copying them from EC2 instance to EBS volume, and