Merge pull request #8 from lgtml/lgtml/assume_role

Added handy CLI arg to assume role

README.rst

@@ -77,6 +77,7 @@ Only the Instance ID is mandatory, but there are also other options you can use:
 * -o, --output : by default the script outputs log files and reports to the ``out`` folder. This option allows you to change it.
 * -n, --no-process : download log file(s), but do not process them with pgBadger.
 * -X, --pgbadger-args : command-line arguments to pass to pgBadger
+* --assume-role : By specifying a role you can use STS to assume a role, which is useful for cross account access with out having to setup the `.config` file. Format ``arn:aws:iam::<account_id>:<role_name>``
 
 Contribute
 ----------

package/rdspgbadger.py

@@ -47,6 +47,7 @@ def valid_date(s):
                     action='store_true')
 parser.add_argument('-d', '--date', help="get logs for given YYYY-MM-DD date",
                     type=valid_date)
+parser.add_argument('--assume-role', help="AWS STS AssumeRole")
 parser.add_argument('-r', '--region', help="AWS region")
 parser.add_argument('-o', '--output', help="Output folder for logs and report",
                     default='out')
@@ -71,11 +72,27 @@ def define_logger(verbose=False):
     logger.addHandler(consoleHandler)
 
 
-def get_all_logs(dbinstance_id, output, date=None, region=None):
+def get_all_logs(dbinstance_id, output,
+                 date=None, region=None, assume_role=None):
+
+    boto_args = {}
     if region:
-        client = boto3.client("rds", region_name=region)
-    else:
-        client = boto3.client("rds")
+        boto_args['region_name'] = region
+
+    if assume_role:
+        sts_client = boto3.client('sts')
+        assumedRoleObject = sts_client.assume_role(
+                RoleArn=assume_role,
+                RoleSessionName="RDSPGBadgerSession1"
+        )
+
+        credentials = assumedRoleObject['Credentials']
+        boto_args['aws_access_key_id'] = credentials['AccessKeyId']
+        boto_args['aws_secret_access_key'] = credentials['SecretAccessKey']
+        boto_args['aws_session_token'] = credentials['SessionToken']
+        logger.info('STS Assumed role {}'.format(assume_role))
+
+    client = boto3.client("rds", **boto_args)
     paginator = client.get_paginator("describe_db_log_files")
     response_iterator = paginator.paginate(
         DBInstanceIdentifier=dbinstance_id,
@@ -108,7 +125,8 @@ def write_log(client, dbinstance_id, filename, logfilename):
                 if exc.errno != errno.EEXIST:
                     raise
         with open(filename, "a") as logfile:
-            logfile.write(response["LogFileData"])
+            if 'LogFileData' in response:
+                logfile.write(response["LogFileData"])
 
         if not response["AdditionalDataPending"]:
             break
@@ -135,7 +153,13 @@ def main():
     logger.debug("pgbadger found")
 
     try:
-        get_all_logs(args.instance, args.output, args.date, args.region)
+        get_all_logs(
+                args.instance,
+                args.output,
+                date=args.date,
+                region=args.region,
+                assume_role=args.assume_role
+            )
     except (EndpointConnectionError, ClientError) as e:
         logger.error(e)
         exit(1)