feat(assume_role): support specifying session policy arns

trek10inc · Oct 16, 2024 · 3bd9759 · 3bd9759
1 parent 84a0ff9
commit 3bd9759
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/awsume/awsumepy/default_plugins.py b/awsume/awsumepy/default_plugins.py
@@ -149,6 +149,13 @@ def add_arguments(config: dict, parser: argparse.ArgumentParser):
         metavar='session_policy',
         help='Custom session policy JSON',
     )
+    parser.add_argument('--session-policy-arns',
+        nargs='+',
+        default=[],
+        dest='session_policy_arns',
+        metavar='session_policy_arns',
+        help='List of policy ARNs',
+    )
     parser.add_argument('--role-duration',
         action='store',
         dest='role_duration',
@@ -345,7 +352,17 @@ def assume_role_from_cli(config: dict, arguments: dict, profiles: dict):
     logger.debug('Session name: {}'.format(session_name))
     if not arguments.source_profile:
         logger.debug('Using current credentials to assume role')
-        role_session = aws_lib.assume_role({}, arguments.role_arn, session_name, session_policy=arguments.session_policy, region=region, external_id=arguments.external_id, role_duration=role_duration, tags=arguments.session_tags)
+        role_session = aws_lib.assume_role(
+            {},
+            arguments.role_arn,
+            session_name,
+            session_policy=arguments.session_policy,
+            session_policy_arns=arguments.session_policy_arns,
+            region=region,
+            external_id=arguments.external_id,
+            role_duration=role_duration,
+            tags=arguments.session_tags,
+        )
     else:
         logger.debug('Using the source_profile from the cli to call assume_role')
         source_profile = profiles.get(arguments.source_profile)
@@ -364,6 +381,7 @@ def assume_role_from_cli(config: dict, arguments: dict, profiles: dict):
                     arguments.role_arn,
                     session_name,
                     session_policy=arguments.session_policy,
+                    session_policy_arns=arguments.session_policy_arns,
                     region=region,
                     external_id=arguments.external_id,
                     role_duration=role_duration,
@@ -378,6 +396,7 @@ def assume_role_from_cli(config: dict, arguments: dict, profiles: dict):
                     arguments.role_arn,
                     session_name,
                     session_policy=arguments.session_policy,
+                    session_policy_arns=arguments.session_policy_arns,
                     region=region,
                     external_id=arguments.external_id,
                     role_duration=role_duration,
@@ -403,6 +422,7 @@ def assume_role_from_cli(config: dict, arguments: dict, profiles: dict):
                 arguments.role_arn,
                 session_name,
                 session_policy=arguments.session_policy,
+                session_policy_arns=arguments.session_policy_arns,
                 region=region,
                 external_id=arguments.external_id,
                 role_duration=role_duration,
@@ -423,6 +443,7 @@ def get_assume_role_credentials(config: dict, arguments: argparse.Namespace, pro
         target_profile.get('role_arn'),
         profile_lib.get_session_name(config, arguments, profiles, target_profile_name),
         session_policy=arguments.session_policy,
+        session_policy_arns=arguments.session_policy_arns,
         region=region,
         external_id=external_id,
         role_duration=role_duration,
@@ -469,6 +490,7 @@ def get_assume_role_credentials_mfa_required(config: dict, arguments: argparse.N
         target_profile.get('role_arn'),
         profile_lib.get_session_name(config, arguments, profiles, target_profile_name),
         session_policy=arguments.session_policy,
+        session_policy_arns=arguments.session_policy_arns,
         region=region,
         external_id=external_id,
         role_duration=role_duration,
@@ -500,6 +522,7 @@ def get_assume_role_credentials_mfa_required_large_custom_duration(config: dict,
         target_profile.get('role_arn'),
         profile_lib.get_session_name(config, arguments, profiles, target_profile_name),
         session_policy=arguments.session_policy,
+        session_policy_arns=arguments.session_policy_arns,
         region=region,
         external_id=external_id,
         role_duration=role_duration,

diff --git a/awsume/awsumepy/lib/aws.py b/awsume/awsumepy/lib/aws.py
@@ -1,5 +1,5 @@
 import os
-from typing import Union
+from typing import List, Union
 
 import boto3
 import botocore
@@ -33,6 +33,7 @@ def assume_role(
     role_arn: str,
     session_name: str,
     session_policy: str = None,
+    session_policy_arns: List[str] = [],
     external_id: str = None,
     region: str = None,
     role_duration: int = None,
@@ -56,6 +57,8 @@ def assume_role(
         kwargs = { 'RoleSessionName': session_name, 'RoleArn': role_arn }
         if session_policy:
             kwargs['Policy'] = session_policy
+        if session_policy_arns:
+            kwargs['PolicyArns'] = [{'arn': arn} for arn in session_policy_arns]
         if external_id:
             kwargs['ExternalId'] = external_id
         if role_duration: