vantage-sdkpy 0.1.1__py3-none-any.whl

vantage_sdk/__init__.py

@@ -0,0 +1,44 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Standalone Vantage SDK public exports."""
+
+# Import base classes
+from .base import BaseCRUDSDK, BaseGraphQLResourceSDK
+from .cloud.cloud_account_crud import cloud_account_sdk
+
+from .cluster.crud import cluster_sdk
+from .job import JobScript, JobScriptFile, JobSubmission, JobTemplate
+from .job import job_script_sdk, job_submission_sdk, job_template_sdk
+from .schemas import CliContext, SDKContext
+from .support_ticket.crud import support_ticket_sdk
+from .team.crud import team_sdk
+from .vantage_rest_api_client import VantageRestApiClient, create_vantage_rest_client
+
+__all__ = [
+    "BaseCRUDSDK",
+    "BaseGraphQLResourceSDK",
+    "SDKContext",
+    "CliContext",
+    "cluster_sdk",
+    "cloud_account_sdk",
+    "team_sdk",
+    "support_ticket_sdk",
+    "job_script_sdk",
+    "job_template_sdk",
+    "job_submission_sdk",
+    "JobScript",
+    "JobScriptFile",
+    "JobSubmission",
+    "JobTemplate",
+    "VantageRestApiClient",
+    "create_vantage_rest_client",
+]

vantage_sdk/admin/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Admin SDK package for administrative operations."""
+
+from .management import get_extra_attributes
+
+__all__ = ["get_extra_attributes"]

vantage_sdk/admin/management/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Management SDK module for administrative operations."""
+
+from .organizations import get_extra_attributes
+
+__all__ = ["get_extra_attributes"]

vantage_sdk/admin/management/organizations.py

@@ -0,0 +1,60 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Organizations management SDK functions.
+
+This module provides SDK functions for organization-level administrative operations
+such as managing extra attributes and other organization configurations.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+import typer
+
+logger = logging.getLogger(__name__)
+
+
+async def get_extra_attributes(ctx: typer.Context) -> Optional[Dict[Any, Any]]:
+    """Get organization extra attributes.
+
+    This function retrieves the list of extra attributes configured for the organization.
+    Extra attributes are custom fields that can be added to various entities.
+
+    The REST client is automatically initialized and attached to ctx.obj.rest_client
+    by the @attach_vantage_rest_client decorator.
+
+    Args:
+        ctx: Typer context with rest_client, settings, and persona already attached
+
+    Returns:
+        Dict of extra attribute data containing the organization's extra attributes,
+        or None if the request fails
+
+    Raises:
+        httpx.HTTPStatusError: If the API request fails
+        Exception: For other request failures
+
+    Example:
+        >>> import typer
+        >>> ctx = typer.Context(...)  # Context with settings and persona
+        >>> attributes = await get_extra_attributes(ctx)
+        >>> print(attributes)
+    """
+    path = "/admin/management/organizations/extra-attributes"
+    try:
+        # The VantageRestApiClient.get() method returns the JSON data directly
+        # (not a response object), so we can return it as-is
+        data = await ctx.obj.rest_client.get(path)
+        return data
+    except Exception as e:
+        logger.warning("Failed to retrieve extra attributes: %s", e)
+        return None

vantage_sdk/auth.py

@@ -0,0 +1,517 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Authentication and authorization functionality for the Vantage CLI."""
+
+import asyncio
+import base64
+import datetime
+import hashlib
+import logging
+import secrets
+from textwrap import dedent
+from typing import Any, Callable
+
+import httpx
+import snick
+from jose import jwt
+from jose.exceptions import ExpiredSignatureError
+from pydantic import ValidationError
+
+from vantage_sdk.client import make_oauth_request
+from vantage_sdk.config import Settings
+from vantage_sdk.constants import (
+    OIDC_DEVICE_PATH,
+    OIDC_SCOPES,
+    OIDC_TOKEN_PATH,
+    TOKEN_REFRESH_THRESHOLD_SECONDS,
+)
+from vantage_sdk.exceptions import Abort
+from vantage_sdk.schemas import CliContext, DeviceCodeData, IdentityData, Persona, TokenSet
+
+logger = logging.getLogger(__name__)
+
+
+def generate_pkce_pair() -> tuple[str, str]:
+    """Generate a PKCE code verifier and code challenge pair.
+
+    Returns:
+        Tuple of (code_verifier, code_challenge)
+    """
+    # Generate a random code verifier (43-128 characters)
+    code_verifier = secrets.token_urlsafe(64)
+
+    # Create code challenge using SHA256 and base64url encoding
+    code_challenge_bytes = hashlib.sha256(code_verifier.encode("ascii")).digest()
+    code_challenge = base64.urlsafe_b64encode(code_challenge_bytes).rstrip(b"=").decode("ascii")
+
+    return code_verifier, code_challenge
+
+
+def extract_persona(token_set: TokenSet, settings: Settings | None = None) -> Persona:
+    """Extract a persona from an explicitly provided token set."""
+    token_set = refresh_token_if_needed(token_set, settings=settings)
+
+    # Now validate and extract identity from the (potentially refreshed) token
+    identity_data = validate_token_and_extract_identity(token_set)
+
+    logger.debug(f"Persona created with identity_data: {identity_data}")
+
+    return Persona(
+        token_set=token_set,
+        identity_data=identity_data,
+    )
+
+
+def validate_token_and_extract_identity(token_set: TokenSet) -> IdentityData:
+    """Validate access token and extract user identity information."""
+    logger.debug("Validating access token")
+
+    token_file_is_empty = not token_set.access_token
+    if token_file_is_empty:
+        logger.debug("Access token file exists but it is empty")
+        raise Abort(
+            """
+            Access token file exists but it is empty.
+
+            Please try logging in again.
+            """,
+            subject="Empty access token file",
+            log_message="Empty access token file",
+        )
+
+    with Abort.handle_errors(
+        """
+        There was an unknown error while validating the access token.
+
+        Please try logging in again.
+        """,
+        ignore_exc_class=ExpiredSignatureError,  # Will be handled in calling context
+        raise_kwargs={
+            "subject": "Invalid access token",
+            "log_message": "Unknown error while validating access access token",
+        },
+    ):
+        token_data = jwt.decode(
+            token_set.access_token,
+            "",  # Empty key is acceptable when verify_signature is False
+            options={
+                "verify_signature": False,
+                "verify_aud": False,
+                "verify_exp": True,
+            },
+        )
+
+    logger.debug("Extracting identity data from the access token")
+    with Abort.handle_errors(
+        """
+        There was an error extracting the user's identity from the access token.
+
+        Please try logging in again.
+        """,
+        handle_exc_class=ValidationError,
+        raise_kwargs={
+            "subject": "Missing user data",
+            "log_message": "Token data could not be extracted to identity",
+        },
+    ):
+        if "organization" not in token_data or not token_data["organization"]:
+            raise Abort(
+                """
+                The access token is missing organization information.
+
+                Please ensure your user account is associated with an organization
+                and try logging in again.
+                """,
+                subject="Missing organization info",
+                log_message="Access token missing organization information",
+            )
+
+        # Extract org_id from organization structure
+        # Organization is typically: {"org-uuid": {"id": "org-uuid", ...}}
+        organization = token_data.get("organization", {})
+        logger.debug(f"Organization data extracted from token: {organization}")
+        org_key = next(iter(organization), None)
+        logger.debug(f"Organization key identified: {org_key}")
+        org_id = organization.get(org_key, {}).get("id", "") if org_key else ""
+        logger.debug(f"Extracted org_id: {org_id}")
+
+        email = token_data.get("email") or ""
+        # Extract username from email (part before @) or use preferred_username claim
+        username = token_data.get("preferred_username") or email.split("@")[0] if email else ""
+        identity = IdentityData(
+            email=email,
+            client_id=token_data.get("azp") or "unknown",
+            org_id=org_id,
+            org_name=org_key or "",
+            username=username,
+        )
+        logger.debug(f"Extracted identity data: {identity}")
+
+    return identity
+
+
+def is_token_expired(token: str, buffer_seconds: int = TOKEN_REFRESH_THRESHOLD_SECONDS) -> bool:
+    """Check if a JWT token is expired or will expire within buffer_seconds.
+
+    Args:
+        token: JWT access token
+        buffer_seconds: Number of seconds before actual expiry to consider token expired.
+            Defaults to TOKEN_REFRESH_THRESHOLD_SECONDS (300s) for proactive refresh.
+
+    Returns:
+        True if token is expired or will expire soon, False otherwise
+    """
+    try:
+        # Decode token without verification to get expiration
+        token_data = jwt.decode(
+            token,
+            "",  # Empty key is acceptable when verify_signature is False
+            options={
+                "verify_signature": False,
+                "verify_aud": False,
+                "verify_exp": False,  # Don't verify expiration here, we want to check manually
+            },
+        )
+
+        if "exp" not in token_data:
+            logger.debug("Token does not contain expiration claim")
+            return True  # Consider token expired if no expiration claim
+
+        exp_timestamp = token_data["exp"]
+        exp_datetime = datetime.datetime.fromtimestamp(exp_timestamp)
+        now_with_buffer = datetime.datetime.now() + datetime.timedelta(seconds=buffer_seconds)
+
+        is_expired = exp_datetime <= now_with_buffer
+
+        if is_expired:
+            logger.debug(
+                f"Token expired or will expire soon. Expires at: {exp_datetime}, Current time + buffer: {now_with_buffer}"
+            )
+        else:
+            logger.debug(
+                f"Token is valid. Expires at: {exp_datetime}, Current time + buffer: {now_with_buffer}"
+            )
+
+        return is_expired
+
+    except Exception as e:
+        logger.debug(f"Error checking token expiration: {e}")
+        return True  # Consider token expired if we can't parse it
+
+
+def refresh_token_if_needed(token_set: TokenSet, settings: Settings | None = None) -> TokenSet:
+    """Check if the access token is expired and refresh it if needed.
+
+    Args:
+        token_set: Current token set
+        settings: Explicit settings used for refresh requests
+
+    Returns:
+        Updated token set with refreshed tokens if refresh was needed
+    """
+    if not token_set.access_token:
+        logger.debug("No access token available")
+        return token_set
+
+    if not is_token_expired(token_set.access_token):
+        logger.debug("Access token is still valid")
+        return token_set
+
+    logger.debug("Access token is expired, attempting refresh")
+
+    if not token_set.refresh_token:
+        logger.debug("No refresh token available")
+        raise Abort(
+            "The access token is expired and no refresh token is available. Please log in again.",
+            subject="Token expired",
+            log_message="Token expired and no refresh token available",
+        )
+
+    try:
+        if settings is None:
+            raise Abort(
+                "The access token is expired and SDK settings were not provided for refresh.",
+                subject="Missing SDK settings",
+                log_message="Token refresh requested without explicit settings",
+            )
+
+        # Attempt to refresh the token
+        refresh_success = refresh_access_token_standalone(token_set, settings)
+
+        if refresh_success:
+            logger.debug("Successfully refreshed access token")
+            return token_set
+        else:
+            logger.warning("Token refresh failed - check error logs above for details")
+            raise Exception("Token refresh failed")
+
+    except Exception as e:
+        logger.debug(f"Token refresh error: {e}")
+        raise Abort(
+            dedent(
+                """\
+                Your authentication session has expired.
+
+                Please log in again by running:
+
+                    vantage login
+                """
+            ),
+            subject="Authentication Required",
+            log_message=f"Token refresh failed: {e}",
+        )
+
+
+def init_persona(
+    *,
+    access_token: str,
+    refresh_token: str | None = None,
+    settings: Settings | None = None,
+) -> Persona:
+    """Initialize a persona from explicit token values without filesystem access."""
+    token_set = TokenSet(access_token=access_token, refresh_token=refresh_token)
+
+    try:
+        identity_data = validate_token_and_extract_identity(token_set)
+    except ExpiredSignatureError:
+        Abort.require_condition(
+            token_set.refresh_token is not None,
+            "The auth token is expired. Please retrieve a new and log in again.",
+            raise_kwargs={
+                "subject": "Expired access token",
+            },
+        )
+
+        logger.debug("The access token is expired. Attempting to refresh token")
+        resolved_settings = settings or Settings()
+        refresh_success = refresh_access_token_standalone(token_set, resolved_settings)
+        if not refresh_success:
+            raise Exception("Failed to refresh access token")
+        identity_data = validate_token_and_extract_identity(token_set)
+
+    logger.debug(f"Persona created with identity_data: {identity_data}")
+
+    return Persona(
+        token_set=token_set,
+        identity_data=identity_data,
+    )
+
+
+def refresh_access_token_standalone(token_set: TokenSet, settings: "Settings") -> bool:
+    """Attempt to fetch a new access token given a refresh token.
+
+    Returns True if refresh was successful, False otherwise.
+    Sets the access token in-place.
+    """
+    if not token_set.refresh_token:
+        return False
+
+    url = f"{settings.get_auth_url()}{OIDC_TOKEN_PATH}"
+    logger.debug(f"Requesting refreshed access token from {url}")
+
+    try:
+        with httpx.Client() as client:
+            response = client.post(
+                url,
+                data={
+                    "client_id": settings.oidc_client_id,
+                    "grant_type": "refresh_token",
+                    "refresh_token": token_set.refresh_token,
+                },
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=30.0,
+            )
+            response.raise_for_status()
+
+            token_data = response.json()
+            token_set.access_token = token_data["access_token"]
+
+            # Update refresh token if provided
+            if "refresh_token" in token_data:
+                token_set.refresh_token = token_data["refresh_token"]
+
+            logger.debug("Successfully refreshed access token")
+            return True
+
+    except httpx.HTTPStatusError as e:
+        logger.debug(
+            f"Token refresh failed with status {e.response.status_code}: {e.response.text}"
+        )
+        return False
+    except httpx.TimeoutException as e:
+        logger.debug(f"Token refresh timed out: {e}")
+        return False
+    except httpx.RequestError as e:
+        logger.debug(f"Token refresh request error: {e}")
+        return False
+    except KeyError as e:
+        logger.debug(f"Token refresh response missing required field: {e}")
+        return False
+    except Exception as e:
+        logger.debug(f"Unexpected error during token refresh: {type(e).__name__}: {e}")
+        return False
+
+
+async def refresh_access_token(ctx: CliContext, token_set: TokenSet):
+    """Attempt to fetch a new access token given a refresh token in a token_set.
+
+    Sets the access token in-place.
+
+    If refresh fails, notify the user that they need to log in again.
+    """
+    if ctx.client is None:
+        raise RuntimeError("HTTP client not initialized")
+    if ctx.settings is None:
+        raise RuntimeError("Settings not initialized")
+
+    url = "/realms/vantage/protocol/openid-connect/token"
+    logger.debug(f"Requesting refreshed access token from {url}")
+
+    refreshed_token_set: TokenSet = await make_oauth_request(
+        ctx.client,
+        url,
+        data={
+            "client_id": ctx.settings.oidc_client_id,
+            "grant_type": "refresh_token",
+            "refresh_token": token_set.refresh_token,
+        },
+        response_model_cls=TokenSet,
+        abort_message="The auth token could not be refreshed. Please try logging in again.",
+        abort_subject="EXPIRED ACCESS TOKEN",
+    )
+
+    token_set.access_token = refreshed_token_set.access_token
+
+
+async def fetch_auth_tokens(
+    ctx: CliContext, status_callback: Callable[[str], None] | None = None
+) -> TokenSet:
+    """Fetch an access token (and possibly a refresh token) from Auth0.
+
+    Prints out a URL for the user to use to authenticate and polls the token endpoint to fetch it
+    when the browser-based process finishes.
+    """
+    if ctx.client is None:
+        raise RuntimeError("HTTP client not initialized")
+    if ctx.settings is None:
+        raise RuntimeError("Settings not initialized")
+
+    # Use console from context - it should always be available
+    console = ctx.console
+
+    # Generate PKCE code verifier and challenge
+    code_verifier, code_challenge = generate_pkce_pair()
+
+    device_code_data: DeviceCodeData = await make_oauth_request(
+        ctx.client,
+        OIDC_DEVICE_PATH,
+        data={
+            "client_id": ctx.settings.oidc_client_id,
+            "code_challenge": code_challenge,
+            "code_challenge_method": "S256",
+            "scope": OIDC_SCOPES,
+        },
+        response_model_cls=DeviceCodeData,
+        abort_message=(
+            """
+            There was a problem retrieving a device verification code from
+            the auth provider
+            """
+        ),
+        abort_subject="COULD NOT RETRIEVE TOKEN",
+    )
+
+    max_poll_time = 5 * 60  # 5 minutes
+    login_message = dedent(
+        f"""
+        To complete login, please open the following link in a browser:
+
+          {device_code_data.verification_uri_complete}
+
+        Waiting up to {max_poll_time / 60} minutes for you to complete the process...
+        """
+    ).strip()
+
+    if status_callback is not None:
+        status_callback(login_message)
+    else:
+        logger.info(login_message)
+
+    # Calculate timeout and start time
+    start_time = datetime.datetime.now()
+    timeout_seconds = ctx.settings.oidc_max_poll_time  # This is already in seconds (int)
+    attempt = 0
+
+    while True:
+        attempt += 1
+        elapsed = (datetime.datetime.now() - start_time).total_seconds()
+
+        if elapsed >= timeout_seconds:
+            break
+
+        response_data: dict[str, Any] = {}
+        try:
+            token_data = await make_oauth_request(
+                ctx.client,
+                OIDC_TOKEN_PATH,
+                data={
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    "device_code": device_code_data.device_code,
+                    "client_id": ctx.settings.oidc_client_id,
+                    "code_verifier": code_verifier,
+                },
+                response_model_cls=TokenSet,
+                abort_message="IGNORE",
+                abort_subject="IGNORE",
+            )
+            return token_data
+        except Exception:
+            response = await ctx.client.post(
+                f"{ctx.settings.get_auth_url()}{OIDC_TOKEN_PATH}",
+                data={
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    "device_code": device_code_data.device_code,
+                    "client_id": ctx.settings.oidc_client_id,
+                    "code_verifier": code_verifier,
+                },
+            )
+            response_data = response.json()
+
+        if "error" in response_data:
+            if response_data["error"] == "authorization_pending":
+                logger.debug(f"Token fetch attempt #{attempt} failed")
+                logger.debug(f"Will try again in {device_code_data.interval} seconds")
+                await asyncio.sleep(device_code_data.interval)
+            elif response_data["error"] == "slow_down":
+                logger.debug(f"Server requested slow down on attempt #{attempt}")
+                logger.debug(f"Will try again in {device_code_data.interval * 2} seconds")
+                await asyncio.sleep(device_code_data.interval * 2)
+            else:
+                raise Abort(
+                    snick.unwrap(
+                        """
+                        There was a problem retrieving a device verification code
+                        from the auth provider:
+                        Unexpected failure retrieving access token.
+                        """
+                    ),
+                    subject="Unexpected error",
+                    log_message=f"Unexpected error response: {response_data}",
+                )
+        else:
+            return TokenSet(**response_data)
+
+    raise Abort(
+        "Login process was not completed in time. Please try again.",
+        subject="Timed out",
+        log_message="Timed out while waiting for user to complete login",
+    )

vantage_sdk/base/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (C) 2025 Vantage Compute Corporation
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the Free Software
+# Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+"""Base SDK classes for CRUD operations."""
+
+from .crud import BaseCRUDSDK, BaseGraphQLResourceSDK, BaseLocalResourceSDK, BaseRestApiResourceSDK
+
+__all__ = [
+    "BaseCRUDSDK",
+    "BaseLocalResourceSDK",
+    "BaseGraphQLResourceSDK",
+    "BaseRestApiResourceSDK",
+]