skyflow-flowvault-python 1.0.0__py3-none-any.whl

common/generated/rest/errors/unauthorized_error.py

@@ -0,0 +1,14 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.api_error import ApiError
+
+
+class UnauthorizedError(ApiError):
+    def __init__(
+        self,
+        body: typing.Dict[str, typing.Optional[typing.Any]],
+        headers: typing.Optional[typing.Dict[str, str]] = None,
+    ):
+        super().__init__(status_code=401, headers=headers, body=body)

common/generated/rest/types/__init__.py

@@ -0,0 +1,9 @@
+# This file was auto-generated by Fern from our API Definition.
+
+# isort: skip_file
+
+from .googlerpc_status import GooglerpcStatus
+from .protobuf_any import ProtobufAny
+from .v_1_get_auth_token_response import V1GetAuthTokenResponse
+
+__all__ = ["GooglerpcStatus", "ProtobufAny", "V1GetAuthTokenResponse"]

common/generated/rest/types/googlerpc_status.py

@@ -0,0 +1,22 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from .protobuf_any import ProtobufAny
+
+
+class GooglerpcStatus(UniversalBaseModel):
+    code: typing.Optional[int] = None
+    message: typing.Optional[str] = None
+    details: typing.Optional[typing.List[ProtobufAny]] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

common/generated/rest/types/protobuf_any.py

@@ -0,0 +1,21 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class ProtobufAny(UniversalBaseModel):
+    type: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="@type")] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

common/generated/rest/types/v_1_get_auth_token_response.py

@@ -0,0 +1,33 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+import pydantic
+import typing_extensions
+from ..core.pydantic_utilities import IS_PYDANTIC_V2, UniversalBaseModel
+from ..core.serialization import FieldMetadata
+
+
+class V1GetAuthTokenResponse(UniversalBaseModel):
+    access_token: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="accessToken")] = (
+        pydantic.Field(default=None)
+    )
+    """
+    AccessToken.
+    """
+
+    token_type: typing_extensions.Annotated[typing.Optional[str], FieldMetadata(alias="tokenType")] = pydantic.Field(
+        default=None
+    )
+    """
+    TokenType : Bearer.
+    """
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

common/generated/rest/version.py

@@ -0,0 +1,6 @@
+# NOTE: hand-patched, not Fern-generated content. Fern originally emitted a runtime
+# metadata.version("skyflow.generated.rest") lookup here, but this code is bundled into the
+# skyflow (v2) and v3 wheels via a build_py hook rather than published under that distribution
+# name, so the lookup always raised PackageNotFoundError on import. Hardcoded until the Fern
+# generator config (skyflow-fern-config) is updated to stop emitting a runtime lookup here.
+__version__ = "0.0.9"

common/service_account/__init__.py

@@ -0,0 +1 @@
+from ._utils import generate_bearer_token, generate_bearer_token_from_creds, is_expired, generate_signed_data_tokens, generate_signed_data_tokens_from_creds

common/service_account/_utils.py

@@ -0,0 +1,248 @@
+import json
+import datetime
+import re
+import time
+import jwt
+from urllib.parse import urlparse
+from common.errors import SkyflowError
+from common.service_account.client.auth_client import AuthClient
+from common.utils.logger import log_info, log_error_log
+from common.utils import get_base_url, format_scope, SkyflowMessages
+from common.utils.constants import JWT, CredentialField, JwtField, OptionField, ResponseField
+from common.generated.rest.errors.unauthorized_error import UnauthorizedError
+from common.utils import is_valid_url
+from common.utils.constants import CTX_KEY_REGEX
+
+
+invalid_input_error_code = SkyflowMessages.ErrorCodes.INVALID_INPUT.value
+
+_CTX_KEY_PATTERN = re.compile(CTX_KEY_REGEX)
+
+_SNAKE_TO_CAMEL_CRED_MAP = {
+    'private_key': CredentialField.PRIVATE_KEY,
+    'client_id':   CredentialField.CLIENT_ID,
+    'key_id':      CredentialField.KEY_ID,
+    'token_uri':   CredentialField.TOKEN_URI,
+    'client_name': CredentialField.CLIENT_NAME,
+}
+
+
+def _normalize_credentials(credentials):
+    return {_SNAKE_TO_CAMEL_CRED_MAP.get(k, k): v for k, v in credentials.items()}
+
+
+def _validate_and_resolve_ctx(ctx, messages=None):
+    """Validate ctx value and return resolved value for JWT claims.
+    Returns None if ctx should be omitted, the value if valid, or raises SkyflowError if invalid.
+    """
+    messages = messages or SkyflowMessages
+    if ctx is None:
+        return None
+    if isinstance(ctx, str):
+        if ctx.strip() == '':
+            return None
+        return ctx
+    if isinstance(ctx, dict):
+        if len(ctx) == 0:
+            return None
+        for key in ctx:
+            if not isinstance(key, str) or not _CTX_KEY_PATTERN.match(key):
+                raise SkyflowError(
+                    messages.Error.INVALID_CTX_MAP_KEY.value.format(key),
+                    invalid_input_error_code
+                )
+        return ctx
+    if isinstance(ctx, (bool, int, float)):
+        return ctx
+    raise SkyflowError(
+        messages.Error.INVALID_CTX_TYPE.value,
+        invalid_input_error_code
+    )
+
+def is_expired(token, logger = None):
+    if token is None:
+        return True
+    if len(token) == 0:
+        log_error_log(SkyflowMessages.ErrorLogs.INVALID_BEARER_TOKEN.value)
+        return True
+
+    try:
+        decoded = jwt.decode(
+            token, options={OptionField.VERIFY_SIGNATURE: False, OptionField.VERIFY_AUD: False})
+        if time.time() >= decoded[JwtField.EXP]:
+            log_info(SkyflowMessages.Info.BEARER_TOKEN_EXPIRED.value, logger)
+            log_error_log(SkyflowMessages.ErrorLogs.INVALID_BEARER_TOKEN.value)
+            return True
+        return False
+    except jwt.ExpiredSignatureError:
+        return True
+    except Exception:
+        log_error_log(SkyflowMessages.Error.JWT_DECODE_ERROR.value, logger)
+        return True
+
+def generate_bearer_token(credentials_file_path, options = None, logger = None):
+    log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_TRIGGERED.value, logger)
+    try:
+        with open(credentials_file_path, 'r') as credentials_file:
+            try:
+                credentials = json.load(credentials_file)
+            except Exception:
+                log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value, logger=logger)
+                raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path), invalid_input_error_code)
+    except SkyflowError:
+        raise
+    except Exception:
+        raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
+    result = get_service_account_token(credentials, options, logger)
+    return result
+
+def generate_bearer_token_from_creds(credentials, options = None, logger = None):
+    log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_TRIGGERED.value, logger)
+    credentials = credentials.strip()
+    try:
+        json_credentials = json.loads(credentials.replace('\n', '\\n'))
+    except Exception:
+        raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS_STRING.value, invalid_input_error_code)
+    result = get_service_account_token(json_credentials, options, logger)
+    return result
+
+def get_service_account_token(credentials, options, logger):
+    credentials = _normalize_credentials(credentials)
+    try:
+        private_key = credentials[CredentialField.PRIVATE_KEY]
+    except KeyError:
+        log_error_log(SkyflowMessages.ErrorLogs.PRIVATE_KEY_IS_REQUIRED.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.MISSING_PRIVATE_KEY.value, invalid_input_error_code)
+    try:
+        client_id = credentials[CredentialField.CLIENT_ID]
+    except KeyError:
+        log_error_log(SkyflowMessages.ErrorLogs.CLIENT_ID_IS_REQUIRED.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.MISSING_CLIENT_ID.value, invalid_input_error_code)
+    try:
+        key_id = credentials[CredentialField.KEY_ID]
+    except KeyError:
+        log_error_log(SkyflowMessages.ErrorLogs.KEY_ID_IS_REQUIRED.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.MISSING_KEY_ID.value, invalid_input_error_code)
+    try:
+        token_uri = credentials[CredentialField.TOKEN_URI]
+    except KeyError:
+        log_error_log(SkyflowMessages.ErrorLogs.TOKEN_URI_IS_REQUIRED.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.MISSING_TOKEN_URI.value, invalid_input_error_code)
+
+    if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+        log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+    if options and CredentialField.TOKEN_URI_OPTION in options:
+        token_uri = options[CredentialField.TOKEN_URI_OPTION]
+        if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+            log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value, logger=logger)
+            raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+    signed_token = get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger)
+    base_url = get_base_url(token_uri)
+    auth_client = AuthClient(base_url)
+    auth_api = auth_client.get_auth_api()
+
+    formatted_scope = None
+    if options and OptionField.ROLE_IDS in options:
+        formatted_scope = format_scope(options.get(OptionField.ROLE_IDS))
+
+    try:
+        response = auth_api.authentication_service_get_auth_token(assertion = signed_token,
+                                    grant_type=JWT.GRANT_TYPE_JWT_BEARER,
+                                    scope=formatted_scope)
+        log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_SUCCESS.value, logger)
+    except UnauthorizedError:
+        log_error_log(SkyflowMessages.ErrorLogs.UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.UNAUTHORIZED_ERROR_IN_GETTING_BEARER_TOKEN.value, invalid_input_error_code)
+    except Exception:
+        log_error_log(SkyflowMessages.ErrorLogs.FAILED_TO_GET_BEARER_TOKEN.value, logger=logger)
+        raise SkyflowError(SkyflowMessages.Error.FAILED_TO_GET_BEARER_TOKEN.value, invalid_input_error_code)
+    return response.access_token, response.token_type
+
+def get_signed_jwt(options, client_id, key_id, token_uri, private_key, logger):
+    payload = {
+        JwtField.ISS: client_id,
+        JwtField.KEY: key_id,
+        JwtField.AUD: token_uri,
+        JwtField.SUB: client_id,
+        JwtField.EXP: datetime.datetime.utcnow() + datetime.timedelta(minutes=60)
+    }
+    if options and OptionField.CTX in options:
+        resolved_ctx = _validate_and_resolve_ctx(options.get(OptionField.CTX))
+        if resolved_ctx is not None:
+            payload[JwtField.CTX] = resolved_ctx
+    try:
+        return jwt.encode(payload=payload, key=private_key, algorithm=JWT.ALGORITHM_RS256)
+    except Exception:
+        raise SkyflowError(SkyflowMessages.Error.JWT_INVALID_FORMAT.value, invalid_input_error_code)
+
+
+
+def get_signed_tokens(credentials_obj, options):
+    options = options if options is not None else {}
+    credentials_obj = _normalize_credentials(credentials_obj)
+    expiry_time = int(time.time()) + options.get(OptionField.TIME_TO_LIVE, 60)
+    prefix = JWT.SIGNED_TOKEN_PREFIX
+
+    token_uri = credentials_obj.get(CredentialField.TOKEN_URI)
+    if not isinstance(token_uri, str) or not is_valid_url(token_uri):
+        log_error_log(SkyflowMessages.ErrorLogs.INVALID_TOKEN_URI.value)
+        raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_URI.value, invalid_input_error_code)
+
+    resolved_ctx = None
+    if OptionField.CTX in options:
+        resolved_ctx = _validate_and_resolve_ctx(options[OptionField.CTX])
+
+    results = []
+    if options and options.get(OptionField.DATA_TOKENS):
+        for token in options[OptionField.DATA_TOKENS]:
+            claims = {
+                JwtField.ISS: JWT.ISSUER_SDK,
+                JwtField.KEY: credentials_obj.get(CredentialField.KEY_ID),
+                JwtField.EXP: expiry_time,
+                JwtField.SUB: credentials_obj.get(CredentialField.CLIENT_ID),
+                JwtField.TOK: token,
+                JwtField.IAT: int(time.time()),
+            }
+            if resolved_ctx is not None:
+                claims[JwtField.CTX] = resolved_ctx
+            private_key = credentials_obj.get(CredentialField.PRIVATE_KEY)
+            try:
+                signed_jwt = jwt.encode(claims, private_key, algorithm=JWT.ALGORITHM_RS256)
+            except Exception:
+                raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
+            results.append(get_signed_data_token_response_object(prefix + signed_jwt, token))
+    log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKEN_SUCCESS.value)
+    return results
+
+
+def generate_signed_data_tokens(credentials_file_path, options):
+    log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKENS_TRIGGERED.value)
+    try:
+        with open(credentials_file_path, 'r') as credentials_file:
+            try:
+                credentials = json.load(credentials_file)
+            except Exception:
+                raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path),
+                                   invalid_input_error_code)
+    except SkyflowError:
+        raise
+    except Exception:
+        raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
+    return get_signed_tokens(credentials, options)
+
+def generate_signed_data_tokens_from_creds(credentials, options):
+    log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKENS_TRIGGERED.value)
+    credentials = credentials.strip()
+    try:
+        json_credentials = json.loads(credentials.replace('\n', '\\n'))
+    except Exception:
+        log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value)
+        raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS_STRING.value, invalid_input_error_code)
+    return get_signed_tokens(json_credentials, options)
+
+
+def get_signed_data_token_response_object(signed_token, actual_token):
+    return actual_token, signed_token

common/service_account/client/auth_client.py

@@ -0,0 +1,13 @@
+from common.generated.rest.client import SkyflowAuth
+from common.utils.constants import OPTIONAL_TOKEN
+
+class AuthClient:
+    def __init__(self, url):
+        self.__url = url
+        self.__api_client = self.initialize_api_client()
+
+    def initialize_api_client(self):
+        return SkyflowAuth(base_url=self.__url, token=OPTIONAL_TOKEN)
+
+    def get_auth_api(self):
+        return self.__api_client.authentication

common/utils/__init__.py

@@ -0,0 +1,4 @@
+from .enums import LogLevel, Env, TokenType
+from ._skyflow_messages import SkyflowMessages
+from ._helpers import get_base_url, format_scope, is_valid_url
+from ._utils import get_credentials, get_vault_url, validate_api_key

common/utils/_helpers.py

@@ -0,0 +1,18 @@
+from urllib.parse import urlparse
+
+def get_base_url(url):
+    parsed_url = urlparse(url)
+    base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+    return base_url
+
+def format_scope(scopes):
+    if not scopes:
+        return None
+    return " ".join([f"role:{scope}" for scope in scopes])
+
+def is_valid_url(url):
+    try:
+        result = urlparse(url)
+        return all([result.scheme == "https", result.netloc])
+    except Exception:
+        return False