Source code for schrodinger.stepper.cloud.aws_client

"""
Helper module for returning an AWS client object.
"""

try:
    import boto3
    from botocore.exceptions import ProfileNotFound
except ImportError:
    boto3 = None
    ProfileNotFound = Exception

try:
    from ec2_metadata import ec2_metadata
except ImportError:
    ec2_metadata = None

from schrodinger.application.steps import env_keys

CONNECTION_REDIRECT_MSG = """
AWS default connection setup failed to find credentials or config file. 
Attempting to retrieve environment variables."""


def get_client(aws_service):
    """
    Retrieve an AWS client for the requested AWS service. AWS credentials
    are obtained in one of the following manner and order:
    1. Shared credential file (~/.aws/credentials)
    2. AWS config file (~/.aws/config)
    3. Environment variables: `SCHRODINGER_AWS_KEY_ID` and `SCHRODINGER_AWS_KEY`

    See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
    for more details.

    :param aws_service: name of AWS service to connect to.
    :type aws_service: str

    :return: an active client connected to requested AWS service.
    :rtype: `boto3.Session.client`
    """
    try:
        session = get_session()
        client = session.client(aws_service)
    except ProfileNotFound:
        print(CONNECTION_REDIRECT_MSG)
        client = boto3.client(
            aws_service,
            aws_access_key_id=env_keys.SCHRODINGER_AWS_KEY,
            aws_secret_access_key=env_keys.SCHRODINGER_AWS_SECRET_KEY,
            region_name=env_keys.REGION)
    return client


def get_resource(aws_service):
    """
    Retrieve an AWS resource for the requested AWS service. See
    `aws_client.get_client()` for details on how credentials are obtained.

    :param aws_service: name of AWS service to connect to.
    :type aws_service: str

    :return: an active resource connected to requested AWS service.
    :rtype: `boto3.Session.resource`
    """
    try:
        session = get_session()
        resource = session.resource(aws_service)
    except ProfileNotFound:
        print(CONNECTION_REDIRECT_MSG)
        resource = boto3.resource(
            aws_service,
            aws_access_key_id=env_keys.SCHRODINGER_AWS_KEY,
            aws_secret_access_key=env_keys.SCHRODINGER_AWS_SECRET_KEY,
            region_name=env_keys.REGION)
    return resource


def get_session():
    """
    Retrieve an AWS active session. AWS credentials are obtained in one of the
    following manner and order:
    1. Shared credential file (~/.aws/credentials)
    2. AWS config file (~/.aws/config)

    :return: an active AWS session.
    :rtype: `boto3.Session`
    """
    return boto3.Session(profile_name=env_keys.AWS_PROFILE,
                         region_name=get_region())


def get_credentials():
    """
    This helper method usage is meant to retrieve aws credentials when they are
    set under ~/.aws. Note that this is typically encouraged only for
    development purposes and production services should instead use IAM roles
    for authentication across aws services.

    :return: boto3 active credentials object
    :rtype: `botocore.credential.Credential`
    """
    session = boto3.Session(profile_name=env_keys.AWS_PROFILE)
    return session.get_credentials()


def get_region():
    """
    Retrieves the AWS region such as `us-east-1` in the following manner and
    order:
    1. EC2 instance metadata (most production services will use this feature).
    2. Custom schrodinger environment variables
    If neither are available then `None` is returned. Note that most often
    credential profiles also encode the region property so a None value should
    not be a cause for concern.

    :return: AWS region
    :rtype: str or None
    """
    if ec2_metadata and ec2_metadata.region:
        return ec2_metadata.region
    return env_keys.REGION