sso.py

# Copyright 2020 Ben Kehoe
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This code is based on the code for the AWS CLI v2"s `aws sso login` functionality
# https://github.com/aws/aws-cli/tree/v2/awscli/customizations/sso

import os
import sys
import webbrowser
import logging
import datetime
import uuid
import numbers
import typing
import json

import boto3
import botocore
from botocore.credentials import JSONFileCache
from botocore.credentials import SSOCredentialFetcher

from .format import format_account_id

from .vendored_botocore.utils import SSOTokenFetcher

from .exceptions import InvalidSSOConfigError, AuthDispatchError, AuthenticationNeededError
from .browser import OpenBrowserHandler, non_interactive_auth_raiser

SSO_TOKEN_DIR = os.path.expanduser(
    os.path.join("~", ".aws", "sso", "cache")
)

CREDENTIALS_CACHE_DIR = os.path.expanduser(
    os.path.join("~", ".aws", "cli", "cache")
)

LOGGER = logging.getLogger(__name__)

__all__ = ["get_boto3_session", "login", "list_available_accounts", "list_available_roles"]

# from customizations/sso/utils.py in AWS CLI v2
def _serialize_utc_timestamp(obj):
    if isinstance(obj, datetime.datetime):
        return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
    return obj
def _sso_json_dumps(obj):
    return json.dumps(obj, default=_serialize_utc_timestamp)

def get_token_fetcher(session, sso_region, *, interactive=False, sso_cache=None,
                     on_pending_authorization=None, message=None, outfile=None,
                     disable_browser=None, expiry_window=None):
    if hasattr(session, "_session"): #boto3 Session
        session = session._session

    if sso_cache is None:
        sso_cache = JSONFileCache(SSO_TOKEN_DIR, dumps_func=_sso_json_dumps)

    if on_pending_authorization is None:
        if interactive:
            on_pending_authorization = OpenBrowserHandler(
                outfile=outfile,
                message=message,
                disable_browser=disable_browser,
            )
        else:
            on_pending_authorization = non_interactive_auth_raiser

    token_fetcher = SSOTokenFetcher(
        sso_region=sso_region,
        client_creator=session.create_client,
        cache=sso_cache,
        on_pending_authorization=on_pending_authorization,
        expiry_window=expiry_window,
    )
    return token_fetcher

def _get_token_loader_from_token_fetcher(token_fetcher, force_refresh=False):
    def token_loader(start_url):
        token_response = token_fetcher.fetch_token(
            start_url=start_url,
            force_refresh=force_refresh
        )
        LOGGER.debug("TOKEN: {}".format(token_response))
        return token_response

    return token_loader


def get_credentials(session, start_url, sso_region, account_id, role_name, *,
        token_fetcher=None,
        force_refresh=False,
        sso_cache=None,
        credential_cache=None):
    """Return credentials for the given role.

    The return value is a dict containing the assumed role credentials.

    You probably want to use get_boto3_session() instead, which returns
    a boto3 session that caches its credentials and automatically refreshes them.
    """
    if hasattr(session, "_session"): #boto3 Session
        session = session._session
    if not token_fetcher:
        token_fetcher = get_token_fetcher(session, sso_region, sso_cache=sso_cache)

    token_loader = _get_token_loader_from_token_fetcher(token_fetcher=token_fetcher, force_refresh=force_refresh)

    if credential_cache is None:
        credential_cache = JSONFileCache(CREDENTIALS_CACHE_DIR)

    credential_fetcher = SSOCredentialFetcher(
        start_url=start_url,
        sso_region=sso_region,
        role_name=role_name,
        account_id=account_id,
        client_creator=session.create_client,
        cache=credential_cache,
        token_loader=token_loader,
    )

    return credential_fetcher.fetch_credentials()

def _get_botocore_session(
        start_url,
        sso_region,
        account_id,
        role_name,
        credential_cache=None,
        sso_cache=None,
        ):
    profile_name = str(uuid.uuid4())
    botocore_session = botocore.session.Session(session_vars={
        'profile': (None, None, None, None),
        # 'region': (None, None, None, None), # allow region to be set by env var
    })

    load_config = lambda: {
        "profiles": {
            profile_name: {
                "sso_start_url": start_url,
                "sso_region": sso_region,
                "sso_account_id": account_id,
                "sso_role_name": role_name,
            }
        }
    }
    sso_provider = botocore.credentials.SSOProvider(
        load_config=load_config,
        client_creator=botocore_session.create_client,
        profile_name=profile_name,
        cache=credential_cache,
        token_cache=sso_cache,
    )

    botocore_session.register_component(
        "credential_provider",
        botocore.credentials.CredentialResolver([sso_provider])
    )

    return botocore_session

def get_boto3_session(
        start_url: str,
        sso_region: str,
        account_id: typing.Union[str, int],
        role_name: str,
        *,
        region: str,
        login: bool=False,
        sso_cache=None,
        credential_cache=None) -> boto3.Session:
    """Get a boto3 session with the input configuration.

    Args:
        start_url (str): The start URL for the AWS SSO instance.
        sso_region (str): The AWS region for the AWS SSO instance.
        account_id (str): The AWS account ID to use.
        role_name (str): The AWS SSO role (aka Permission Set) name to use.
        region (str): The AWS region for the boto3 session.
        login (bool): Interactively log in the user if their AWS SSO credentials have expired.
        sso_cache: A dict-like object for AWS SSO credential caching to replace
            the default file cache in ~/.aws/sso/cache .
        credential_cache: A dict-like object to cache the role credentials in to
            replace the default in-memory cache.

    Returns:
        A boto3 Session object configured for the account and role.
    """
    account_id = format_account_id(account_id)

    if login:
        _login(start_url, sso_region, sso_cache=sso_cache)

    botocore_session = _get_botocore_session(start_url, sso_region, account_id, role_name,
        credential_cache=credential_cache,
        sso_cache=sso_cache)

    session = boto3.Session(botocore_session=botocore_session, region_name=region)

    return session

def login(
        start_url: str,
        sso_region: str,
        *,
        force_refresh: bool=False,
        expiry_window=None,
        disable_browser: bool=None,
        message: str=None,
        outfile: typing.Union[typing.TextIO, bool]=None,
        user_auth_handler=None,
        sso_cache=None,) -> typing.Dict:
    """Interactively log in the user if their AWS SSO credentials have expired.

    If the user is not logged in or force_refresh is True, it will attempt to log in.
    If the user is logged in and force_refresh is False, no action is taken.

    If disable_browser is True, a message will be printed to stderr
    with a URL and code for the user to log in with.
    Otherwise, it will attempt to automatically open the user's browser
    to log in, as well as printing the URL and code to stderr as a fallback.

    A custom message can be printed by setting message to a template string
    using {url} and {code} as placeholders.
    The message can be suppressed by setting message to False.

    Args:
        start_url (str): The start URL for the AWS SSO instance.
        sso_region (str): The AWS region for the AWS SSO instance.
        force_refresh (bool): Always go through the authentication process.
        expiry_window: A datetime.timedelta (or number of seconds),
            or callable returning such, specifying the minimum duration
            any existing token must be valid for.
        disable_browser (bool): Skip the browser popup
            and only print a message with the URL and code.
        message (str): A message template to print with the fallback URL and code.
        outfile (file): The file-like object to print the message to,
            or False to suppress the message.
        user_auth_handler (callable): override browser popup and message printing and
            use the given function instead. It must be a callable taking four
            keyword arguments: verificationUri, userCode, verificationUriComplete,
            and expiresAt (a datetime).
            Provide the verificationUri and userCode if the user is expected to
            type them in; verificationUriComplete has the userCode embedded in it,
            suitable for copying or browser popup. This function must return
            promptly or it will block the login process.
        sso_cache: A dict-like object for AWS SSO credential caching to replace
            the default file cache in ~/.aws/sso/cache .

    Returns:
        The token dict as returned by sso-oidc:CreateToken,
        which contains the actual authorization token, as well as the expiration.
        """
    session = botocore.session.Session(session_vars={
        'profile': (None, None, None, None),
        'region': (None, None, None, None),
    })

    on_pending_authorization = None
    if user_auth_handler:
        def on_pending_authorization(**kwargs):
            kwargs_to_pass = {}
            for key in ['verificationUri', 'userCode', 'verificationUriComplete', 'expiresAt']:
                if key in kwargs:
                    kwargs_to_pass[key] = kwargs[key]
            return user_auth_handler(**kwargs_to_pass)

    token_fetcher = get_token_fetcher(
        session=session,
        sso_region=sso_region,
        interactive=True,
        message=message,
        outfile=outfile,
        disable_browser=disable_browser,
        sso_cache=sso_cache,
        expiry_window=expiry_window,
        on_pending_authorization=on_pending_authorization)

    token = token_fetcher.fetch_token(
        start_url=start_url,
        force_refresh=force_refresh
    )
    token['expiresAt'] = _serialize_utc_timestamp(token['expiresAt'])
    if 'receivedAt' in token:
        token['receivedAt'] = _serialize_utc_timestamp(token['receivedAt'])

    return token

_login = login

def logout(
        start_url: str,
        sso_region: str,
        *,
        sso_cache=None):
    """Log out of the given AWS SSO instance.

    Args:
        start_url (str): The start URL for the AWS SSO instance.
        sso_region (str): The AWS region for the AWS SSO instance.
        sso_cache: A dict-like object for AWS SSO credential caching.

    Returns:
        Never raises.
        Returns True if a token was found and successfully logged out.
        Returns False if no token was found.
        If any exception is raised during the logout process,
            it is caught and returned.
    """

    session = botocore.session.Session(session_vars={
        'profile': (None, None, None, None),
        'region': (None, None, None, None),
    })

    token_fetcher = get_token_fetcher(
        session=session,
        sso_region=sso_region,
        sso_cache=sso_cache)

    try:
        token = token_fetcher.pop_token_from_cache(
            start_url=start_url
        )

        if not token:
            return False
        else:
            config = botocore.config.Config(
                region_name=sso_region,
                signature_version=botocore.UNSIGNED,
            )
            client = session.create_client("sso", config=config)

            client.logout(accessToken=token["accessToken"])
            return True
    except Exception as e:
        LOGGER.debug("Exception during logout", exc_info=True)
        return e

def list_available_accounts(
        start_url: str,
        sso_region: str,
        *,
        login: bool=False,
        sso_cache=None) -> typing.Iterator[typing.Tuple[str, str]]:
    """Iterate over the available accounts the user has access to through AWS SSO.

    Args:
        start_url (str): The start URL for the AWS SSO instance.
        sso_region (str): The AWS region for the AWS SSO instance.
        login (bool): Interactively log in the user if their AWS SSO credentials have expired.
        sso_cache: A dict-like object for AWS SSO credential caching.

    Returns:
        An iterator that yields account id and account name.
    """
    session = botocore.session.Session(session_vars={
        'profile': (None, None, None, None),
        'region': (None, None, None, None),
    })

    token_fetcher = get_token_fetcher(session, sso_region, interactive=login, sso_cache=sso_cache)

    token = token_fetcher.fetch_token(start_url)

    config = botocore.config.Config(
        region_name=sso_region,
        signature_version=botocore.UNSIGNED,
    )
    client = session.create_client("sso", config=config)

    list_accounts_args = {"accessToken": token["accessToken"]}
    while True:
        response = client.list_accounts(**list_accounts_args)

        for account in response["accountList"]:
            yield account["accountId"], account["accountName"]

        next_token = response.get("nextToken")
        if not next_token:
            break
        else:
            list_accounts_args["nextToken"] = response["nextToken"]

def list_available_roles(
        start_url: str,
        sso_region: str,
        account_id: typing.Union[str, int, typing.Iterable[typing.Union[str, int]]]=None,
        *,
        login: bool=False,
        sso_cache=None) -> typing.Iterator[typing.Tuple[str, str, str]]:
    """Iterate over the available accounts and roles the user has access to through AWS SSO.

    Args:
        start_url (str): The start URL for the AWS SSO instance.
        sso_region (str): The AWS region for the AWS SSO instance.
        account_id: Optional account id or list of account ids to check.
            If not set, all accounts available to the user are listed.
        login (bool): Interactively log in the user if their AWS SSO credentials have expired.
        sso_cache: A dict-like object for AWS SSO credential caching.

    Returns:
        An iterator that yields account id, account name, and role name.
        If the account(s) were provided in the input, the account name is always "UNKNOWN".
    """
    if account_id:
        if isinstance(account_id, (str, numbers.Number)):
            account_id_list = [format_account_id(account_id)]
        else:
            account_id_list = [format_account_id(v) for v in account_id] # type: ignore
    else:
        account_id_list = None

    session = botocore.session.Session(session_vars={
        'profile': (None, None, None, None),
        'region': (None, None, None, None),
    })

    token_fetcher = get_token_fetcher(session, sso_region, interactive=login, sso_cache=sso_cache)

    token = token_fetcher.fetch_token(start_url)

    config = botocore.config.Config(
        region_name=sso_region,
        signature_version=botocore.UNSIGNED,
    )
    client = session.create_client("sso", config=config)

    if account_id_list:
        def account_iterator():
            for acct in account_id_list:
                yield acct, "UNKNOWN"
    else:
        def account_iterator():
            list_accounts_args = {"accessToken": token["accessToken"]}
            while True:
                response = client.list_accounts(**list_accounts_args)

                for account in response["accountList"]:
                    yield account["accountId"], account["accountName"]

                next_token = response.get("nextToken")
                if not next_token:
                    break
                else:
                    list_accounts_args["nextToken"] = response["nextToken"]

    for account_id, account_name in account_iterator():
        list_role_args = {
            "accessToken": token["accessToken"],
            "accountId": account_id,
        }

        while True:
            response = client.list_account_roles(**list_role_args)

            for role in response["roleList"]:
                role_name = role["roleName"]

                yield account_id, account_name, role_name # type: ignore

            next_token = response.get("nextToken")
            if not next_token:
                break
            else:
                list_role_args["nextToken"] = response["nextToken"]