"""AWS assume role utils."""

import boto3

from feed_ingestion.conf.config import BOTO3_CONFIG


def assumed_session(role_arn, session_name='feed_ingestion'):
    """Assume role_arn and return a boto3 Session with temporary credentials.

    Args:
        role_arn (str): The role ARN to assume.
        session_name (str): Assigned session name.

    Returns:
        boto3.Session with credentials from the assumed role.
    """
    session = boto3.session.Session()
    client = session.client('sts', config=BOTO3_CONFIG)
    credentials = client.assume_role(
        RoleArn=role_arn,
        RoleSessionName=session_name)['Credentials']
    return boto3.Session(
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken'],
    )


def get_s3_client_assume_role(role_arn):
    """Assume role_arn and return new boto3 Client connection.

    Args:
        role_arn (str): The role_arn.

    Returns:
        Connection boto3.session.Session.client.
    """
    session = assumed_session(role_arn)
    return session.client('s3')