outerbounds 0.3.55rc4__py3-none-any.whl → 0.3.89__py3-none-any.whl

outerbounds/command_groups/perimeters_cli.py

@@ -1,36 +1,33 @@
-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 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
 
-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 +56,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 +72,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 +91,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 +116,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 +165,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 +211,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 +232,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 +289,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 +354,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 +388,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 +428,20 @@ 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)
+
+
 def confirm_user_has_access_to_perimeter_or_fail(
     perimeter_id: str,
     perimeters: Dict[str, Any],
@@ -372,3 +464,46 @@ 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)
+
+    # GOOGLE_APPLICATION_CREDENTIALS is a well known gcloud environment variable
+    credentials_file_loc = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+
+    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 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/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,7 +512,7 @@ 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):
@@ -559,7 +555,24 @@ def show_relevant_links(config_dir=None, profile=None, perimeter_id="", output="
     show_links_response.add_or_update_data("links", [])
     links = []
     try:
-        metaflow_config = metaflowconfig.init_config(config_dir, profile)
+        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",

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,17 +1,66 @@
+from outerbounds._vendor import click
 import json
 import os
 import requests
 from os import path
-import requests
+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)
+
+    # 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
+
+    return config
+
+
+def init_config_from_url(config_dir, profile, url) -> Dict[str, str]:
+    config = read_metaflow_config_from_filesystem(config_dir, profile)
 
-def init_config(config_dir="", profile="") -> dict:
-    profile = profile or os.environ.get("METAFLOW_PROFILE")
-    config_dir = config_dir or os.path.expanduser(
-        os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")
+    if config is None or "METAFLOW_SERVICE_AUTH_KEY" not in config:
+        raise Exception("METAFLOW_SERVICE_AUTH_KEY not found in config file")
+
+    config_response = requests.get(
+        url,
+        headers={"x-api-key": f'{config["METAFLOW_SERVICE_AUTH_KEY"]}'},
     )
+    config_response.raise_for_status()
+    remote_config = config_response.json()["config"]
+    return remote_config
+
 
+def read_metaflow_config_from_filesystem(config_dir, profile) -> Dict[str, str]:
     config_filename = f"config_{profile}.json" if profile else "config.json"
+
     path_to_config = os.path.join(config_dir, config_filename)
 
     if os.path.exists(path_to_config):
@@ -19,22 +68,6 @@ def init_config(config_dir="", profile="") -> dict:
             config = json.load(json_file)
     else:
         raise Exception("Unable to locate metaflow config at '%s')" % (path_to_config))
-
-    # This is new remote-metaflow config; fetch it from the URL
-    if "OBP_METAFLOW_CONFIG_URL" in config:
-        if config is None or "METAFLOW_SERVICE_AUTH_KEY" not in config:
-            raise Exception("METAFLOW_SERVICE_AUTH_KEY not found in config file")
-
-        config_response = requests.get(
-            config["OBP_METAFLOW_CONFIG_URL"],
-            headers={"x-api-key": f'{config["METAFLOW_SERVICE_AUTH_KEY"]}'},
-        )
-        config_response.raise_for_status()
-        remote_config = config_response.json()["config"]
-        remote_config["METAFLOW_SERVICE_AUTH_KEY"] = config["METAFLOW_SERVICE_AUTH_KEY"]
-        return remote_config
-
-    # Legacy config, use from filesystem
     return config
 
 
@@ -70,3 +103,60 @@ def get_sanitized_url_from_config(config_dir: str, profile: str, key: str) -> st
 
     url_in_config = url_in_config.rstrip("/")
     return url_in_config
+
+
+def get_remote_metaflow_config_for_perimeter(
+    origin_token: str, perimeter: str, api_server: str
+):
+    try:
+        response = requests.get(
+            f"{api_server}/v1/perimeters/{perimeter}/metaflowconfigs/default",
+            headers={"x-api-key": origin_token},
+        )
+        response.raise_for_status()
+        config = response.json()["config"]
+        config["METAFLOW_SERVICE_AUTH_KEY"] = origin_token
+        return config
+    except Exception as e:
+        click.secho(
+            f"Failed to get metaflow config from {api_server}. Error: {str(e)}",
+            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

@@ -5,6 +5,7 @@ class OuterboundsCommandStatus(Enum):
     OK = "OK"
     FAIL = "FAIL"
     WARN = "WARN"
+    NOT_SUPPORTED = "NOT_SUPPORTED"
 
 
 class CommandStatus:
@@ -39,6 +40,11 @@ class OuterboundsCommandResponse:
         self.metadata = {}
         self._data = {}
 
+    def update(self, status, code, message):
+        self.status = status
+        self._code = code
+        self._message = message
+
     def add_or_update_metadata(self, key, value):
         self.metadata[key] = value
 

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)