From 2155cdae6282d58869fd2debfcf6004d4f0b1040 Mon Sep 17 00:00:00 2001 From: Preston Tamkin Date: Tue, 16 May 2023 12:05:34 -0700 Subject: [PATCH] feat!: Better integration with Cloud Companion: Login/Logout support and 'profile' override in bealine:// (#7) * feat: Add profile override support to the bealine:// protocol The 'profile picker' in the downloader looks at the Bealine config and matches the first profile(in the order they are in .bealine/config) that has that farm as default. This is not-precise. I saw some cases where a user had multiple profiles, but only the second one was valid. This is not likely to happen with 'real' customers, but is completely a non-obvious behavior to anyone. With this change, Cloud Companion(when its used) and set this profile name, ensuring the downloader is using the same profile as Cloud Companion. Of course, longer term we hope to put this downloader right into Cloud Companion, but for now seems worthwhile. Testing: - Added a unit test to validate the profile is overriden as expected * feat!: Add CloudCompanion support for Login/Logout BREAKING CHANGE: api.login and api.logout signatures have changed slightly 1. The return value of api.Login has changed. Itaswere and still is a string, but instead of assuming it's a SSO start URL it is now a more complete message for the user. Users no longer have to frame as 'SSO Start URL {return_value}' but use the value as the message 2. onPendingAuthorization now only uses kwargs. This is because SSO/CloudCompanion login types provide different information. The Login/Logout button now understand both Permission Sets and CloudCompanion. Cloud Companion will now 'owns' a section in .bealine/config, [cloud-companion]. It will write a key 'path' that is the path to its executable when it is used. Bealine already assumes that any AWS profile with the 'studio_id' key is managed by Cloud Companion. If the user attempts to 'login' then we use that path from the config and invoke CloudCompanion with a known set of args(login --profile {profileName}). We then wait for the credentials to become valid. Logout works similarly CloudCompanion has a corresponding change to add these login/logout commands. When 'login' is used, CloudCompanion pops up and logs into that profile. Once it logs in, it shows a dialog telling the user it will minimize which will then bring BeaLine back into focus There's a slight change to the return values of api.login/api.logout. Testing: * Only tested on mac so far. I will test on windows during PR review * I tested login/logout through the dev GUI for both a permissionset and Cloud Companion * New unit tests added and some updated * login/logout also work via the CLI --- src/bealine/api/_loginout.py | 162 +++++++++++++++--- src/bealine/api/_session.py | 13 +- .../cli/groups/handle_web_url_command.py | 21 ++- src/bealine/cli/groups/loginout_commands.py | 34 ++-- src/bealine/config/config_file.py | 4 + .../ui/dialogs/bealine_login_dialog.py | 62 ++++--- test/bealine/cli/test_cli_config.py | 3 +- test/bealine/cli/test_cli_handle_web_url.py | 19 +- test/bealine/cli/test_cli_loginout.py | 69 +++++++- test/bealine/shared_constants.py | 1 + 10 files changed, 306 insertions(+), 82 deletions(-) diff --git a/src/bealine/api/_loginout.py b/src/bealine/api/_loginout.py index 5f1cdcab6..40cfd0a25 100644 --- a/src/bealine/api/_loginout.py +++ b/src/bealine/api/_loginout.py @@ -8,6 +8,7 @@ from configparser import ConfigParser from logging import getLogger from typing import Callable, Optional +import subprocess from botocore.utils import JSONFileCache # type: ignore[import] @@ -20,40 +21,45 @@ sso_json_dumps, ) from ._awscli_botocore import SSOTokenFetcher -from ._session import get_boto3_session, invalidate_boto3_session_cache +from ._session import ( + get_boto3_session, + invalidate_boto3_session_cache, + get_credentials_type, + check_credentials_status, + AwsCredentialsType, + AwsCredentialsStatus, +) +from ..config import get_setting +import time logger = getLogger(__name__) -def login( +class UnsupportedProfileTypeForLoginLogout(Exception): + pass + + +def _login_sso( on_pending_authorization: Optional[Callable], on_cancellation_check: Optional[Callable], force_refresh: bool, config: Optional[ConfigParser] = None, -) -> str: - """ - Logs in to the provided session, if it is an Identity Center AWS profile - created with `aws configure sso`, or follows the instructions documented in - https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-sso.html#sso-configure-profile-manual. - - Args: - on_pending_authorization (Callable): A callback that receives the URI and code for the user to login. - on_cancellation_check (Callable): A callback that allows the operation to cancel before the token is received. - force_refresh (bool): If True, forces a credential refresh even if already logged in. - config (ConfigParser, optional): The BeaLine configuration - object to use instead of the config file. - """ +): session = get_boto3_session(config=config) base_sso = BaseSSOCommand(session._session) sso_config = base_sso._get_sso_config() token_cache = JSONFileCache(SSO_TOKEN_DIR, dumps_func=sso_json_dumps) + def pending_sso_authz(**kwargs): + if on_pending_authorization: + on_pending_authorization(credential_type=AwsCredentialsType.SSO_LOGIN, **kwargs) + token_fetcher = SSOTokenFetcher( sso_region=sso_config["sso_region"], client_creator=session._session.create_client, cache=token_cache, - on_pending_authorization=on_pending_authorization, + on_pending_authorization=pending_sso_authz, on_cancellation_check=on_cancellation_check, ) token_fetcher.fetch_token( @@ -62,10 +68,93 @@ def login( force_refresh=force_refresh, registration_scopes=sso_config.get("registration_scopes"), ) - return sso_config["sso_start_url"] + return f"AWS IAM Identity Center Start URL: {sso_config['sso_start_url']}" + + +def _login_cloud_companion( + on_pending_authorization: Optional[Callable], + on_cancellation_check: Optional[Callable], + config: Optional[ConfigParser] = None, +): + # Cloud Companion writes the absolute path to itself to the config file + cloud_companion_path = get_setting("cloud-companion.path", config=config) + profile_name = get_setting("defaults.aws_profile_name", config=config) + args = [cloud_companion_path, "login", "--profile", profile_name] + + # Open CloudCompanion, non-blocking the user will keep CloudCompanion running in the background. + try: + # We don't hookup to stdin but do this to avoid issues on windows + # See https://docs.python.org/3/library/subprocess.html#subprocess.STARTUPINFO.lpAttributeList + + p = subprocess.Popen( + args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, stdin=subprocess.PIPE + ) + except FileNotFoundError: + raise Exception( + f"Could not find Cloud Companion at {cloud_companion_path}. Please ensure Cloud Companion is installed correctly and setup the {profile_name} profile again." + ) + if on_pending_authorization: + on_pending_authorization(credential_type=AwsCredentialsType.CLOUD_COMPANION_LOGIN) + # And wait for the user to complete login + while True: + # Cloud Companion is a GUI app that will keep on running + # So we sit here and test that profile for validity until it works + if check_credentials_status(config) == AwsCredentialsStatus.AUTHENTICATED: + return f"Cloud Companion Profile: {profile_name}" + if on_cancellation_check: + # Check if the UI has signaled a cancel + if on_cancellation_check(): + p.kill() + raise Exception() + if p.poll(): + # Cloud Companion has stopped, we assume it returned us an error on one line on stderr + # but let's be specific about Cloud Companion failing incase the error is non-obvious + # and let's tack on stdout incase it helps + err_prefix = f"Cloud Companion was not able to log into the {profile_name} profile: " + out = p.stdout.read().decode("utf-8") if p.stdout else "" + raise Exception(f"{err_prefix}: {out}") + + time.sleep(0.5) + + +def login( + on_pending_authorization: Optional[Callable], + on_cancellation_check: Optional[Callable], + force_refresh: bool, + config: Optional[ConfigParser] = None, +) -> str: + """ + Logs in to the provided session if supported. + + This method AWS Identity Center Permission sets(e.g. aws configure sso): https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-sso.html#sso-configure-profile-manual. + and Nimble Cloud Companion + + If BeaLine doesn't know how to login to the the requested session Profile UnsupportedProfileTypeForLoginLogout is thrown + + Args: + on_pending_authorization (Callable): A callback that receives method-specific information to continue login. + All methods: 'credential_type' parameter of type AwsCredentialsType + For Permission Sets: + userCode: A one-time user verification code. This is needed to authorize an in-use device. + verificationUri: The URI of the verification page that takes the userCode to authorize the device. + verificationUriComplete: An alternate URL that the client can use to automatically launch a browser. This process skips the manual step in which the user visits the verification page and enters their code. + For Cloud Companion: No additional parameters + on_cancellation_check (Callable): A callback that allows the operation to cancel before the token is received. + force_refresh (bool): If True, forces a credential refresh even if already logged in. This only has an effect for AwsCredentialsType.SSO_LOGIN + config (ConfigParser, optional): The BeaLine configuration + object to use instead of the config file. + """ + credentials_type = get_credentials_type(config) + if credentials_type == AwsCredentialsType.SSO_LOGIN: + return _login_sso(on_pending_authorization, on_cancellation_check, force_refresh, config) + if credentials_type == AwsCredentialsType.CLOUD_COMPANION_LOGIN: + return _login_cloud_companion(on_pending_authorization, on_cancellation_check, config) + raise UnsupportedProfileTypeForLoginLogout( + "This action is only valid for AWS IAM Identity Center and Cloud Companion Profiles" + ) -def logout(config: Optional[ConfigParser] = None) -> None: +def logout(config: Optional[ConfigParser] = None) -> str: """ Removes all cached AWS SSO access tokens and any cached temporary AWS credentials retrieved with SSO access tokens across all @@ -76,8 +165,37 @@ def logout(config: Optional[ConfigParser] = None) -> None: object to use instead of the config file. """ session = get_boto3_session(config=config) - SSOTokenSweeper(session._session).delete_credentials(SSO_TOKEN_DIR) - SSOCredentialSweeper().delete_credentials(AWS_CREDS_CACHE_DIR) - # Force a refresh of the cached boto3 Session - invalidate_boto3_session_cache() + credentials_type = get_credentials_type(config) + if credentials_type == AwsCredentialsType.SSO_LOGIN: + SSOTokenSweeper(session._session).delete_credentials(SSO_TOKEN_DIR) + SSOCredentialSweeper().delete_credentials(AWS_CREDS_CACHE_DIR) + # Force a refresh of the cached boto3 Session + invalidate_boto3_session_cache() + return "You have been logged out" + if credentials_type == AwsCredentialsType.CLOUD_COMPANION_LOGIN: + # Cloud Companion writes the absolute path to itself to the config file + cloud_companion_path = get_setting("cloud-companion.path", config=config) + profile_name = get_setting("defaults.aws_profile_name", config=config) + args = [cloud_companion_path, "logout", "--profile", profile_name] + + # Open CloudCompanion, blocking + # Unlike login, that opens the regular Cloud Companion GUI, logout is a CLI command that clears the profile + # This makes it easier as we can execute and look at the return cdoe + try: + output = subprocess.check_output(args) + except FileNotFoundError: + raise Exception( + f"Could not find Cloud Companion at {cloud_companion_path}. Please ensure Cloud Companion is installed correctly and setup the {profile_name} profile again." + ) + except subprocess.CalledProcessError as e: + raise Exception( + f"Cloud Companion was unable to logout the profile {profile_name}. Return code {e.returncode}: {e.output}" + ) + + # Force a refresh of the cached boto3 Session + invalidate_boto3_session_cache() + return output.decode("utf8") + raise UnsupportedProfileTypeForLoginLogout( + "This action is only valid for AWS IAM Identity Center and Cloud Companion Profiles" + ) diff --git a/src/bealine/api/_session.py b/src/bealine/api/_session.py index bae4c360a..2e5bb98e5 100644 --- a/src/bealine/api/_session.py +++ b/src/bealine/api/_session.py @@ -30,7 +30,7 @@ class AwsCredentialsType(Enum): NOT_VALID = 0 SSO_LOGIN = 1 HOST_PROVIDED = 2 - CTDX_LOGIN = 3 + CLOUD_COMPANION_LOGIN = 3 class AwsCredentialsStatus(Enum): @@ -107,7 +107,7 @@ def get_boto3_client(service_name: str, config: Optional[ConfigParser] = None) - def get_credentials_type(config: Optional[ConfigParser] = None) -> AwsCredentialsType: """ Returns SSO_LOGIN if BeaLine login (i.e. IAM Identity Center - login) is available, CTDX_LOGIN if CTDX wrote the credentials, HOST_PROVIDED otherwise. + login) is available, CLOUD_COMPANION_LOGIN if CTDX wrote the credentials, HOST_PROVIDED otherwise. Args: config (ConfigParser, optional): The BeaLine configuration @@ -125,7 +125,7 @@ def get_credentials_type(config: Optional[ConfigParser] = None) -> AwsCredential return AwsCredentialsType.SSO_LOGIN if "studio_id" in profile_config: # CTDX adds some Nimble-specific keys here which we can use to know that this came from CTDX - return AwsCredentialsType.CTDX_LOGIN + return AwsCredentialsType.CLOUD_COMPANION_LOGIN else: return AwsCredentialsType.HOST_PROVIDED @@ -148,7 +148,7 @@ def user_identities(config: Optional[ConfigParser] = None) -> bool: # for IAM Identity Center login. return get_credentials_type(config=config) in [ AwsCredentialsType.SSO_LOGIN, - AwsCredentialsType.CTDX_LOGIN, + AwsCredentialsType.CLOUD_COMPANION_LOGIN, ] else: from ..config.config_file import str2bool @@ -198,4 +198,9 @@ def check_credentials_status(config: Optional[ConfigParser] = None) -> AwsCreden except (SSOTokenLoadError, UnauthorizedSSOTokenError): return AwsCredentialsStatus.NEEDS_LOGIN except Exception: + # We assume that the presence of a CloudCompanion profile + # means we will know everything necessary to start it and login. + + if get_credentials_type(config) == AwsCredentialsType.CLOUD_COMPANION_LOGIN: + return AwsCredentialsStatus.NEEDS_LOGIN return AwsCredentialsStatus.CONFIGURATION_ERROR diff --git a/src/bealine/cli/groups/handle_web_url_command.py b/src/bealine/cli/groups/handle_web_url_command.py index 16beb2c6d..983356411 100644 --- a/src/bealine/cli/groups/handle_web_url_command.py +++ b/src/bealine/cli/groups/handle_web_url_command.py @@ -100,26 +100,25 @@ def cli_handle_web_url( if split_url.netloc == "download-output": url_queries = parse_query_string( split_url.query, - parameter_names=[ - "farm-id", - "queue-id", - "job-id", - "step-id", - "task-id", - ], + parameter_names=["farm-id", "queue-id", "job-id", "step-id", "task-id", "profile"], required_parameter_names=["farm-id", "queue-id", "job-id"], ) # Validate the IDs - validate_resource_ids(url_queries) + # We copy the dict without the 'profile' key as that isn't a resource ID + validate_resource_ids({k: url_queries[k] for k in url_queries.keys() - {"profile"}}) job_id = url_queries.pop("job_id") step_id = url_queries.pop("step_id", None) task_id = url_queries.pop("task_id", None) - # Add the standard option "profile", choosing the best one based on farm and queue IDs - url_queries["profile"] = config_file.get_best_profile_for_farm( - url_queries["farm_id"], url_queries["queue_id"] + # Add the standard option "profile", using the one provided by the url(set by Cloud Companion) + # or choosing a best guess based on farm and queue IDs + url_queries["profile"] = url_queries.pop( + "profile", + config_file.get_best_profile_for_farm( + url_queries["farm_id"], url_queries["queue_id"] + ), ) # Get a temporary config object with the remaining standard options handled diff --git a/src/bealine/cli/groups/loginout_commands.py b/src/bealine/cli/groups/loginout_commands.py index a071f310e..4d1f2963a 100644 --- a/src/bealine/cli/groups/loginout_commands.py +++ b/src/bealine/cli/groups/loginout_commands.py @@ -9,22 +9,32 @@ import click from ... import api +from ...api._session import AwsCredentialsType from ...config import config_file from .._common import handle_error -def _cli_on_pending_authorization(userCode, verificationUri, verificationUriComplete, **kwargs): +def _cli_on_pending_authorization(**kwargs): """ Callback for `login`, to print out the Identity Center login uri and launch it in a web browser. """ - click.echo( - "\nAttempting to automatically open the SSO authorization page in your default browser.\n" - + "To authorize this request manually, open the following link:\n" - + f"\n{verificationUriComplete}\n" - + f"\nOr open {verificationUri} and enter the code: {userCode}" - ) - webbrowser.open_new_tab(verificationUriComplete) + + if kwargs["credential_type"] == AwsCredentialsType.SSO_LOGIN: + userCode, verificationUri, verificationUriComplete = ( + kwargs["userCode"], + kwargs["verificationUri"], + kwargs["verificationUriComplete"], + ) + click.echo( + "\nAttempting to automatically open the SSO authorization page in your default browser.\n" + + "To authorize this request manually, open the following link:\n" + + f"\n{verificationUriComplete}\n" + + f"\nOr open {verificationUri} and enter the code: {userCode}" + ) + webbrowser.open_new_tab(verificationUriComplete) + if kwargs["credential_type"] == AwsCredentialsType.CLOUD_COMPANION_LOGIN: + click.echo("Opening Cloud Companion. Please login before returning here.") @click.command(name="login") @@ -40,20 +50,20 @@ def cli_login(force): """ Logs in to the BeaLine configured AWS profile. - This is an IAM Permission Set configured for the BeaLine - CLI and submitters embedded in applications. + This is for any profile type that BeaLine knows how to login to + including AWS IAM IDC Permission Sets and Nimble Cloud Companion """ click.echo( f"Logging into AWS Profile {config_file.get_setting('defaults.aws_profile_name')!r} for BeaLine" ) - sso_start_url = api.login( + message = api.login( on_pending_authorization=_cli_on_pending_authorization, on_cancellation_check=None, force_refresh=force, ) - click.echo(f"\nSuccessfully logged into Start URL: {sso_start_url}\n") + click.echo(f"\nSuccessfully logged in: {message}\n") @click.command(name="logout") diff --git a/src/bealine/config/config_file.py b/src/bealine/config/config_file.py index 3ceb5fd97..a8591dc17 100644 --- a/src/bealine/config/config_file.py +++ b/src/bealine/config/config_file.py @@ -50,6 +50,10 @@ # section [profile-AmazonBealineCliAccess default] # "section_format" - How its value gets formatted into config file sections. SETTINGS: Dict[str, Dict[str, Any]] = { + # This is written by Cloud Companion and read by bealine + "cloud-companion.path": { + "default": "", + }, "defaults.aws_profile_name": { "default": DEFAULT_BEALINE_AWS_PROFILE_NAME, "section_format": "profile-{}", diff --git a/src/bealine/ui/dialogs/bealine_login_dialog.py b/src/bealine/ui/dialogs/bealine_login_dialog.py index 20886ffd9..20806e1f7 100644 --- a/src/bealine/ui/dialogs/bealine_login_dialog.py +++ b/src/bealine/ui/dialogs/bealine_login_dialog.py @@ -27,6 +27,7 @@ from ... import api from ...api._awscli_botocore import CanceledAuthorizationError +from ...api._session import AwsCredentialsType class BealineLoginDialog(QMessageBox): @@ -96,7 +97,7 @@ def __init__( self.buttonClicked.connect(self.on_button_clicked) self.setWindowTitle("Log in to BeaLine") - self.setText("Getting SSO authorization page from BeaLine...") + self.setText("Logging you in...") self.setStandardButtons(QMessageBox.Cancel) self._start_login() @@ -109,37 +110,45 @@ def _login_background_thread(self): """ try: - def on_pending_authorization( - userCode, verificationUri, verificationUriComplete, **kwargs - ): - # Qt doesn't theme URLs with style sheets, so grab the QPalette value - # to set it directly. - link_color = ( - QApplication.instance().palette().color(QPalette.Link).name(QColor.HexRgb) - ) - self.login_thread_message.emit( - "Attempting to automatically open the SSO authorization page " - + "in your default browser.

" - + "To authorize this request manually, open the following link:
" - + f'' - + f"{html.escape(verificationUriComplete)}

" - + f'Or open ' - + f"{html.escape(verificationUri)}
and enter the code: {html.escape(userCode)}" - ) - - # Try opening the link - webbrowser.open_new_tab(verificationUriComplete) + def on_pending_authorization(**kwargs): + if kwargs["credential_type"] == AwsCredentialsType.SSO_LOGIN: + userCode, verificationUri, verificationUriComplete = ( + kwargs["userCode"], + kwargs["verificationUri"], + kwargs["verificationUriComplete"], + ) + # Qt doesn't theme URLs with style sheets, so grab the QPalette value + # to set it directly. + link_color = ( + QApplication.instance().palette().color(QPalette.Link).name(QColor.HexRgb) + ) + self.login_thread_message.emit( + "Attempting to automatically open the SSO authorization page " + + "in your default browser.

" + + "To authorize this request manually, open the following link:
" + + f'' + + f"{html.escape(verificationUriComplete)}

" + + f'Or open ' + + f"{html.escape(verificationUri)}
and enter the code: {html.escape(userCode)}" + ) + + # Try opening the link + webbrowser.open_new_tab(verificationUriComplete) + elif kwargs["credential_type"] == AwsCredentialsType.CLOUD_COMPANION_LOGIN: + self.login_thread_message.emit( + "Opening Cloud Companion. Please login before returning here." + ) def on_cancellation_check(): return self.canceled - sso_start_url = api.login( + success_message = api.login( on_pending_authorization, on_cancellation_check, self.force_refresh, config=self.config, ) - self.login_thread_succeeded.emit(sso_start_url) + self.login_thread_succeeded.emit(success_message) except CanceledAuthorizationError as e: # If it wasn't canceled, send the exception to the dialog if not self.canceled: @@ -152,6 +161,7 @@ def _start_login(self) -> None: """ Starts the background login thread. """ + self.__login_thread = threading.Thread( target=self._login_background_thread, name="BeaLine Login Thread" ) @@ -173,7 +183,7 @@ def handle_login_thread_message(self, message: str) -> None: """ self.setText(message) - def handle_login_thread_succeeded(self, sso_start_url: str) -> None: + def handle_login_thread_succeeded(self, success_message: str) -> None: """ Handles the signal sent from the background login thread when the login has succeeded. @@ -184,9 +194,7 @@ def handle_login_thread_succeeded(self, sso_start_url: str) -> None: else: self.setStandardButtons(QMessageBox.Ok) self.setIcon(QMessageBox.Information) - self.setText( - f"Successfully logged into Start URL:

{html.escape(sso_start_url)}" - ) + self.setText(f"Successfully logged into:

{html.escape(success_message)}") def on_button_clicked(self, button): if self.standardButton(button) == QMessageBox.Cancel: diff --git a/test/bealine/cli/test_cli_config.py b/test/bealine/cli/test_cli_config.py index a3b06ffc2..c4436f5f7 100644 --- a/test/bealine/cli/test_cli_config.py +++ b/test/bealine/cli/test_cli_config.py @@ -34,7 +34,7 @@ def test_cli_config_show_defaults(fresh_bealine_config): assert fresh_bealine_config in result.output # Assert the expected number of settings - assert len(settings.keys()) == 8 + assert len(settings.keys()) == 9 for setting_name in settings.keys(): assert setting_name in result.output @@ -88,6 +88,7 @@ def test_cli_config_show_modified_config(fresh_bealine_config): Confirm that the CLI interface prints out all the configuration file data, when the configuration is default """ + config.set_setting("cloud-companion.path", "/User/auser/bin/CloudCompanion") config.set_setting("defaults.aws_profile_name", "EnvVarOverrideProfile") config.set_setting("settings.job_history_dir", "~/alternate/job_history") config.set_setting("settings.user_identities", "False") diff --git a/test/bealine/cli/test_cli_handle_web_url.py b/test/bealine/cli/test_cli_handle_web_url.py index 600ad1483..93b357189 100644 --- a/test/bealine/cli/test_cli_handle_web_url.py +++ b/test/bealine/cli/test_cli_handle_web_url.py @@ -23,7 +23,14 @@ from bealine_job_attachments.models import JobAttachmentS3Settings from ..api.test_job_bundle_submission import MOCK_GET_QUEUE_RESPONSE -from ..shared_constants import MOCK_FARM_ID, MOCK_JOB_ID, MOCK_QUEUE_ID, MOCK_STEP_ID, MOCK_TASK_ID +from ..shared_constants import ( + MOCK_FARM_ID, + MOCK_JOB_ID, + MOCK_QUEUE_ID, + MOCK_STEP_ID, + MOCK_TASK_ID, + MOCK_PROFILE_NAME, +) def test_parse_query_string(): @@ -278,11 +285,19 @@ def test_cli_handle_web_url_download_output_with_optional_input(fresh_bealine_co web_url = ( f"bealine://download-output?farm-id={MOCK_FARM_ID}&queue-id={MOCK_QUEUE_ID}&job-id={MOCK_JOB_ID}" - + f"&step-id={MOCK_STEP_ID}&task-id={MOCK_TASK_ID}" + + f"&step-id={MOCK_STEP_ID}&task-id={MOCK_TASK_ID}&profile={MOCK_PROFILE_NAME}" ) runner = CliRunner() result = runner.invoke(bealine_cli.cli, ["handle-web-url", web_url]) + assert result.exit_code == 0, result.output + + # call_args[1] is kwargs, but we need to use the tuple interface to be compatible back to 3.7 + captured_config = session_mock.call_args[1]["config"] + + assert ( + config.get_setting("defaults.aws_profile_name", captured_config) == "my-studio-profile" + ) mock_download.assert_called_once_with( s3_settings=JobAttachmentS3Settings(**MOCK_GET_QUEUE_RESPONSE["jobAttachmentSettings"]), # type: ignore diff --git a/test/bealine/cli/test_cli_loginout.py b/test/bealine/cli/test_cli_loginout.py index 1c5da3f4a..3c902cda1 100644 --- a/test/bealine/cli/test_cli_loginout.py +++ b/test/bealine/cli/test_cli_loginout.py @@ -5,6 +5,7 @@ """ import os import webbrowser +import subprocess from datetime import datetime from tempfile import TemporaryDirectory from unittest.mock import patch @@ -12,7 +13,7 @@ import pytest from click.testing import CliRunner -from bealine import api +from bealine import api, config from bealine.api import _awscli, _awscli_botocore, _loginout from bealine.cli import bealine_cli @@ -51,7 +52,9 @@ ), ], ) -def test_cli_loginout(fresh_bealine_config, profile_name, scoped_config, sessions_config): +def test_cli_permissionset_loginout( + fresh_bealine_config, profile_name, scoped_config, sessions_config +): """ Confirm that the CLI login command performs the expected sso-oicd handshake, then tries to cache the registration and token. @@ -59,6 +62,7 @@ def test_cli_loginout(fresh_bealine_config, profile_name, scoped_config, session Then, with the cached mocked credentials, calls the CLI logout command to confirm it clears them. """ + with patch.object(api._session, "get_boto3_session") as session_mock, patch.object( api, "get_boto3_session", new=session_mock ), patch.object(_loginout, "get_boto3_session", new=session_mock), patch.object( @@ -155,7 +159,7 @@ class SlowDownException(Exception): assert "https://verification-url/" in result.output assert "https://complete-verification-url.com/" in result.output assert ( - "Successfully logged into Start URL: https://d-012345abcd.awsapps.com/start" + "AWS IAM Identity Center Start URL: https://d-012345abcd.awsapps.com/start" in result.output ) assert result.exit_code == 0 @@ -180,3 +184,62 @@ class SlowDownException(Exception): # Verify that the logout call resets the cached session to None assert api._session.__cached_boto3_session is None + + +def test_cli_cloud_companion_login(fresh_bealine_config): + """ + Confirm that the CLI login/logout command invokes Cloud Companion as expected + """ + scoped_config = { + "credential_process": "/bin/NimbleCloudCompanion get-credentials --profile sandbox-us-west-2", + "studio_id": "us-west-2:stid-1g9neezauta8ease", + "region": "us-west-2", + } + + profile_name = "sandbox-us-west-2" + config.set_setting("cloud-companion.path", "/bin/NimbleCloudCompanion") + config.set_setting("defaults.aws_profile_name", profile_name) + + with patch.object(api._session, "get_boto3_session") as session_mock, patch.object( + api, "get_boto3_session", new=session_mock + ), patch.object(_loginout, "get_boto3_session", new=session_mock), patch.object( + subprocess, "Popen" + ) as popen_mock, patch.object( + subprocess, "check_output" + ) as check_output_mock: + # The profile name + session_mock().profile_name = profile_name + # This configuration includes the IdC profile + session_mock()._session.get_scoped_config.return_value = scoped_config + session_mock()._session.full_config = {"profiles": {profile_name: scoped_config}} + check_output_mock.return_value = bytes("Successfully logged out", "utf8") + + runner = CliRunner() + result = runner.invoke(bealine_cli.cli, ["login"]) + + popen_mock.assert_called_once_with( + ["/bin/NimbleCloudCompanion", "login", "--profile", "sandbox-us-west-2"], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + ) + + assert result.exit_code == 0 + # Sanity check that the session mock is being used + session_mock.assert_any_call() + + assert "Successfully logged in: Cloud Companion Profile: sandbox-us-west-2" in result.output + assert result.exit_code == 0 + + # Now lets logout + runner = CliRunner() + result = runner.invoke(bealine_cli.cli, ["logout"]) + + check_output_mock.assert_called_once_with( + ["/bin/NimbleCloudCompanion", "logout", "--profile", "sandbox-us-west-2"] + ) + + assert "Successfully logged out" in result.output + + # Verify that the logout call resets the cached session to None + assert api._session.__cached_boto3_session is None diff --git a/test/bealine/shared_constants.py b/test/bealine/shared_constants.py index 5ce68bbc7..e6736ffd7 100644 --- a/test/bealine/shared_constants.py +++ b/test/bealine/shared_constants.py @@ -6,3 +6,4 @@ MOCK_JOB_ID = "job-0123456789abcdefabcdefabcdefabcd" MOCK_STEP_ID = "step-0123456789abcdefabcdefabcdefabcd" MOCK_TASK_ID = "task-0123456789abcdefabcdefabcdefabcd" +MOCK_PROFILE_NAME = "my-studio-profile"