def assume_role()

in functions/source/c1c_controltower_lifecycle.py [0:0]

• 33 lines of code
• 3 McCabe index (conditional complexity)

def assume_role(aws_account_number, role_name) -> boto3.Session:
    try:
        sts_client = boto3.client("sts")
        logger.info(f"Retrieving session for operation")
        logger.info(
            f"currently executing in "
            f'{sts_client.get_caller_identity()["Account"]};'
            f" called account is {aws_account_number}"
        )
        if sts_client.get_caller_identity()["Account"] == aws_account_number:
            logger.info(
                f"Target account is Control Tower Management; returning local credentials session"
            )
            return boto3.session.Session()
        partition = sts_client.get_caller_identity()["Arn"].split(":")[1]

        assume_role_response = sts_client.assume_role(
            RoleArn="arn:{}:iam::{}:role/{}".format(
                partition, aws_account_number, role_name
            ),
            RoleSessionName=str(aws_account_number + "-" + role_name),
        )
        sts_session = boto3.Session(
            aws_access_key_id=assume_role_response["Credentials"]["AccessKeyId"],
            aws_secret_access_key=assume_role_response["Credentials"][
                "SecretAccessKey"
            ],
            aws_session_token=assume_role_response["Credentials"]["SessionToken"],
        )
        logger.info(f"Assumed session for {aws_account_number} - {role_name}.")
        return sts_session
    except Exception as e:
        logger.info(f"Could not assume role : {e}")
        raise e