import boto3


def aws_client(client_type, role_arn=None, session_name='my_session', session=None, region: str='us-east-1'):
    if session is None:
        _session = aws_session(role_arn=role_arn, session_name=session_name, region=region)
    else:
        _session = session
    client = _session.client(client_type, region_name=region)
    return client


def aws_session(role_arn=None, session_name='my_session', region: str='us-east-1', profile_name=None):
    """
    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(region_name=region, profile_name=profile_name)