Source code for latch_cli.utils

"""Utility functions for services."""

import hashlib
import os
import stat
import subprocess
import urllib.parse
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import List
from urllib.parse import urljoin

import click
import jwt
from latch_sdk_config.user import user_config

from latch_cli.constants import latch_constants
from latch_cli.tinyrequests import get

# todo(ayush): need a better way to check if "latch" has been appended to urllib
if "latch" not in urllib.parse.uses_netloc:
    urllib.parse.uses_netloc.append("latch")
    urllib.parse.uses_relative.append("latch")


[docs]def urljoins(*args: str, dir: bool = False) -> str: """Construct a URL by appending paths Paths are always joined, with extra `/`s added if missing. Does not allow overriding basenames as opposed to normal `urljoin`. Whether the final path ends in a `/` is still significant and will be preserved in the output >>> urljoin("latch:///directory/", "another_directory") latch:///directory/another_directory >>> # No slash means "another_directory" is treated as a filename >>> urljoin(urljoin("latch:///directory/", "another_directory"), "file") latch:///directory/file >>> # Unintentionally overrode the filename >>> urljoins("latch:///directory/", "another_directory", "file") latch:///directory/another_directory/file >>> # Joined paths as expected Args: args: Paths to join dir: If true, ensure the output ends with a `/` """ res = args[0] for x in args[1:]: if res[-1] != "/": res = f"{res}/" res = urljoin(res, x) if dir and res[-1] != "/": res = f"{res}/" return res
[docs]class AuthenticationError(RuntimeError): ...
[docs]def get_auth_header() -> str: sdk_token = user_config.token execution_token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID") if sdk_token is not None and sdk_token != "": header = f"Latch-SDK-Token {sdk_token}" elif execution_token is not None: header = f"Latch-Execution-Token {execution_token}" else: raise AuthenticationError("Unable to find authentication credentials.") return header
[docs]def retrieve_or_login() -> str: """Returns a valid JWT to access Latch, prompting a login flow if needed. Returns: A JWT """ from latch_cli.services.login import login token = user_config.token if token == "": token = login() return token
[docs]def current_workspace() -> str: ws = user_config.workspace_id if ws == "": ws = account_id_from_token(retrieve_or_login()) user_config.update_workspace(ws, "Personal Workspace") return ws
[docs]def sub_from_jwt(token: str) -> str: """Extract a user sub (UUID) from a JWT minted by auth0. Args: token: JWT Returns: The user sub contained within the JWT. """ payload = jwt.decode(token, options={"verify_signature": False}) try: sub = payload["sub"] except KeyError: raise ValueError( "Provided token lacks a user sub in the data payload" " and is not a valid token." ) return sub
[docs]def account_id_from_token(token: str) -> str: """Exchanges a valid JWT for a Latch account ID. Latch account IDs are needed for any user-specific request, eg. register workflows or copy files to Latch. Args: token: JWT Returns: A Latch account ID (UUID). """ decoded_jwt = jwt.decode(token, options={"verify_signature": False}) try: return decoded_jwt.get("id") except KeyError as e: raise ValueError("Your Latch access token is malformed") from e
def _normalize_remote_path(remote_path: str): if remote_path.startswith("latch://"): remote_path = remote_path[len("latch://") :] if ( not remote_path.startswith("/") and not remote_path.startswith("shared") and not remote_path.startswith("account") ): remote_path = f"/{remote_path}" return remote_path def _si_number_strings(num): for unit in ["", "k", "M", "G", "T", "P", "E", "Z"]: if abs(num) < 1000: # `rstrip` remoes trailing zeroes return f"{num:3.1f}".rstrip("0").rstrip("."), unit num /= 1000 return f"{num:.1f}", "Yi"
[docs]def with_si_suffix(num, suffix="B", styled=False): num, unit = _si_number_strings(num) if styled: import click return click.style(num, fg="bright_green", bold=True) + click.style( f"{unit}{suffix}", fg="green" ) return f"{num}{unit}{suffix}"
[docs]def hash_directory(dir_path: Path) -> str: # todo(maximsmol): store per-file hashes to show which files triggered a version change click.secho("Calculating workflow version based on file content hash", bold=True) click.secho(" Disable with --disable-auto-version/-d", italic=True, dim=True) m = hashlib.new("sha256") m.update(current_workspace().encode("utf-8")) ignore_file = dir_path / ".dockerignore" exclude: List[str] = ["/.latch", ".git"] try: with ignore_file.open("r") as f: click.secho(" Using .dockerignore", italic=True) for l in f: l = l.strip() if l == "": continue if l[0] == "#": continue exclude.append(l) except FileNotFoundError: ... from docker.utils import exclude_paths paths = list(exclude_paths(dir_path, exclude)) paths.sort() for item in paths: p = Path(dir_path / item) p_stat = p.stat() if stat.S_ISDIR(p_stat.st_mode): m.update(p.name.encode("utf-8")) continue m.update(p.name.encode("utf-8")) file_size = p_stat.st_size if not stat.S_ISREG(p_stat.st_mode): click.secho( f"{p.relative_to(dir_path.resolve())} is not a regular file. Ignoring" " contents", fg="yellow", bold=True, ) continue if file_size > latch_constants.file_max_size: click.secho( f"{p.relative_to(dir_path.resolve())} is too large" f" ({with_si_suffix(file_size)}) to checksum. Ignoring contents", fg="yellow", bold=True, ) continue m.update(p.read_bytes()) return m.hexdigest()
[docs]def generate_temporary_ssh_credentials(ssh_key_path: Path) -> str: # check if there is already a valid key at that path, and if so, use that # otherwise, if its not valid, remove it if ssh_key_path.exists(): try: # check if file is valid + print out a fingerprint for the key cmd = ["ssh-keygen", "-l", "-f", ssh_key_path] valid_private_key = subprocess.run(cmd, check=True, capture_output=True) cmd = ["ssh-keygen", "-l", "-f", ssh_key_path.with_suffix(".pub")] valid_public_key = subprocess.run(cmd, check=True, capture_output=True) if valid_private_key.stdout != valid_public_key.stdout: raise # if both files are valid and their fingerprints match, use them instead of generating a new pair print(f"Found existing key pair at {ssh_key_path}.") except: print(f"Found malformed key-pair at {ssh_key_path}. Overwriting.") ssh_key_path.unlink(missing_ok=True) ssh_key_path.with_suffix(".pub").unlink(missing_ok=True) # generate private key if not ssh_key_path.exists(): ssh_key_path.parent.mkdir(parents=True, exist_ok=True) cmd = ["ssh-keygen", "-f", ssh_key_path, "-N", "", "-q"] try: subprocess.run(cmd, check=True) except subprocess.CalledProcessError as e: raise ValueError( "There was a problem creating temporary SSH credentials. Please ensure" " that `ssh-keygen` is installed and available in your PATH." ) from e os.chmod(ssh_key_path, 0o700) # make key available to ssh-agent daemon cmd = ["ssh-add", ssh_key_path] try: subprocess.run(cmd, check=True, stderr=subprocess.DEVNULL) except subprocess.CalledProcessError as e: raise ValueError( "There was an issue adding temporary SSH credentials to your SSH Agent." " Please ensure that your SSH Agent is running" ) from e # decode private key into public key cmd = ["ssh-keygen", "-y", "-f", ssh_key_path] try: out = subprocess.run(cmd, check=True, capture_output=True) except subprocess.CalledProcessError as e: raise ValueError( "There was a problem decoding your temporary credentials. Please ensure" " that `ssh-keygen` is installed and available in your PATH." ) from e public_key = out.stdout.decode("utf-8").strip("\n") return public_key
[docs]def get_local_package_version() -> str: try: from importlib import metadata except ImportError: import importlib_metadata as metadata return metadata.version(latch_constants.pkg_name)
[docs]def get_latest_package_version_request() -> str: cache_location = user_config.root / "cached_version" resp = get(f"https://pypi.org/pypi/{latch_constants.pkg_name}/json") version = resp.json()["info"]["version"] with open(cache_location, "w") as f: f.write(f"{version} {datetime.now().isoformat()}") return version
[docs]def get_latest_package_version() -> str: version = None cache_location = user_config.root / "cached_version" try: with open(cache_location, "r") as f: version, timestamp = f.read().split(" ") if datetime.now() > datetime.fromisoformat(timestamp) + timedelta(days=1): version = get_latest_package_version_request() except FileNotFoundError: version = get_latest_package_version_request() return version
[docs]class TemporarySSHCredentials: def __init__(self, ssh_key_path: Path): self._ssh_key_path = ssh_key_path self._public_key = None
[docs] def generate(self): if self._public_key is not None: return self._public_key = generate_temporary_ssh_credentials(self._ssh_key_path)
[docs] def cleanup(self): if ( self._ssh_key_path.exists() and self._ssh_key_path.with_suffix(".pub").exists() ): subprocess.run( ["ssh-add", "-d", self._ssh_key_path], check=True, capture_output=True, ) self._ssh_key_path.unlink(missing_ok=True) self._ssh_key_path.with_suffix(".pub").unlink(missing_ok=True)
@property def public_key(self): self.generate() return self._public_key @property def private_key(self): self.generate() with open(self._ssh_key_path, "r") as f: return f.read() def __enter__(self): self.generate() return self def __exit__(self, type, value, traceback): self.cleanup()
[docs]class WorkflowType(Enum): latchbiosdk = "latchbiosdk" snakemake = "snakemake"
[docs]def identifier_suffix_from_str(x: str) -> str: res = "" for c in x: res += c if f"_{c}".isidentifier() else "_" return res
[docs]def identifier_from_str(x: str) -> str: return (x[0] if x[0].isidentifier() else "_") + identifier_suffix_from_str(x[1:])