frogml 1.1.62__py3-none-any.whl → 1.2.19__py3-none-any.whl

frogml/core/inner/tool/auth/auth_client.py

@@ -0,0 +1,192 @@
+import os
+from typing import Callable, Optional, Union
+from urllib.parse import urljoin
+
+import requests
+from cachetools import LRUCache, TTLCache, cached, cachedmethod
+from frogml.core.exceptions import FrogmlLoginException
+from frogml.core.exceptions.frogml_token_exception import FrogMLTokenException
+from frogml.core.inner.tool.auth.token_maintainer import TokenMaintainer
+from frogml.storage.authentication.models import AuthConfig, BearerAuth, EmptyAuth
+from frogml.storage.authentication.utils import get_credentials
+from frogml.storage.constants import (
+    JF_ACCESS_TOKEN_RELOAD_INTERVAL_SECONDS,
+    JF_SHORT_TOKEN_ROTATION_GENERATION_TIMEOUT_SECONDS,
+    JF_SHORT_TOKEN_ROTATION_GENERATION_TIMEOUT_SECONDS_DEFAULT,
+    JF_SHORT_TOKEN_ROTATION_GRACE_PERIOD_SECONDS,
+    JF_SHORT_TOKEN_ROTATION_GRACE_PERIOD_SECONDS_DEFAULT,
+    JF_SHORT_TOKEN_ROTATION_INTERVAL_SECONDS,
+    JF_SHORT_TOKEN_ROTATION_INTERVAL_SECONDS_DEFAULT,
+    JF_USE_SHORT_TOKEN_ROTATION,
+    JF_USE_SHORT_TOKEN_ROTATION_DEFAULT,
+)
+from requests import Response
+from requests.exceptions import RequestException
+
+
+class FrogMLAuthClient:
+    __MIN_TOKEN_LENGTH: int = 64
+
+    def __init__(self, auth_config: Optional[AuthConfig] = None):
+        self.__auth_config: Optional[AuthConfig] = auth_config
+        self.__token: Optional[str] = None
+        self.__tenant_id: Optional[str] = None
+
+        # Initialize instance-level caches
+        self.__token_maintainer_cache = LRUCache(maxsize=1)
+        self.__tenant_id_cache = LRUCache(maxsize=1)
+        self.__jpd_id_cache = LRUCache(maxsize=10)
+
+        self.__base_token_retriever: Callable[[], str] = (
+            self.__get_base_token_retriever()
+        )
+        self.__use_token_maintainer: bool = self.__should_use_token_maintainer()
+
+        # best-effort to prime the token maintainer (if used)
+        if self.__use_token_maintainer:
+            self.get_token()
+
+    def __should_use_token_maintainer(self) -> bool:
+        return (
+            str(
+                os.getenv(
+                    JF_USE_SHORT_TOKEN_ROTATION, JF_USE_SHORT_TOKEN_ROTATION_DEFAULT
+                )
+            ).upper()
+            == "TRUE"
+        )
+
+    def get_token(self) -> str:
+        base_token: str = self.__base_token_retriever()
+
+        if not self.__use_token_maintainer:
+            return base_token
+
+        return self._get_token_maintainer(base_token).get_token()
+
+    @cachedmethod(
+        cache=lambda self: self.__token_maintainer_cache,
+    )
+    def _get_token_maintainer(self, base_token: str) -> TokenMaintainer:
+        token_rotation_interval_seconds: int = int(
+            os.getenv(
+                JF_SHORT_TOKEN_ROTATION_INTERVAL_SECONDS,
+                JF_SHORT_TOKEN_ROTATION_INTERVAL_SECONDS_DEFAULT,
+            )
+        )
+        token_rotation_grace_period_seconds: int = int(
+            os.getenv(
+                JF_SHORT_TOKEN_ROTATION_GRACE_PERIOD_SECONDS,
+                JF_SHORT_TOKEN_ROTATION_GRACE_PERIOD_SECONDS_DEFAULT,
+            )
+        )
+        token_rotation_generation_timeout_seconds: int = int(
+            os.getenv(
+                JF_SHORT_TOKEN_ROTATION_GENERATION_TIMEOUT_SECONDS,
+                JF_SHORT_TOKEN_ROTATION_GENERATION_TIMEOUT_SECONDS_DEFAULT,
+            )
+        )
+
+        return TokenMaintainer(
+            base_token=base_token,
+            base_url=self.get_base_url(),
+            token_rotation_interval_seconds=token_rotation_interval_seconds,
+            token_rotation_grace_period_seconds=token_rotation_grace_period_seconds,
+            token_rotation_generation_timeout_seconds=token_rotation_generation_timeout_seconds,
+        )
+
+    def get_base_token(self) -> str:
+        auth: Union[EmptyAuth, BearerAuth] = self.get_auth()
+
+        if isinstance(auth, BearerAuth):
+            self.validate_token(auth.token)
+            self.__token = auth.token
+
+        else:
+            raise FrogMLTokenException(
+                message="Token not found in the authentication configurations."
+            )
+
+        return self.__token
+
+    @cachedmethod(lambda self: self.__tenant_id_cache)
+    def get_tenant_id(self) -> str:
+        base_url: str = self.get_base_url()
+        url: str = urljoin(base_url, "/ui/api/v1/system/auth/screen/footer")
+
+        try:
+            response: Response = requests.get(url, timeout=15, auth=self.get_auth())
+            response.raise_for_status()  # Raises an HTTPError for bad responses
+            response_data: dict = response.json()
+
+            if "serverId" in response_data:
+                self.__tenant_id = response_data["serverId"]
+            else:
+                self.__tenant_id = self.__get_jpd_id(base_url)
+
+            return self.__tenant_id
+
+        except (RequestException, ValueError) as exc:
+            raise FrogmlLoginException(
+                "Failed to authenticate with JFrog. Please check your credentials"
+            ) from exc
+
+    def get_base_url(self) -> str:
+        artifactory_url, _ = get_credentials(self.__auth_config)
+        return self.__remove_artifactory_path_from_url(artifactory_url)
+
+    def get_auth(self) -> Union[EmptyAuth, BearerAuth]:
+        return get_credentials(self.__auth_config)[1]
+
+    def validate_token(self, token: Optional[str]):
+        if token is None or len(token) <= self.__MIN_TOKEN_LENGTH or token.isspace():
+            raise FrogmlLoginException(
+                "Authentication with JFrog failed: Only JWT Access Tokens are supported. "
+                "Please ensure you are using a valid JWT Access Token."
+            )
+
+    @cachedmethod(cache=lambda self: self.__jpd_id_cache)
+    def __get_jpd_id(self, base_url: str) -> str:
+        url: str = urljoin(base_url, "/jfconnect/api/v1/system/jpd_id")
+        response: Response = requests.get(url=url, timeout=15, auth=self.get_auth())
+
+        if response.status_code == 200:
+            return response.text
+
+        if response.status_code == 401:
+            raise FrogmlLoginException(
+                "Failed to authenticate with JFrog. Please check your credentials"
+            )
+        else:
+            raise FrogmlLoginException(
+                "Failed to authenticate with JFrog. Please check your artifactory configuration"
+            )
+
+    def __get_base_token_retriever(self) -> Callable[[], str]:
+        if JF_ACCESS_TOKEN_RELOAD_INTERVAL_SECONDS not in os.environ:
+            # no reload interval configured, preserve the original behavior of reading each time
+            return self.get_base_token
+
+        # base token reload interval is configured, wrap the basic retrieval with a cache
+        try:
+            ttl: int = int(os.getenv(JF_ACCESS_TOKEN_RELOAD_INTERVAL_SECONDS))
+        except ValueError:
+            raise FrogmlLoginException(
+                f"Value of {JF_ACCESS_TOKEN_RELOAD_INTERVAL_SECONDS} should be an integer"
+            )
+
+        cache: TTLCache = TTLCache(ttl=ttl, maxsize=1)
+
+        # Use a constant key since this is a nested function with no parameters
+        @cached(cache=cache)
+        def get_token() -> str:
+            return self.get_base_token()
+
+        return get_token
+
+    @staticmethod
+    def __remove_artifactory_path_from_url(artifactory_url: str) -> str:
+        # Remove '/artifactory' from the URL
+        base_url: str = artifactory_url.replace("/artifactory", "", 1)
+        # Remove trailing slash if exists
+        return base_url.rstrip("/")

frogml/core/inner/tool/auth/token_maintainer.py

@@ -0,0 +1,119 @@
+import threading
+import time
+from typing import Any, Optional
+from urllib.parse import urljoin
+
+import requests
+from frogml.core.exceptions import FrogmlLoginException
+from frogml.storage.authentication.models import BearerAuth
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from requests import Response
+
+
+class TokenMaintainer(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    token_rotation_interval_seconds: int = Field(ge=0)
+    token_rotation_grace_period_seconds: int = Field(ge=0)
+    token_rotation_generation_timeout_seconds: int = Field(ge=0)
+    base_token: str
+    base_url: str
+    __token_scope: str = ""
+    __curr_token: Optional[str] = None
+    __curr_token_cache_expiration: float
+    __lock: threading.RLock = PrivateAttr(default_factory=threading.RLock)
+    __refreshing: bool = False
+
+    def model_post_init(self, context: Any, /) -> None:
+        self.__curr_token_cache_expiration = (
+            -2
+            * (
+                self.token_rotation_interval_seconds
+                + self.token_rotation_grace_period_seconds
+            )
+            - 1.0
+        )
+
+        self.__extract_token_scope()
+
+    def get_token(self) -> str:
+        curr_time = time.monotonic()
+
+        if curr_time > self.__curr_token_cache_expiration:
+            # need to refresh (either because it's about to expire it was never loaded)
+            # note that this does not necessarily mean the current value can't be used
+            self.trigger_background_load()
+
+        if (
+            curr_time
+            > self.__curr_token_cache_expiration
+            + self.token_rotation_grace_period_seconds
+        ):
+            # the current value is either no longer valid (passed its grace) or was never loaded, return the base token
+            return self.base_token
+
+        return self.__curr_token
+
+    def __refresh(self):
+        next_curr_value_expiration: float = (
+            time.monotonic() + self.token_rotation_interval_seconds
+        )
+        try:
+            token = self.__generate_token()
+            with self.__lock:
+                self.__curr_token = token
+                self.__curr_token_cache_expiration = next_curr_value_expiration
+        finally:
+            with self.__lock:
+                self.__refreshing = False
+
+    def __generate_token(self) -> str:
+        try:
+            url: str = urljoin(self.base_url, "/access/api/v1/tokens")
+            token_life_duration: int = int(
+                self.token_rotation_interval_seconds
+                + self.token_rotation_grace_period_seconds
+            )
+            resp: Response = requests.post(
+                url,
+                data={"expires_in": token_life_duration, "scope": self.__token_scope},
+                auth=BearerAuth(self.base_token),
+                timeout=self.token_rotation_generation_timeout_seconds,
+            )
+            resp.raise_for_status()
+            return resp.json()["access_token"]
+        except Exception as e:
+            raise FrogmlLoginException("Failed to generate token: {}".format(e)) from e
+
+    def trigger_background_load(self):
+        with self.__lock:
+            if self.__refreshing:
+                # another refresh is ongoing, do not disturb it
+                return
+            # first start a thread for refreshing the token,
+            # then mark self as refreshing (to prevent a phantom refresh)
+            threading.Thread(target=self.__refresh, daemon=True).start()
+            self.__refreshing = True
+
+    def __extract_token_scope(self):
+        # parsing the base token to get its scope (which is used for generating new ones)
+        # note that here it does not matter whether the token is valid or not, we care
+        # only about the scope
+        try:
+            import jwt
+
+            parsed_token = jwt.decode(
+                self.base_token,
+                options={
+                    "verify_signature": False
+                },  # NOSONAR(S5659) we only parse the token, no authentication happens here
+            )
+        except Exception as e:
+            raise FrogmlLoginException(
+                "Failed to decode JWT token: {}".format(e)
+            ) from e
+
+        if "scp" not in parsed_token:
+            raise FrogmlLoginException("Invalid token - missing scp claim")
+
+        self.__token_scope = parsed_token["scp"]

frogml/core/inner/tool/grpc/grpc_auth.py

@@ -1,42 +1,41 @@
-from typing import TYPE_CHECKING, Callable, Optional, Tuple
-
 import grpc
+from typing_extensions import Self
 
 from frogml.core.inner.const import FrogMLConstants
-
-if TYPE_CHECKING:
-    from frogml.core.inner.tool.auth import FrogMLAuthClient
-
-_SIGNATURE_HEADER_KEY = "authorization"
+from frogml.core.inner.tool.auth import FrogMLAuthClient
 
 
 class FrogMLAuthMetadataPlugin(grpc.AuthMetadataPlugin):
-    def __init__(self):
-        self._auth_client: Optional["FrogMLAuthClient"] = None
+    def __init__(self: Self, should_pass_jf_tenant_id: bool = True):
+        self.__auth_client = FrogMLAuthClient()
+        self.__should_pass_jf_tenant_id: bool = should_pass_jf_tenant_id
 
     def __call__(
         self,
         context: grpc.AuthMetadataContext,
-        callback: Callable[[Tuple[Tuple[str, str], Tuple[str, str]], None], None],
+        callback: grpc.AuthMetadataPluginCallback,
     ):
         """Implements authentication by passing metadata to a callback.
 
         Args:
             context: An AuthMetadataContext providing information on the RPC that
                 the plugin is being called to authenticate.
-            callback: A callback that accepts a tuple of metadata key/value pairs and a None
+            callback: A callback that accepts a sequence of metadata key/value pairs and a None
                 parameter.
         """
-        # Get token from FrogML client
-        if not self._auth_client:
-            from frogml.core.inner.tool.auth import FrogMLAuthClient
-
-            self._auth_client = FrogMLAuthClient()
-
-        token = self._auth_client.get_token()
-        jfrog_tenant_id = self._auth_client.get_tenant_id()
-        metadata = (
-            (_SIGNATURE_HEADER_KEY, f"Bearer {token}"),
-            (FrogMLConstants.JFROG_TENANT_HEADER_KEY.lower(), jfrog_tenant_id),
-        )
+        token: str = self.__auth_client.get_token()
+        metadata: list[tuple[str, str]] = [
+            (FrogMLConstants.SIGNATURE_HEADER_KEY, f"Bearer {token}")
+        ]
+
+        if self.__should_pass_jf_tenant_id:
+            jfrog_tenant_id: str = self.__auth_client.get_tenant_id()
+            metadata.append(
+                (FrogMLConstants.JFROG_TENANT_HEADER_KEY.lower(), jfrog_tenant_id)
+            )
+
         callback(metadata, None)
+
+    @property
+    def should_pass_jf_tenant_id(self: Self) -> bool:
+        return self.__should_pass_jf_tenant_id

frogml/core/inner/tool/grpc/grpc_tools.py

@@ -1,15 +1,18 @@
 import logging
+import re
 import time
 from abc import ABC, abstractmethod
 from random import randint
 from typing import Callable, Optional, Tuple
+from urllib.parse import urlparse, ParseResult
 
 import grpc
 
-from frogml.core.exceptions import FrogmlException
+from frogml.core.exceptions import FrogmlGrpcAddressException, FrogmlException
 from frogml.core.inner.tool.grpc.grpc_auth import FrogMLAuthMetadataPlugin
 
 logger = logging.getLogger()
+HOSTNAME_REGEX: str = r"^(?!-)(?:[A-Za-z0-9-]{1,63}\.)*[A-Za-z0-9-]{1,63}(?<!-)$"
 
 
 def create_grpc_channel(
@@ -23,6 +26,7 @@ def create_grpc_channel(
     max_attempts=4,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
+    should_pass_jf_tenant_id: bool = True,
 ) -> grpc.Channel:
     """
     Create a gRPC channel
@@ -36,16 +40,17 @@ def create_grpc_channel(
         options:  An optional list of key-value pairs (channel_arguments in gRPC Core runtime) to configure the channel.
         backoff_options: dictionary - init_backoff_ms: default 50, max_backoff_ms: default 500, multiplier: default 2
         max_attempts: max number of retry attempts
-        attempt: current retry attempts
         status_for_retry: grpc statuses to retry upon
+        attempt: current retry attempts
+        should_pass_jf_tenant_id: Whether to pass JFrog tenant ID in metadata (if auth is enabled)
     Returns: Returns a grpc.Channel
     """
     if backoff_options is None:
         backoff_options = {}
 
     if not url:
-        raise FrogmlException(
-            "Unable to create gRPC channel. URL has not been defined."
+        raise FrogmlGrpcAddressException(
+            "Unable to create gRPC channel. URL has not been defined.", url
         )
 
     if enable_ssl or url.endswith(":443"):
@@ -53,7 +58,9 @@ def create_grpc_channel(
 
         if enable_auth:
             if auth_metadata_plugin is None:
-                auth_metadata_plugin = FrogMLAuthMetadataPlugin()
+                auth_metadata_plugin = FrogMLAuthMetadataPlugin(
+                    should_pass_jf_tenant_id
+                )
 
             credentials = grpc.composite_channel_credentials(
                 credentials, grpc.metadata_call_credentials(auth_metadata_plugin)
@@ -99,7 +106,7 @@ def create_grpc_channel(
         return intercept_channel
     except grpc.FutureTimeoutError as e:
         raise FrogmlException(
-            f"Connection timed out while attempting to connect to {url}, with: {repr(e)}"
+            f"Connection timed out while attempting to connect to {url}, with: {repr(e)}",
         ) from e
 
 
@@ -114,6 +121,7 @@ def create_grpc_channel_or_none(
     max_attempts=2,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
+    should_pass_jf_tenant_id: bool = True,
 ) -> Callable[[Optional[str], Optional[bool]], Optional[grpc.Channel]]:
     final_backoff_options = backoff_options if backoff_options else {}
 
@@ -132,6 +140,7 @@ def create_grpc_channel_or_none(
                 max_attempts,
                 status_for_retry,
                 attempt,
+                should_pass_jf_tenant_id,
             )
         except Exception as e:
             logger.debug(f"create_grpc_channel error: {repr(e)}")
@@ -140,6 +149,117 @@ def create_grpc_channel_or_none(
     return deferred_channel
 
 
+def validate_grpc_address(
+    grpc_address: str,
+    is_port_specification_allowed: bool = False,
+    is_url_scheme_allowed: bool = False,
+):
+    """
+    Validate gRPC address format
+    Args:
+        grpc_address (str): gRPC address to validate
+        is_port_specification_allowed (bool): Whether to allow port specification in the address
+        is_url_scheme_allowed (bool): Whether to allow URL scheme in the address
+    Raises:
+        FrogmlGrpcAddressException: If the gRPC address is invalid
+    """
+    parsed_grpc_address: ParseResult = parse_address(grpc_address)
+    hostname: str = get_hostname_from_address(parsed_grpc_address)
+    validate_paths_are_not_included_in_address(parsed_grpc_address)
+
+    if not is_url_scheme_allowed:
+        __validate_url_scheme_not_included_in_address(parsed_grpc_address)
+
+    if not is_port_specification_allowed:
+        __validate_port_not_included_in_address(parsed_grpc_address)
+
+    if not is_valid_hostname(hostname):
+        raise FrogmlGrpcAddressException(
+            "gRPC address must be a simple hostname or fully qualified domain name.",
+            parsed_grpc_address,
+        )
+
+
+def validate_paths_are_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    has_invalid_path: bool = (
+        parsed_grpc_address.path not in {"", "/"}
+        or parsed_grpc_address.query
+        or parsed_grpc_address.fragment
+    )
+
+    if has_invalid_path:
+        raise FrogmlGrpcAddressException(
+            "gRPC address must not contain paths, queries, or fragments.",
+            parsed_grpc_address,
+        )
+
+
+def get_hostname_from_address(parsed_grpc_address: ParseResult) -> str:
+    hostname: Optional[str] = parsed_grpc_address.hostname
+    if not hostname:
+        raise FrogmlGrpcAddressException(
+            "gRPC address must contain a valid hostname.", parsed_grpc_address
+        )
+
+    return hostname
+
+
+def __validate_url_scheme_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    if parsed_grpc_address.scheme:
+        raise FrogmlGrpcAddressException(
+            "URL scheme is not allowed in the gRPC address.", parsed_grpc_address
+        )
+
+
+def __validate_port_not_included_in_address(parsed_grpc_address: ParseResult):
+    try:
+        port: Optional[int] = parsed_grpc_address.port
+    except ValueError as exc:
+        raise FrogmlGrpcAddressException(
+            "Invalid port specification in the gRPC address.", parsed_grpc_address
+        ) from exc
+
+    if port:
+        raise FrogmlGrpcAddressException(
+            "Port specification is not allowed in the gRPC address.",
+            parsed_grpc_address,
+        )
+
+
+def parse_address(grpc_address: str) -> ParseResult:
+    if not grpc_address or not grpc_address.strip():
+        raise FrogmlGrpcAddressException(
+            "gRPC address must not be empty or whitespace.", grpc_address
+        )
+
+    trimmed_address: str = grpc_address.strip()
+    parsed_address: ParseResult = urlparse(
+        trimmed_address if "://" in trimmed_address else f"//{trimmed_address}"
+    )
+
+    return parsed_address
+
+
+def is_valid_hostname(hostname: str) -> bool:
+    """
+    Validate that the supplied hostname conforms to RFC-style label rules:
+    anchored pattern enforces full-string validation, negative lookahead/lookbehind block
+    leading or trailing hyphens per label, and each dot-separated label must be 1-63
+    alphanumeric/hyphen characters.
+
+    Args:
+        hostname (str): The hostname to validate.
+    Returns:
+        bool: True if the hostname is valid, False otherwise.
+    """
+    hostname_pattern: re.Pattern = re.compile(HOSTNAME_REGEX)
+    return bool(hostname_pattern.fullmatch(hostname))
+
+
 class SleepingPolicy(ABC):
     @abstractmethod
     def sleep(self, try_i: int):