outerbounds 0.3.68__py3-none-any.whl → 0.3.104__py3-none-any.whl

outerbounds/command_groups/perimeters_cli.py

@@ -1,22 +1,15 @@
-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 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,
@@ -81,7 +74,7 @@ def switch(config_dir=None, profile=None, output="", id=None, force=False):
         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
 
@@ -125,9 +118,31 @@ def switch(config_dir=None, profile=None, output="", id=None, force=False):
         )
 
     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
 
 
 @perimeter.command(help="Show current perimeter")
@@ -276,6 +291,58 @@ def list(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(
@@ -289,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,
@@ -332,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(
@@ -372,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],
@@ -396,4 +471,102 @@ def confirm_user_has_access_to_perimeter_or_fail(
         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

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

outerbounds/utils/metaflowconfig.py

@@ -1,23 +1,45 @@
-import click
+from outerbounds._vendor import click
 import json
 import os
 import requests
 from os import path
-from typing import Dict
+from typing import Dict, Union
 import sys
 
+"""
+key: perimeter specific URL to fetch the remote metaflow config from
+value: the remote metaflow config
+"""
+CACHED_REMOTE_METAFLOW_CONFIG: Dict[str, Dict[str, str]] = {}
+
+
+CURRENT_PERIMETER_KEY = "OB_CURRENT_PERIMETER"
+CURRENT_PERIMETER_URL = "OB_CURRENT_PERIMETER_MF_CONFIG_URL"
+CURRENT_PERIMETER_URL_LEGACY_KEY = (
+    "OB_CURRENT_PERIMETER_URL"  # For backwards compatibility with workstations.
+)
+
 
 def init_config(config_dir, profile) -> Dict[str, str]:
+    global CACHED_REMOTE_METAFLOW_CONFIG
     config = read_metaflow_config_from_filesystem(config_dir, profile)
 
-    # This is new remote-metaflow config; fetch it from the URL
-    if "OBP_METAFLOW_CONFIG_URL" in config:
-        remote_config = init_config_from_url(
-            config_dir, profile, config["OBP_METAFLOW_CONFIG_URL"]
-        )
-        remote_config["OBP_METAFLOW_CONFIG_URL"] = config["OBP_METAFLOW_CONFIG_URL"]
+    # Either user has an ob_config.json file with the perimeter URL
+    # or the default config on the filesystem has the config URL in it.
+    perimeter_specifc_url = get_perimeter_config_url_if_set_in_ob_config(
+        config_dir, profile
+    ) or config.get("OBP_METAFLOW_CONFIG_URL", "")
+
+    if perimeter_specifc_url != "":
+        if perimeter_specifc_url in CACHED_REMOTE_METAFLOW_CONFIG:
+            return CACHED_REMOTE_METAFLOW_CONFIG[perimeter_specifc_url]
+
+        remote_config = init_config_from_url(config_dir, profile, perimeter_specifc_url)
+        remote_config["OBP_METAFLOW_CONFIG_URL"] = perimeter_specifc_url
+
+        CACHED_REMOTE_METAFLOW_CONFIG[perimeter_specifc_url] = remote_config
         return remote_config
-    # Legacy config, use from filesystem
+
     return config
 
 
@@ -101,3 +123,40 @@ def get_remote_metaflow_config_for_perimeter(
             fg="red",
         )
         sys.exit(1)
+
+
+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_perimeter_config_url_if_set_in_ob_config(
+    config_dir: str, profile: str
+) -> Union[str, None]:
+    file_path = get_ob_config_file_path(config_dir, profile)
+
+    if os.path.exists(file_path):
+        with open(file_path, "r") as f:
+            ob_config = json.loads(f.read())
+
+        if CURRENT_PERIMETER_URL in ob_config:
+            return ob_config[CURRENT_PERIMETER_URL]
+        elif CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config:
+            return ob_config[CURRENT_PERIMETER_URL_LEGACY_KEY]
+        else:
+            raise ValueError(
+                "{} does not contain the key {}".format(
+                    file_path, CURRENT_PERIMETER_KEY
+                )
+            )
+    elif "OBP_CONFIG_DIR" in os.environ:
+        raise FileNotFoundError(
+            "Environment variable OBP_CONFIG_DIR is set to {} but this directory does not contain an ob_config.json file.".format(
+                os.environ["OBP_CONFIG_DIR"]
+            )
+        )
+    return None

outerbounds/utils/schema.py

@@ -59,14 +59,14 @@ class OuterboundsCommandResponse:
         if step.status == OuterboundsCommandStatus.FAIL:
             self.status = OuterboundsCommandStatus.FAIL
             self._code = 500
-            self._message = "We found one or more errors with your installation."
+            self._message = "Encountered an error when trying to run command."
         elif (
             step.status == OuterboundsCommandStatus.WARN
             and self.status != OuterboundsCommandStatus.FAIL
         ):
             self.status = OuterboundsCommandStatus.WARN
             self._code = 200
-            self._message = "We found one or more warnings with your installation."
+            self._message = "Encountered one or more warnings when running the command."
 
     def as_dict(self):
         self._data["steps"] = [step.as_dict() for step in self._steps]

outerbounds/utils/utils.py

@@ -0,0 +1,19 @@
+import tempfile
+import os
+
+"""
+Writes the given data to file_loc. Ensures that the directory exists
+and uses a temporary file to ensure that the file is written atomically.
+"""
+
+
+def safe_write_to_disk(file_loc, data):
+    # Ensure the directory exists
+    os.makedirs(os.path.dirname(file_loc), exist_ok=True)
+
+    with tempfile.NamedTemporaryFile(
+        "w", dir=os.path.dirname(file_loc), delete=False
+    ) as f:
+        f.write(data)
+        tmp_file = f.name
+    os.rename(tmp_file, file_loc)