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

@@ -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"
+    )

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

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

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")

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-{}",