Assume Role within A Lambda function (Python)

lambda_assume_role.py

import boto3

# you can assign role in the function like below
# ROLE_ARN = 'arn:aws:iam::01234567890:role/my_role'
#
# or you can pass role as an evironment varibale
# ROLE_ARN = os.environ['role_arn']

ROLE_ARN = = os.environ['role_arn']

def aws_session(role_arn=None, session_name='my_session'):
    """
    If role_arn is given assumes a role and returns boto3 session
    otherwise return a regular session with the current IAM user/role
    """
    if role_arn:
        client = boto3.client('sts')
        response = client.assume_role(RoleArn=role_arn, RoleSessionName=session_name)
        session = boto3.Session(
            aws_access_key_id=response['Credentials']['AccessKeyId'],
            aws_secret_access_key=response['Credentials']['SecretAccessKey'],
            aws_session_token=response['Credentials']['SessionToken'])
        return session
    else:
        return boto3.Session()

def lambda_handler(event, context):
    session_assumed = aws_session(role_arn=ROLE_ARN, session_name='my_lambda')
    session_regular = aws_session()
  
    print(session_assumed.client('sts').get_caller_identity()['Account'])
    print(session_regular.client('sts').get_caller_identity()['Account'])

I found this very useful and well written. Thank you 👍

I am getting following error can you help me?

{
    "errorMessage": "An error occurred (InvalidClientTokenId) when calling the AssumeRole operation: The security token included in the request is invalid",
    "errorType": "ClientError",
    "stackTrace": [
        [
            "/var/task/lambda_function.py",
            28,
            "lambda_handler",
            "session_assumed = aws_session(role_arn=ROLE_ARN, session_name='my_lambda')"
        ],
        [
            "/var/task/lambda_function.py",
            18,
            "aws_session",
            "response = client.assume_role(RoleArn=role_arn, RoleSessionName=session_name)"
        ],
        [
            "/var/runtime/botocore/client.py",
            314,
            "_api_call",
            "return self._make_api_call(operation_name, kwargs)"
        ],
        [
            "/var/runtime/botocore/client.py",
            612,
            "_make_api_call",
            "raise error_class(parsed_response, operation_name)"
        ]
    ]
}

Could you share how the role arn:aws:iam::01234567890:role/my_role looks like, please?