outerbounds 0.3.55rc3__py3-none-any.whl → 0.3.133__py3-none-any.whl

outerbounds/command_groups/perimeters_cli.py

@@ -1,36 +1,35 @@
-import base64
-import hashlib
 import json
 import os
-import re
-import subprocess
 import sys
-import zlib
-from base64 import b64decode, b64encode
-from importlib.machinery import PathFinder
+from io import StringIO
 from os import path
-from pathlib import Path
-from typing import Any, Callable, Dict, List
-
-import boto3
-import click
+from typing import Any, Dict
+from outerbounds._vendor import click
 import requests
-from requests.exceptions import HTTPError
+import configparser
 
-from ..utils import kubeconfig, metaflowconfig
+from ..utils import metaflowconfig
+from ..utils.utils import safe_write_to_disk
 from ..utils.schema import (
     CommandStatus,
     OuterboundsCommandResponse,
     OuterboundsCommandStatus,
 )
 
+from .local_setup_cli import PERIMETER_CONFIG_URL_KEY
+
 
 @click.group()
 def cli(**kwargs):
     pass
 
 
-@cli.command(help="Switch current perimeter", hidden=True)
+@click.group(help="Manage perimeters")
+def perimeter(**kwargs):
+    pass
+
+
+@perimeter.command(help="Switch current perimeter")
 @click.option(
     "-d",
     "--config-dir",
@@ -59,7 +58,7 @@ def cli(**kwargs):
     help="Force change the existing perimeter",
     default=False,
 )
-def switch_perimeter(config_dir=None, profile=None, output="", id=None, force=False):
+def switch(config_dir=None, profile=None, output="", id=None, force=False):
     switch_perimeter_response = OuterboundsCommandResponse()
 
     switch_perimeter_step = CommandStatus(
@@ -75,7 +74,7 @@ def switch_perimeter(config_dir=None, profile=None, output="", id=None, force=Fa
         id, perimeters, output, switch_perimeter_response, switch_perimeter_step
     )
 
-    path_to_config = get_ob_config_file_path(config_dir, profile)
+    path_to_config = metaflowconfig.get_ob_config_file_path(config_dir, profile)
 
     import fcntl
 
@@ -94,7 +93,7 @@ def switch_perimeter(config_dir=None, profile=None, output="", id=None, force=Fa
 
         ob_config_dict = {
             "OB_CURRENT_PERIMETER": str(id),
-            "OB_CURRENT_PERIMETER_URL": perimeters[id]["remote_config_url"],
+            PERIMETER_CONFIG_URL_KEY: perimeters[id]["remote_config_url"],
         }
 
         # Now that we have the lock, we can safely write to the file
@@ -119,12 +118,34 @@ def switch_perimeter(config_dir=None, profile=None, output="", id=None, force=Fa
         )
 
     switch_perimeter_response.add_step(switch_perimeter_step)
+
+    ensure_cloud_creds_step = CommandStatus(
+        "EnsureCloudCredentials",
+        OuterboundsCommandStatus.OK,
+        "Cloud credentials were successfully updated.",
+    )
+
+    try:
+        ensure_cloud_credentials_for_shell(config_dir, profile)
+    except:
+        click.secho(
+            "Failed to update cloud credentials.",
+            fg="red",
+            err=True,
+        )
+        ensure_cloud_creds_step.update(
+            status=OuterboundsCommandStatus.FAIL,
+            reason="Failed to update cloud credentials.",
+            mitigation="",
+        )
+
+    switch_perimeter_response.add_step(ensure_cloud_creds_step)
+
     if output == "json":
         click.echo(json.dumps(switch_perimeter_response.as_dict(), indent=4))
-        return
 
 
-@cli.command(help="Show current perimeter", hidden=True)
+@perimeter.command(help="Show current perimeter")
 @click.option(
     "-d",
     "--config-dir",
@@ -146,7 +167,7 @@ def switch_perimeter(config_dir=None, profile=None, output="", id=None, force=Fa
     help="Show output in the specified format.",
     type=click.Choice(["json", ""]),
 )
-def show_current_perimeter(config_dir=None, profile=None, output=""):
+def show_current(config_dir=None, profile=None, output=""):
     show_current_perimeter_response = OuterboundsCommandResponse()
 
     show_current_perimeter_step = CommandStatus(
@@ -192,7 +213,7 @@ def show_current_perimeter(config_dir=None, profile=None, output=""):
         click.echo(json.dumps(show_current_perimeter_response.as_dict(), indent=4))
 
 
-@cli.command(help="List all available perimeters", hidden=True)
+@perimeter.command(help="List all available perimeters")
 @click.option(
     "-d",
     "--config-dir",
@@ -213,13 +234,29 @@ def show_current_perimeter(config_dir=None, profile=None, output=""):
     help="Show output in the specified format.",
     type=click.Choice(["json", ""]),
 )
-def list_perimeters(config_dir=None, profile=None, output=""):
+def list(config_dir=None, profile=None, output=""):
     list_perimeters_response = OuterboundsCommandResponse()
 
     list_perimeters_step = CommandStatus(
         "ListPerimeters", OuterboundsCommandStatus.OK, "Perimeter Fetch Successful."
     )
 
+    if "WORKSTATION_ID" in os.environ and (
+        "OBP_DEFAULT_PERIMETER" not in os.environ
+        or "OBP_DEFAULT_PERIMETER_URL" not in os.environ
+    ):
+        list_perimeters_response.update(
+            OuterboundsCommandStatus.NOT_SUPPORTED,
+            500,
+            "Perimeters are not supported on old workstations.",
+        )
+        click.secho(
+            "Perimeters are not supported on old workstations.", err=True, fg="red"
+        )
+        if output == "json":
+            click.echo(json.dumps(list_perimeters_response.as_dict(), indent=4))
+        return
+
     ob_config_dict = get_ob_config_or_fail_command(
         config_dir, profile, output, list_perimeters_response, list_perimeters_step
     )
@@ -254,6 +291,58 @@ def list_perimeters(config_dir=None, profile=None, output=""):
         click.echo(json.dumps(list_perimeters_response.as_dict(), indent=4))
 
 
+@perimeter.command(
+    help="Ensure credentials for cloud are synced with perimeter", hidden=True
+)
+@click.option(
+    "-d",
+    "--config-dir",
+    default=path.expanduser(os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")),
+    help="Path to Metaflow configuration directory",
+    show_default=True,
+)
+@click.option(
+    "-p",
+    "--profile",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
+    help="The named metaflow profile in which your workstation exists",
+)
+@click.option(
+    "-o",
+    "--output",
+    default="",
+    help="Show output in the specified format.",
+    type=click.Choice(["json", ""]),
+)
+def ensure_cloud_creds(config_dir=None, profile=None, output=""):
+    ensure_cloud_creds_step = CommandStatus(
+        "EnsureCloudCredentials",
+        OuterboundsCommandStatus.OK,
+        "Cloud credentials were successfully updated.",
+    )
+
+    ensure_cloud_creds_response = OuterboundsCommandResponse()
+
+    try:
+        ensure_cloud_credentials_for_shell(config_dir, profile)
+        click.secho("Cloud credentials updated successfully.", fg="green", err=True)
+    except:
+        click.secho(
+            "Failed to update cloud credentials.",
+            fg="red",
+            err=True,
+        )
+        ensure_cloud_creds_step.update(
+            status=OuterboundsCommandStatus.FAIL,
+            reason="Failed to update cloud credentials.",
+            mitigation="",
+        )
+
+    ensure_cloud_creds_response.add_step(ensure_cloud_creds_step)
+    if output == "json":
+        click.echo(json.dumps(ensure_cloud_creds_response.as_dict(), indent=4))
+
+
 def get_list_perimeters_api_response(config_dir, profile):
     metaflow_token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
     api_url = metaflowconfig.get_sanitized_url_from_config(
@@ -267,15 +356,6 @@ def get_list_perimeters_api_response(config_dir, profile):
     return perimeters_response.json()["perimeters"]
 
 
-def get_ob_config_file_path(config_dir: str, profile: str) -> str:
-    # If OBP_CONFIG_DIR is set, use that, otherwise use METAFLOW_HOME
-    # If neither are set, use ~/.metaflowconfig
-    obp_config_dir = path.expanduser(os.environ.get("OBP_CONFIG_DIR", config_dir))
-
-    ob_config_filename = f"ob_config_{profile}.json" if profile else "ob_config.json"
-    return os.path.expanduser(os.path.join(obp_config_dir, ob_config_filename))
-
-
 def get_perimeters_from_api_or_fail_command(
     config_dir: str,
     profile: str,
@@ -310,7 +390,7 @@ def get_ob_config_or_fail_command(
     command_response: OuterboundsCommandResponse,
     command_step: CommandStatus,
 ) -> Dict[str, str]:
-    path_to_config = get_ob_config_file_path(config_dir, profile)
+    path_to_config = metaflowconfig.get_ob_config_file_path(config_dir, profile)
 
     if not os.path.exists(path_to_config):
         click.secho(
@@ -350,6 +430,23 @@ def get_ob_config_or_fail_command(
     return ob_config_dict
 
 
+def ensure_cloud_credentials_for_shell(config_dir, profile):
+    if "WORKSTATION_ID" not in os.environ:
+        # Naive check to see if we're running in workstation. No need to ensure anything
+        # if this is not a workstation.
+        return
+
+    mf_config = metaflowconfig.init_config(config_dir, profile)
+
+    # Currently we only support GCP. TODO: utkarsh to add support for AWS and Azure
+    if "METAFLOW_DEFAULT_GCP_CLIENT_PROVIDER" in mf_config:
+        # This is a GCP deployment.
+        ensure_gcp_cloud_creds(config_dir, profile)
+    elif "METAFLOW_DEFAULT_AWS_CLIENT_PROVIDER" in mf_config:
+        # This is an AWS deployment.
+        ensure_aws_cloud_creds(config_dir, profile)
+
+
 def confirm_user_has_access_to_perimeter_or_fail(
     perimeter_id: str,
     perimeters: Dict[str, Any],
@@ -372,3 +469,104 @@ def confirm_user_has_access_to_perimeter_or_fail(
         if output == "json":
             click.echo(json.dumps(command_response.as_dict(), indent=4))
         sys.exit(1)
+
+
+def ensure_gcp_cloud_creds(config_dir, profile):
+    token_info = get_gcp_auth_credentials(config_dir, profile)
+    auth_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+    metaflow_token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+
+    try:
+        # GOOGLE_APPLICATION_CREDENTIALS is a well known gcloud environment variable
+        credentials_file_loc = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+    except KeyError:
+        # This is most likely an old workstation when these params weren't set. Do nothing.
+        # Alternatively, user might have deliberately unset it to use their own auth.
+        return
+
+    credentials_json = {
+        "type": "external_account",
+        "audience": f"//iam.googleapis.com/projects/{token_info['gcpProjectNumber']}/locations/global/workloadIdentityPools/{token_info['gcpWorkloadIdentityPool']}/providers/{token_info['gcpWorkloadIdentityPoolProvider']}",
+        "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+        "token_url": "https://sts.googleapis.com/v1/token",
+        "service_account_impersonation_url": f"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{token_info['gcpServiceAccountEmail']}:generateAccessToken",
+        "credential_source": {
+            "url": f"{auth_url}/generate/gcp",
+            "headers": {"x-api-key": metaflow_token},
+            "format": {"type": "json", "subject_token_field_name": "token"},
+        },
+    }
+
+    safe_write_to_disk(credentials_file_loc, json.dumps(credentials_json))
+
+
+def ensure_aws_cloud_creds(config_dir, profile):
+    token_info = get_aws_auth_credentials(config_dir, profile)
+
+    try:
+        token_file_loc = os.environ["OBP_AWS_WEB_IDENTITY_TOKEN_FILE"]
+
+        # AWS_CONFIG_FILE is a well known aws cli environment variable
+        config_file_loc = os.environ["AWS_CONFIG_FILE"]
+    except KeyError:
+        # This is most likely an old workstation when these params weren't set. Do nothing.
+        # Alternatively, user might have deliberately unset it to use their own auth.
+        return
+
+    aws_config = configparser.ConfigParser()
+    aws_config.read(config_file_loc)
+
+    aws_config["profile task"] = {
+        "role_arn": token_info["role_arn"],
+        "web_identity_token_file": token_file_loc,
+    }
+
+    if token_info.get("cspr_role_arn"):
+        # If CSPR role is present, then we need to use the task role (in the task profile)
+        # to assume the CSPR role.
+        aws_config["profile outerbounds"] = {
+            "role_arn": token_info["cspr_role_arn"],
+            "source_profile": "task",
+        }
+    else:
+        # If no CSPR role is present, just use the task profile as the outerbounds profile.
+        aws_config["profile outerbounds"] = aws_config["profile task"]
+
+    aws_config_string = StringIO()
+    aws_config.write(aws_config_string)
+
+    safe_write_to_disk(token_file_loc, token_info["token"])
+    safe_write_to_disk(config_file_loc, aws_config_string.getvalue())
+
+
+def get_aws_auth_credentials(config_dir, profile):  # pragma: no cover
+    token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+    auth_server_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+
+    response = requests.get(
+        "{}/generate/aws".format(auth_server_url), headers={"x-api-key": token}
+    )
+    response.raise_for_status()
+
+    return response.json()
+
+
+def get_gcp_auth_credentials(config_dir, profile):
+    token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+    auth_server_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+
+    response = requests.get(
+        "{}/generate/gcp".format(auth_server_url), headers={"x-api-key": token}
+    )
+    response.raise_for_status()
+
+    return response.json()
+
+
+cli.add_command(perimeter, name="perimeter")

outerbounds/command_groups/tutorials_cli.py

@@ -0,0 +1,111 @@
+import os
+from outerbounds._vendor import click
+import requests
+
+import tarfile
+import hashlib
+import tempfile
+
+
+@click.group()
+def cli(**kwargs):
+    pass
+
+
+@click.group(help="Manage tutorials curated by Outerbounds.", hidden=True)
+def tutorials(**kwargs):
+    pass
+
+
+@tutorials.command(help="Pull Outerbounds tutorials.")
+@click.option(
+    "--url",
+    required=True,
+    help="URL to pull the tutorials from.",
+    type=str,
+)
+@click.option(
+    "--destination-dir",
+    help="Show output in the specified format.",
+    type=str,
+    required=True,
+)
+@click.option(
+    "--force-overwrite",
+    is_flag=True,
+    help="Overwrite all existing files across all tutorials.",
+    type=bool,
+    required=False,
+    default=False,
+)
+def pull(url="", destination_dir="", force_overwrite=False):
+    try:
+        secure_download_and_extract(
+            url, destination_dir, force_overwrite=force_overwrite
+        )
+        click.secho("Tutorials pulled successfully.", fg="green", err=True)
+    except Exception as e:
+        print(e)
+        click.secho(f"Failed to pull tutorials: {e}", fg="red", err=True)
+
+
+def secure_download_and_extract(
+    url, dest_dir, expected_hash=None, force_overwrite=False
+):
+    """
+    Download a tar.gz file from a URL, verify its integrity, and extract its contents.
+
+    :param url: URL of the tar.gz file to download
+    :param dest_dir: Destination directory to extract the contents
+    :param expected_hash: Expected SHA256 hash of the file (optional)
+    """
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_file = os.path.join(
+            temp_dir, hashlib.md5(url.encode()).hexdigest() + ".tar.gz"
+        )
+
+        # Download the file
+        try:
+            response = requests.get(url, stream=True, verify=True)
+            response.raise_for_status()
+
+            with open(temp_file, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Failed to download file: {e}")
+
+        if expected_hash:
+            with open(temp_file, "rb") as f:
+                file_hash = hashlib.sha256(f.read()).hexdigest()
+            if file_hash != expected_hash:
+                raise Exception("File integrity check failed")
+
+        os.makedirs(dest_dir, exist_ok=True)
+
+        try:
+            with tarfile.open(temp_file, "r:gz") as tar:
+                # Keep track of new journeys to extract.
+                to_extract = []
+                members = tar.getmembers()
+                for member in members:
+                    member_path = os.path.join(dest_dir, member.name)
+                    # Check for any files trying to write outside the destination
+                    if not os.path.abspath(member_path).startswith(
+                        os.path.abspath(dest_dir)
+                    ):
+                        raise Exception("Attempted path traversal in tar file")
+                    if not os.path.exists(member_path):
+                        # The user might have modified the existing files, leave them untouched.
+                        to_extract.append(member)
+
+                if force_overwrite:
+                    tar.extractall(path=dest_dir)
+                else:
+                    tar.extractall(path=dest_dir, members=to_extract)
+        except tarfile.TarError as e:
+            raise Exception(f"Failed to extract tar file: {e}")
+
+
+cli.add_command(tutorials, name="tutorials")

outerbounds/command_groups/workstations_cli.py

@@ -1,5 +1,5 @@
-import click
-import yaml
+from outerbounds._vendor import click
+from outerbounds._vendor import yaml
 import requests
 import base64
 import datetime
@@ -19,6 +19,10 @@ from ..utils.schema import (
     OuterboundsCommandStatus,
 )
 from tempfile import NamedTemporaryFile
+from .perimeters_cli import (
+    get_perimeters_from_api_or_fail_command,
+    confirm_user_has_access_to_perimeter_or_fail,
+)
 
 KUBECTL_INSTALL_MITIGATION = "Please install kubectl manually from https://kubernetes.io/docs/tasks/tools/#kubectl"
 
@@ -89,7 +93,7 @@ def generate_workstation_token(config_dir=None, profile=None):
 @click.option(
     "-p",
     "--profile",
-    default="",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
     help="The named metaflow profile in which your workstation exists",
 )
 @click.option(
@@ -111,9 +115,6 @@ def configure_cloud_workstation(config_dir=None, profile=None, binary=None, outp
         "ConfigureKubeConfig", OuterboundsCommandStatus.OK, "Kubeconfig is configured"
     )
     try:
-        if not profile:
-            profile = metaflowconfig.get_metaflow_profile()
-
         metaflow_token = metaflowconfig.get_metaflow_token_from_config(
             config_dir, profile
         )
@@ -193,7 +194,7 @@ def configure_cloud_workstation(config_dir=None, profile=None, binary=None, outp
 @click.option(
     "-p",
     "--profile",
-    default="",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
     help="The named metaflow profile in which your workstation exists",
 )
 @click.option(
@@ -213,8 +214,6 @@ def list_workstations(config_dir=None, profile=None, output="json"):
     list_response.add_or_update_data("workstations", [])
 
     try:
-        if not profile:
-            profile = metaflowconfig.get_metaflow_profile()
         metaflow_token = metaflowconfig.get_metaflow_token_from_config(
             config_dir, profile
         )
@@ -260,7 +259,7 @@ def list_workstations(config_dir=None, profile=None, output="json"):
 @click.option(
     "-w",
     "--workstation",
-    default="",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
     help="The ID of the workstation to hibernate",
 )
 def hibernate_workstation(config_dir=None, profile=None, workstation=None):
@@ -308,7 +307,7 @@ def hibernate_workstation(config_dir=None, profile=None, workstation=None):
 @click.option(
     "-p",
     "--profile",
-    default="",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
     help="The named metaflow profile in which your workstation exists",
 )
 @click.option(
@@ -322,9 +321,6 @@ def restart_workstation(config_dir=None, profile=None, workstation=None):
         click.secho("Please specify a workstation ID", fg="red")
         return
     try:
-        if not profile:
-            profile = metaflowconfig.get_metaflow_profile()
-
         metaflow_token = metaflowconfig.get_metaflow_token_from_config(
             config_dir, profile
         )
@@ -516,8 +512,85 @@ def add_to_path(program_path, platform):
         with open(path_to_rc_file, "a+") as f:  # Open bashrc file
             if program_path not in f.read():
                 f.write("\n# Added by Outerbounds\n")
-                f.write(program_path)
+                f.write(f"export PATH=$PATH:{program_path}")
 
 
 def to_windows_path(path):
     return os.path.normpath(path).replace(os.sep, "\\")
+
+
+@cli.command(help="Show relevant links for a deployment & perimeter", hidden=True)
+@click.option(
+    "-d",
+    "--config-dir",
+    default=path.expanduser(os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")),
+    help="Path to Metaflow configuration directory",
+    show_default=True,
+)
+@click.option(
+    "-p",
+    "--profile",
+    default="",
+    help="The named metaflow profile in which your workstation exists",
+)
+@click.option(
+    "--perimeter-id",
+    default="",
+    help="The id of the perimeter to use",
+)
+@click.option(
+    "-o",
+    "--output",
+    default="",
+    help="Show output in the specified format.",
+    type=click.Choice(["json", ""]),
+)
+def show_relevant_links(config_dir=None, profile=None, perimeter_id="", output=""):
+    show_links_response = OuterboundsCommandResponse()
+    show_links_step = CommandStatus(
+        "showRelevantLinks",
+        OuterboundsCommandStatus.OK,
+        "Relevant links successfully fetched!",
+    )
+    show_links_response.add_or_update_data("links", [])
+    links = []
+    try:
+        if not perimeter_id:
+            metaflow_config = metaflowconfig.init_config(config_dir, profile)
+        else:
+            perimeters_dict = get_perimeters_from_api_or_fail_command(
+                config_dir, profile, output, show_links_response, show_links_step
+            )
+            confirm_user_has_access_to_perimeter_or_fail(
+                perimeter_id,
+                perimeters_dict,
+                output,
+                show_links_response,
+                show_links_step,
+            )
+
+            metaflow_config = metaflowconfig.init_config_from_url(
+                config_dir, profile, perimeters_dict[perimeter_id]["remote_config_url"]
+            )
+
+        links.append(
+            {
+                "id": "metaflow-ui-url",
+                "url": metaflow_config["METAFLOW_UI_URL"],
+                "label": "Metaflow UI URL",
+            }
+        )
+        show_links_response.add_or_update_data("links", links)
+        if output == "json":
+            click.echo(json.dumps(show_links_response.as_dict(), indent=4))
+    except Exception as e:
+        show_links_step.update(
+            OuterboundsCommandStatus.FAIL, "Failed to show relevant links", ""
+        )
+        show_links_response.add_step(show_links_step)
+        if output == "json":
+            show_links_response.add_or_update_data("error", str(e))
+            click.echo(json.dumps(show_links_response.as_dict(), indent=4))
+        else:
+            click.secho("Failed to show relevant links", fg="red", err=True)
+            click.secho("Error: {}".format(str(e)), fg="red", err=True)

outerbounds/utils/kubeconfig.py

@@ -1,6 +1,6 @@
 import os
-import yaml
-from yaml.scanner import ScannerError
+from outerbounds._vendor import yaml
+from outerbounds._vendor.yaml.scanner import ScannerError
 import subprocess
 from os import path
 import platform