assume_tool.py

import boto3
import argparse
import re
import botocore.exceptions
import itertools
import os

def generate_role_permutations(keywords):
    """
    Generate role name permutations from a list of keywords
    
    :param keywords: List of keywords to generate role names
    :return: List of generated role names
    """
    # Example role name patterns
    patterns = [
        "{0}",                  # Direct keyword
        "{0}Role",              # Append "Role"
        "{0}ExecutionRole",     # Append "ExecutionRole"
        "Admin{0}Role",         # Prefix "Admin"
        "{0}ReadOnlyRole"       # Append "ReadOnlyRole"
    ]
    
    # Generate all permutations
    role_names = []
    for keyword in keywords:
        for pattern in patterns:
            role_names.append(pattern.format(keyword))
    
    return role_names

def parse_role_input(role_input, override_account=None):
    """
    Parse role input to extract account ID and role name.
    
    :param role_input: Role name, full ARN, or account:role
    :param override_account: Optional account ID to override detected/parsed account
    :return: Tuple of (account_id, role_name)
    """
    # Full ARN pattern
    arn_pattern = r'^arn:aws:iam::(\d{12}):role/(.+)$'
    arn_match = re.match(arn_pattern, role_input)
    
    if arn_match:
        account_id = arn_match.group(1)
        role_name = arn_match.group(2)
        return override_account or account_id, role_name
    
    # Check if input is in account:role format
    account_role_pattern = r'^(\d{12}):(.+)$'
    account_role_match = re.match(account_role_pattern, role_input)
    
    if account_role_match:
        account_id = account_role_match.group(1)
        role_name = account_role_match.group(2)
        return override_account or account_id, role_name
    
    # If no account specified, try to get current account ID
    try:
        if override_account:
            return override_account, role_input
        
        sts_client = boto3.client('sts')
        account_id = sts_client.get_caller_identity()['Account']
        return account_id, role_input
    except Exception as e:
        raise ValueError(f"Could not determine account ID: {e}")

def read_file_lines(filepath):
    """
    Read lines from a file, stripping whitespace and removing empty lines
    
    :param filepath: Path to the file
    :return: List of non-empty, stripped lines
    """
    try:
        with open(filepath, 'r') as file:
            return [line.strip() for line in file if line.strip()]
    except FileNotFoundError:
        print(f"[!] ERROR: File {filepath} not found.")
        return []

def assume_role(account_id, role_name, profile):
    """Attempts to assume a given IAM role in the target AWS account using the specified profile."""
    role_arn = f"arn:aws:iam::{account_id}:role/{role_name}"
    
    # Create a session using the specified profile
    try:
        session = boto3.Session(profile_name=profile)
        sts_client = session.client("sts")
        
        try:
            response = sts_client.assume_role(
                RoleArn=role_arn,
                RoleSessionName="SystemHealthCheck"  # Innocuous session name
            )
            print(f"[+] SUCCESS: Assumed role {role_name} in account {account_id}")
            
            # Optionally, print out some details about the assumed role
            print("Temporary Credentials Details:")
            print(f"  Access Key ID: {response['Credentials']['AccessKeyId']}")
            print(f"  Expiration: {response['Credentials']['Expiration']}")
            
            return response
        
        except botocore.exceptions.ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == 'AccessDenied':
                print(f"[!] ACCESS DENIED: Unable to assume role {role_name} in account {account_id}")
            elif error_code == 'MalformedPolicyDocument':
                print(f"[!] ERROR: Malformed policy for role {role_name}")
            elif error_code == 'NoSuchEntity':
                print(f"[!] ERROR: Role {role_name} does not exist in account {account_id}")
            else:
                print(f"[!] ERROR: {e.response['Error']['Message']}")
    
    except Exception as e:
        print(f"[!] UNEXPECTED ERROR: {e}")

def main():
    parser = argparse.ArgumentParser(description="AWS IAM Role Assumption Tool")
    
    # Input sources for roles
    role_group = parser.add_mutually_exclusive_group()
    role_group.add_argument("--role", help="Specific role to assume (can be full ARN, rolename, or account:role)")
    role_group.add_argument("--role-file", help="File containing list of roles (one per line)")
    role_group.add_argument("--wordlist", help="File of keywords to generate role name permutations")
    
    # Input sources for accounts
    account_group = parser.add_mutually_exclusive_group()
    account_group.add_argument("--account", help="AWS Account ID to target")
    account_group.add_argument("--account-file", help="File containing list of account IDs (one per line)")
    
    parser.add_argument("--profile", default="default", help="AWS CLI profile to use (default: 'default').")
    
    args = parser.parse_args()
    
    # Determine how to proceed based on input
    try:
        # Prepare roles and accounts
        roles = []
        accounts = []
        
        # Process roles based on input type
        if args.role:
            roles = [args.role]
        elif args.role_file:
            roles = read_file_lines(args.role_file)
        elif args.wordlist:
            # Read keywords and generate role permutations
            keywords = read_file_lines(args.wordlist)
            roles = generate_role_permutations(keywords)
        
        # Process accounts
        if args.account:
            accounts = [args.account]
        elif args.account_file:
            accounts = read_file_lines(args.account_file)
        
        # If no accounts specified, use current account
        if not accounts:
            sts_client = boto3.client('sts')
            accounts = [sts_client.get_caller_identity()['Account']]
        
        # If no roles specified, request input
        if not roles:
            print("[!] ERROR: No roles specified. Use --role, --role-file, or --wordlist.")
            return
        
        # Print out the number of role variations being tested
        if args.wordlist:
            print(f"[*] Generated {len(roles)} role variations from {len(keywords)} keywords")
        
        # Attempt to assume roles across accounts
        for account, role in itertools.product(accounts, roles):
            try:
                # Parse the role input, using optional account
                account_id, role_name = parse_role_input(role, account)
                
                # Assume the specified role
                assume_role(account_id, role_name, args.profile)
            
            except ValueError as e:
                print(f"[!] ERROR processing {role} in account {account}: {e}")
    
    except Exception as e:
        print(f"[!] UNEXPECTED ERROR: {e}")

if __name__ == "__main__":
    main()