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

outerbounds/_vendor/yaml/serializer.py

@@ -0,0 +1,127 @@
+__all__ = ["Serializer", "SerializerError"]
+
+from .error import YAMLError
+from .events import *
+from .nodes import *
+
+
+class SerializerError(YAMLError):
+    pass
+
+
+class Serializer:
+
+    ANCHOR_TEMPLATE = "id%03d"
+
+    def __init__(
+        self,
+        encoding=None,
+        explicit_start=None,
+        explicit_end=None,
+        version=None,
+        tags=None,
+    ):
+        self.use_encoding = encoding
+        self.use_explicit_start = explicit_start
+        self.use_explicit_end = explicit_end
+        self.use_version = version
+        self.use_tags = tags
+        self.serialized_nodes = {}
+        self.anchors = {}
+        self.last_anchor_id = 0
+        self.closed = None
+
+    def open(self):
+        if self.closed is None:
+            self.emit(StreamStartEvent(encoding=self.use_encoding))
+            self.closed = False
+        elif self.closed:
+            raise SerializerError("serializer is closed")
+        else:
+            raise SerializerError("serializer is already opened")
+
+    def close(self):
+        if self.closed is None:
+            raise SerializerError("serializer is not opened")
+        elif not self.closed:
+            self.emit(StreamEndEvent())
+            self.closed = True
+
+    # def __del__(self):
+    #    self.close()
+
+    def serialize(self, node):
+        if self.closed is None:
+            raise SerializerError("serializer is not opened")
+        elif self.closed:
+            raise SerializerError("serializer is closed")
+        self.emit(
+            DocumentStartEvent(
+                explicit=self.use_explicit_start,
+                version=self.use_version,
+                tags=self.use_tags,
+            )
+        )
+        self.anchor_node(node)
+        self.serialize_node(node, None, None)
+        self.emit(DocumentEndEvent(explicit=self.use_explicit_end))
+        self.serialized_nodes = {}
+        self.anchors = {}
+        self.last_anchor_id = 0
+
+    def anchor_node(self, node):
+        if node in self.anchors:
+            if self.anchors[node] is None:
+                self.anchors[node] = self.generate_anchor(node)
+        else:
+            self.anchors[node] = None
+            if isinstance(node, SequenceNode):
+                for item in node.value:
+                    self.anchor_node(item)
+            elif isinstance(node, MappingNode):
+                for key, value in node.value:
+                    self.anchor_node(key)
+                    self.anchor_node(value)
+
+    def generate_anchor(self, node):
+        self.last_anchor_id += 1
+        return self.ANCHOR_TEMPLATE % self.last_anchor_id
+
+    def serialize_node(self, node, parent, index):
+        alias = self.anchors[node]
+        if node in self.serialized_nodes:
+            self.emit(AliasEvent(alias))
+        else:
+            self.serialized_nodes[node] = True
+            self.descend_resolver(parent, index)
+            if isinstance(node, ScalarNode):
+                detected_tag = self.resolve(ScalarNode, node.value, (True, False))
+                default_tag = self.resolve(ScalarNode, node.value, (False, True))
+                implicit = (node.tag == detected_tag), (node.tag == default_tag)
+                self.emit(
+                    ScalarEvent(alias, node.tag, implicit, node.value, style=node.style)
+                )
+            elif isinstance(node, SequenceNode):
+                implicit = node.tag == self.resolve(SequenceNode, node.value, True)
+                self.emit(
+                    SequenceStartEvent(
+                        alias, node.tag, implicit, flow_style=node.flow_style
+                    )
+                )
+                index = 0
+                for item in node.value:
+                    self.serialize_node(item, node, index)
+                    index += 1
+                self.emit(SequenceEndEvent())
+            elif isinstance(node, MappingNode):
+                implicit = node.tag == self.resolve(MappingNode, node.value, True)
+                self.emit(
+                    MappingStartEvent(
+                        alias, node.tag, implicit, flow_style=node.flow_style
+                    )
+                )
+                for key, value in node.value:
+                    self.serialize_node(key, node, None)
+                    self.serialize_node(value, node, key)
+                self.emit(MappingEndEvent())
+            self.ascend_resolver()

outerbounds/_vendor/yaml/tokens.py

@@ -0,0 +1,129 @@
+class Token(object):
+    def __init__(self, start_mark, end_mark):
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+
+    def __repr__(self):
+        attributes = [key for key in self.__dict__ if not key.endswith("_mark")]
+        attributes.sort()
+        arguments = ", ".join(
+            ["%s=%r" % (key, getattr(self, key)) for key in attributes]
+        )
+        return "%s(%s)" % (self.__class__.__name__, arguments)
+
+
+# class BOMToken(Token):
+#    id = '<byte order mark>'
+
+
+class DirectiveToken(Token):
+    id = "<directive>"
+
+    def __init__(self, name, value, start_mark, end_mark):
+        self.name = name
+        self.value = value
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+
+
+class DocumentStartToken(Token):
+    id = "<document start>"
+
+
+class DocumentEndToken(Token):
+    id = "<document end>"
+
+
+class StreamStartToken(Token):
+    id = "<stream start>"
+
+    def __init__(self, start_mark=None, end_mark=None, encoding=None):
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+        self.encoding = encoding
+
+
+class StreamEndToken(Token):
+    id = "<stream end>"
+
+
+class BlockSequenceStartToken(Token):
+    id = "<block sequence start>"
+
+
+class BlockMappingStartToken(Token):
+    id = "<block mapping start>"
+
+
+class BlockEndToken(Token):
+    id = "<block end>"
+
+
+class FlowSequenceStartToken(Token):
+    id = "["
+
+
+class FlowMappingStartToken(Token):
+    id = "{"
+
+
+class FlowSequenceEndToken(Token):
+    id = "]"
+
+
+class FlowMappingEndToken(Token):
+    id = "}"
+
+
+class KeyToken(Token):
+    id = "?"
+
+
+class ValueToken(Token):
+    id = ":"
+
+
+class BlockEntryToken(Token):
+    id = "-"
+
+
+class FlowEntryToken(Token):
+    id = ","
+
+
+class AliasToken(Token):
+    id = "<alias>"
+
+    def __init__(self, value, start_mark, end_mark):
+        self.value = value
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+
+
+class AnchorToken(Token):
+    id = "<anchor>"
+
+    def __init__(self, value, start_mark, end_mark):
+        self.value = value
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+
+
+class TagToken(Token):
+    id = "<tag>"
+
+    def __init__(self, value, start_mark, end_mark):
+        self.value = value
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+
+
+class ScalarToken(Token):
+    id = "<scalar>"
+
+    def __init__(self, value, plain, start_mark, end_mark, style=None):
+        self.value = value
+        self.plain = plain
+        self.start_mark = start_mark
+        self.end_mark = end_mark
+        self.style = style

outerbounds/command_groups/cli.py

@@ -1,4 +1,4 @@
-import click
+from outerbounds._vendor import click
 from . import local_setup_cli
 from . import workstations_cli
 from . import perimeters_cli

outerbounds/command_groups/local_setup_cli.py

@@ -11,8 +11,7 @@ from importlib.machinery import PathFinder
 from os import path
 from pathlib import Path
 from typing import Any, Callable, Dict, List
-
-import click
+from outerbounds._vendor import click
 import requests
 from requests.exceptions import HTTPError
 
@@ -641,7 +640,6 @@ class ConfigurationWriter:
         if config_type == "inline":
             if "OBP_PERIMETER" in self.decoded_config:
                 self.selected_perimeter = self.decoded_config["OBP_PERIMETER"]
-
             if "OBP_METAFLOW_CONFIG_URL" in self.decoded_config:
                 self.decoded_config = {
                     "OBP_METAFLOW_CONFIG_URL": self.decoded_config[
@@ -702,8 +700,6 @@ class ConfigurationWriter:
         with open(config_path, "w") as fd:
             json.dump(self.existing, fd, indent=4)
 
-        # Every time a config is initialized, we should also reset the corresponding ob_config[_profile].json
-        remote_config = metaflowconfig.init_config(self.out_dir, self.profile)
         if self.selected_perimeter and "OBP_METAFLOW_CONFIG_URL" in self.decoded_config:
             with open(self.ob_config_path, "w") as fd:
                 ob_config_dict = {

outerbounds/command_groups/perimeters_cli.py

@@ -1,22 +1,13 @@
-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 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,
@@ -81,7 +72,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 +116,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 +289,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 +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,
@@ -332,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(
@@ -372,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],
@@ -396,4 +466,44 @@ 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)
+
+    # 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

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