Added support for AWS Credentials profiles and enhanced Ctrl-C handling

enumerate-iam.py

@@ -1,24 +1,43 @@
 #!/usr/bin/env python
+import sys
 import argparse
 
+from boto3 import Session
 from enumerate_iam.main import enumerate_iam
 
-
 def main():
     parser = argparse.ArgumentParser(description='Enumerate IAM permissions')
 
-    parser.add_argument('--access-key', help='AWS access key', required=True)
-    parser.add_argument('--secret-key', help='AWS secret key', required=True)
+    parser.add_argument('--profile', help='AWS profile name fetched from credentials file. Specify this parameter or access-key and secret-key manually.')
+    parser.add_argument('--access-key', help='AWS access key if profile was not used')
+    parser.add_argument('--secret-key', help='AWS secret key if profile was not used')
     parser.add_argument('--session-token', help='STS session token')
     parser.add_argument('--region', help='AWS region to send API requests to', default='us-east-1')
 
     args = parser.parse_args()
 
-    enumerate_iam(args.access_key,
-                  args.secret_key,
-                  args.session_token,
+    if args.profile and (args.access_key or args.secret_key or args.session_token):
+        sys.stderr.write('error: Profile and raw AWS credential options are mutually exclusive.\n')
+        sys.stderr.write('       Please specify either --profile or --access-key and --secret-key.\n\n')
+        parser.print_help()
+        sys.exit(2)
+
+    access_key = args.access_key
+    secret_key = args.secret_key
+    session_token = args.session_token
+
+    if args.profile:
+        session = Session(profile_name = args.profile)
+        credentials = session.get_credentials()
+        currcreds = credentials.get_frozen_credentials()
+        access_key = currcreds.access_key
+        secret_key = currcreds.secret_key
+        session_token = currcreds.token
+
+    enumerate_iam(access_key,
+                  secret_key,
+                  session_token,
                   args.region)
 
-
 if __name__ == '__main__':
     main()

enumerate_iam/main.py

@@ -19,6 +19,7 @@
 import re
 import json
 import logging
+import signal
 import boto3
 import botocore
 import random
@@ -33,6 +34,7 @@
 
 MAX_THREADS = 25
 CLIENT_POOL = {}
+STOP_SIGNAL = False
 
 
 def report_arn(candidate):
@@ -62,17 +64,24 @@ def enumerate_using_bruteforce(access_key, secret_key, session_token, region):
     """
     Attempt to brute-force common describe calls.
     """
+    global STOP_SIGNAL
     output = dict()
 
     logger = logging.getLogger()
     logger.info('Attempting common-service describe / list brute force.')
 
+    # Ignore SIGINT signals so that child processes inherit SIGINT handler
+    original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
     pool = ThreadPool(MAX_THREADS)
+    signal.signal(signal.SIGINT, original_sigint_handler)
+
     args_generator = generate_args(access_key, secret_key, session_token, region)
 
     try:
-        results = pool.map(check_one_permission, args_generator)
+        results = pool.map_async(check_one_permission, args_generator)
+        results.get(600)
     except KeyboardInterrupt:
+        STOP_SIGNAL = True
         print('')
 
         results = []
@@ -85,6 +94,7 @@ def enumerate_using_bruteforce(access_key, secret_key, session_token, region):
             pool.join()
         except KeyboardInterrupt:
             print('')
+            STOP_SIGNAL = True
             return output
 
     for thread_result in results:
@@ -152,6 +162,9 @@ def check_one_permission(arg_tuple):
     access_key, secret_key, session_token, region, service_name, operation_name = arg_tuple
     logger = logging.getLogger()
 
+    if STOP_SIGNAL:
+        return
+
     service_client = get_client(access_key, secret_key, session_token, service_name, region)
     if service_client is None:
         return
@@ -167,6 +180,8 @@ def check_one_permission(arg_tuple):
     logger.debug('Testing %s.%s() in region %s' % (service_name, operation_name, region))
 
     try:
+        if STOP_SIGNAL: 
+            return
         action_response = action_function()
     except (botocore.exceptions.ClientError,
             botocore.exceptions.EndpointConnectionError,