blaxel 0.64.0__py3-none-any.whl

blaxel/authentication/clientcredentials.py

@@ -0,0 +1,103 @@
+"""
+This module provides the ClientCredentials class, which handles client credentials-based
+authentication for Blaxel. It manages token refreshing and authentication flows using
+client credentials and refresh tokens.
+"""
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Generator, Optional
+
+import requests
+from httpx import Auth, Request, Response
+
+from blaxel.aimon.settings import get_settings
+
+
+@dataclass
+class DeviceLoginFinalizeResponse:
+    access_token: str
+    expires_in: int
+    refresh_token: str
+    token_type: str
+
+
+class ClientCredentials(Auth):
+    """
+    A provider that authenticates requests using client credentials.
+    """
+
+    def __init__(self, credentials, workspace_name: str, base_url: str):
+        """
+        Initializes the ClientCredentials provider with the given credentials, workspace name, and base URL.
+
+        Parameters:
+            credentials: Credentials containing access and refresh tokens.
+            workspace_name (str): The name of the workspace.
+            base_url (str): The base URL for authentication.
+        """
+        self.credentials = credentials
+        self.expires_at = None
+        self.workspace_name = workspace_name
+        self.base_url = base_url
+
+    def get_headers(self):
+        """
+        Retrieves the authentication headers after ensuring tokens are valid.
+
+        Returns:
+            dict: A dictionary of headers with Bearer token and workspace.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
+        err = self.get_token()
+        if err:
+            raise err
+        return {
+            "X-Blaxel-Authorization": f"Bearer {self.credentials.access_token}",
+            "X-Blaxel-Workspace": self.workspace_name,
+        }
+
+    def get_token(self) -> Optional[Exception]:
+        """
+        Checks if the access token needs to be refreshed and performs the refresh if necessary.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
+        settings = get_settings()
+        if self.need_token():
+            headers = {"Authorization": f"Basic {self.credentials.client_credentials}", "Content-Type": "application/json"}
+            body = {"grant_type": "client_credentials"}
+            response = requests.post(f"{settings.base_url}/oauth/token", headers=headers, json=body)
+            response.raise_for_status()
+            creds = response.json()
+            self.credentials.access_token = creds["access_token"]
+            self.credentials.refresh_token = creds["refresh_token"]
+            self.credentials.expires_in = creds["expires_in"]
+            self.expires_at = datetime.now() + timedelta(seconds=self.credentials.expires_in)
+        return None
+
+    def need_token(self):
+        if not self.expires_at:
+            return True
+        return datetime.now() > self.expires_at - timedelta(minutes=10)
+
+    def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Processes the authentication flow by ensuring tokens are valid and adding necessary headers.
+
+        Parameters:
+            request (Request): The HTTP request to authenticate.
+
+        Yields:
+            Request: The authenticated request.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
+        self.get_token()
+        request.headers["X-Blaxel-Authorization"] = f"Bearer {self.credentials.access_token}"
+        request.headers["X-Blaxel-Workspace"] = self.workspace_name
+        yield request

blaxel/authentication/credentials.py

@@ -0,0 +1,295 @@
+"""
+This module provides classes and functions for managing credentials and workspace configurations.
+It includes functionalities to load, save, and manage authentication credentials, as well as to handle
+workspace contexts and configurations.
+"""
+
+from dataclasses import dataclass
+from logging import getLogger
+from pathlib import Path
+from typing import List
+
+import yaml
+
+from blaxel.aimon.settings import Settings
+
+logger = getLogger(__name__)
+
+
+@dataclass
+class Credentials:
+    """
+    A dataclass representing user credentials for authentication.
+
+    Attributes:
+        apiKey (str): The API key.
+        access_token (str): The access token.
+        refresh_token (str): The refresh token.
+        expires_in (int): Token expiration time in seconds.
+        device_code (str): The device code for device authentication.
+        client_credentials (str): The client credentials for authentication.
+    """
+    apiKey: str = ""
+    access_token: str = ""
+    refresh_token: str = ""
+    expires_in: int = 0
+    device_code: str = ""
+    client_credentials: str = ""
+
+
+@dataclass
+class WorkspaceConfig:
+    """
+    A dataclass representing the configuration for a workspace.
+
+    Attributes:
+        name (str): The name of the workspace.
+        credentials (Credentials): The credentials associated with the workspace.
+    """
+    name: str
+    credentials: Credentials
+
+
+@dataclass
+class ContextConfig:
+    """
+    A dataclass representing the current context configuration.
+
+    Attributes:
+        workspace (str): The name of the current workspace.
+    """
+    workspace: str = ""
+
+
+@dataclass
+class Config:
+    """
+    A dataclass representing the overall configuration, including workspaces and context.
+
+    Attributes:
+        workspaces (List[WorkspaceConfig]): A list of workspace configurations.
+        context (ContextConfig): The current context configuration.
+    """
+    workspaces: List[WorkspaceConfig] = None
+    context: ContextConfig = None
+
+    def __post_init__(self):
+        """
+        Post-initialization to ensure workspaces and context are initialized.
+        """
+        if self.workspaces is None:
+            self.workspaces = []
+        if self.context is None:
+            self.context = ContextConfig()
+
+    def to_json(self) -> dict:
+        """
+        Converts the Config dataclass to a JSON-compatible dictionary.
+
+        Returns:
+            dict: The JSON representation of the configuration.
+        """
+        return {
+            "workspaces": [
+                {
+                    "name": ws.name,
+                    "credentials": {
+                        "apiKey": ws.credentials.apiKey,
+                        "access_token": ws.credentials.access_token,
+                        "refresh_token": ws.credentials.refresh_token,
+                        "expires_in": ws.credentials.expires_in,
+                        "device_code": ws.credentials.device_code,
+                        "client_credentials": ws.credentials.client_credentials,
+                    },
+                }
+                for ws in self.workspaces
+            ],
+            "context": {
+                "workspace": self.context.workspace,
+            },
+        }
+
+
+def load_config() -> Config:
+    """
+    Loads the configuration from the user's home directory.
+
+    Returns:
+        Config: The loaded configuration.
+    """
+    config = Config()
+    home_dir = Path.home()
+    if home_dir:
+        config_path = home_dir / ".blaxel" / "config.yaml"
+        if config_path.exists():
+            try:
+                with open(config_path) as f:
+                    data = yaml.safe_load(f)
+                    if data:
+                        workspaces = []
+                        for ws in data.get("workspaces", []):
+                            creds = Credentials(**ws.get("credentials", {}))
+                            workspaces.append(WorkspaceConfig(name=ws["name"], credentials=creds))
+                        config.workspaces = workspaces
+                        if "context" in data:
+                            config.context = ContextConfig(workspace=data["context"].get("workspace", ""))
+            except yaml.YAMLError:
+                # Invalid YAML, use empty config
+                pass
+    return config
+
+
+def save_config(config: Config):
+    """
+    Saves the provided configuration to the user's home directory.
+
+    Parameters:
+        config (Config): The configuration to save.
+
+    Raises:
+        RuntimeError: If the home directory cannot be determined.
+    """
+    home_dir = Path.home()
+    if not home_dir:
+        raise RuntimeError("Could not determine home directory")
+
+    config_dir = home_dir / ".blaxel"
+    config_file = config_dir / "config.yaml"
+
+    config_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+    with open(config_file, "w", encoding="utf-8") as f:
+        yaml.dump(config.to_json(), f)
+
+
+def list_workspaces() -> List[str]:
+    """
+    Lists all available workspace names from the configuration.
+
+    Returns:
+        List[str]: A list of workspace names.
+    """
+    config = load_config()
+    return [workspace.name for workspace in config.workspaces]
+
+
+def current_context() -> ContextConfig:
+    """
+    Retrieves the current context configuration.
+
+    Returns:
+        ContextConfig: The current context configuration.
+    """
+    config = load_config()
+    return config.context
+
+
+def set_current_workspace(workspace_name: str):
+    """
+    Sets the current workspace in the configuration.
+
+    Parameters:
+        workspace_name (str): The name of the workspace to set as current.
+    """
+    config = load_config()
+    config.context.workspace = workspace_name
+    save_config(config)
+
+
+def load_credentials(workspace_name: str) -> Credentials:
+    """
+    Loads credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace whose credentials are to be loaded.
+
+    Returns:
+        Credentials: The credentials associated with the workspace. Returns empty credentials if not found.
+    """
+    config = load_config()
+    for workspace in config.workspaces:
+        if workspace.name == workspace_name:
+            return workspace.credentials
+    return Credentials()
+
+
+def load_credentials_from_settings(settings: Settings) -> Credentials:
+    """
+    Loads credentials from the provided settings.
+
+    Parameters:
+        settings (Settings): The settings containing authentication information.
+
+    Returns:
+        Credentials: The loaded credentials from settings.
+    """
+    return Credentials(
+        apiKey=settings.authentication.apiKey,
+        client_credentials=settings.authentication.client.credentials,
+    )
+
+
+def create_home_dir_if_missing():
+    """
+    Creates the Blaxel home directory if it does not exist.
+
+    Logs a warning if credentials already exist or an error if directory creation fails.
+    """
+    home_dir = Path.home()
+    if not home_dir:
+        logger.error("Error getting home directory")
+        return
+
+    credentials_dir = home_dir / ".blaxel"
+    credentials_file = credentials_dir / "credentials.json"
+
+    if credentials_file.exists():
+        logger.warning("You are already logged in. Enter a new API key to overwrite it.")
+    else:
+        try:
+            credentials_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
+        except Exception as e:
+            logger.error(f"Error creating credentials directory: {e}")
+
+
+def save_credentials(workspace_name: str, credentials: Credentials):
+    """
+    Saves the provided credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace.
+        credentials (Credentials): The credentials to save.
+    """
+    create_home_dir_if_missing()
+    if not credentials.access_token and not credentials.apiKey:
+        logger.info("No credentials to save, error")
+        return
+
+    config = load_config()
+    found = False
+
+    for i, workspace in enumerate(config.workspaces):
+        if workspace.name == workspace_name:
+            config.workspaces[i].credentials = credentials
+            found = True
+            break
+
+    if not found:
+        config.workspaces.append(WorkspaceConfig(name=workspace_name, credentials=credentials))
+
+    save_config(config)
+
+
+def clear_credentials(workspace_name: str):
+    """
+    Clears the credentials for the specified workspace.
+
+    Parameters:
+        workspace_name (str): The name of the workspace whose credentials are to be cleared.
+    """
+    config = load_config()
+    config.workspaces = [ws for ws in config.workspaces if ws.name != workspace_name]
+
+    if config.context.workspace == workspace_name:
+        config.context.workspace = ""
+
+    save_config(config)

blaxel/authentication/device_mode.py

@@ -0,0 +1,197 @@
+"""
+This module provides classes for handling device-based authentication,
+including device login processes and bearer token management. It facilitates token refreshing
+and authentication flows using device codes and bearer tokens.
+"""
+
+import base64
+import json
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, Generator, Optional
+
+from httpx import Auth, Request, Response, post
+
+
+@dataclass
+class DeviceLogin:
+    """
+    A dataclass representing a device login request.
+
+    Attributes:
+        client_id (str): The client ID for the device.
+        scope (str): The scope of the authentication.
+    """
+    client_id: str
+    scope: str
+
+
+@dataclass
+class DeviceLoginResponse:
+    """
+    A dataclass representing the response from a device login request.
+
+    Attributes:
+        client_id (str): The client ID associated with the device login.
+        device_code (str): The device code for authentication.
+        user_code (str): The user code for completing authentication.
+        expires_in (int): Time in seconds until the device code expires.
+        interval (int): Polling interval in seconds.
+        verification_uri (str): URI for user to verify device login.
+        verification_uri_complete (str): Complete URI including the user code for verification.
+    """
+    client_id: str
+    device_code: str
+    user_code: str
+    expires_in: int
+    interval: int
+    verification_uri: str
+    verification_uri_complete: str
+
+
+@dataclass
+class DeviceLoginFinalizeRequest:
+    """
+    A dataclass representing a device login finalize request.
+
+    Attributes:
+        grant_type (str): The type of grant being requested.
+        client_id (str): The client ID for finalizing the device login.
+        device_code (str): The device code to finalize login.
+    """
+    grant_type: str
+    client_id: str
+    device_code: str
+
+
+@dataclass
+class DeviceLoginFinalizeResponse:
+    access_token: str
+    expires_in: int
+    refresh_token: str
+    token_type: str
+
+
+class BearerToken(Auth):
+    """
+    A provider that authenticates requests using a Bearer token.
+    """
+
+    def __init__(self, credentials, workspace_name: str, base_url: str):
+        """
+        Initializes the BearerToken provider with the given credentials, workspace name, and base URL.
+
+        Parameters:
+            credentials: Credentials containing the Bearer token and refresh token.
+            workspace_name (str): The name of the workspace.
+            base_url (str): The base URL for authentication.
+        """
+        self.credentials = credentials
+        self.workspace_name = workspace_name
+        self.base_url = base_url
+
+    def get_headers(self) -> Dict[str, str]:
+        """
+        Retrieves the authentication headers containing the Bearer token and workspace information.
+
+        Returns:
+            Dict[str, str]: A dictionary of headers with Bearer token and workspace.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
+        err = self.refresh_if_needed()
+        if err:
+            raise err
+        return {
+            "X-Blaxel-Authorization": f"Bearer {self.credentials.access_token}",
+            "X-Blaxel-Workspace": self.workspace_name,
+        }
+
+    def refresh_if_needed(self) -> Optional[Exception]:
+        """
+        Checks if the Bearer token needs to be refreshed and performs the refresh if necessary.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
+        # Need to refresh token if expires in less than 10 minutes
+        parts = self.credentials.access_token.split(".")
+        if len(parts) != 3:
+            return Exception("Invalid JWT token format")
+
+        try:
+            claims_bytes = base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))
+            claims = json.loads(claims_bytes)
+        except Exception as e:
+            return Exception(f"Failed to decode/parse JWT claims: {str(e)}")
+        exp_time = datetime.fromtimestamp(claims["exp"])
+        current_time = datetime.now()
+        # Refresh if token expires in less than 10 minutes
+        if current_time + timedelta(minutes=10) > exp_time:
+            return self.do_refresh()
+
+        return None
+
+    def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
+        """
+        Processes the authentication flow by ensuring the Bearer token is valid and adding necessary headers.
+
+        Parameters:
+            request (Request): The HTTP request to authenticate.
+
+        Yields:
+            Request: The authenticated request.
+
+        Raises:
+            Exception: If token refresh fails.
+        """
+        err = self.refresh_if_needed()
+        if err:
+            return err
+
+        request.headers["X-Blaxel-Authorization"] = f"Bearer {self.credentials.access_token}"
+        request.headers["X-Blaxel-Workspace"] = self.workspace_name
+        yield request
+
+    def do_refresh(self) -> Optional[Exception]:
+        """
+        Performs the token refresh using the refresh token.
+
+        Returns:
+            Optional[Exception]: An exception if refreshing fails, otherwise None.
+        """
+        if not self.credentials.refresh_token:
+            return Exception("No refresh token to refresh")
+
+        url = f"{self.base_url}/oauth/token"
+        refresh_data = {
+            "grant_type": "refresh_token",
+            "refresh_token": self.credentials.refresh_token,
+            "device_code": self.credentials.device_code,
+            "client_id": "blaxel",
+        }
+
+        try:
+            response = post(url, json=refresh_data, headers={"Content-Type": "application/json"})
+            response.raise_for_status()
+            finalize_response = DeviceLoginFinalizeResponse(**response.json())
+
+            if not finalize_response.refresh_token:
+                finalize_response.refresh_token = self.credentials.refresh_token
+
+            from .credentials import Credentials, save_credentials
+
+            creds = Credentials(
+                access_token=finalize_response.access_token,
+                refresh_token=finalize_response.refresh_token,
+                expires_in=finalize_response.expires_in,
+                device_code=self.credentials.device_code,
+            )
+
+            self.credentials = creds
+            save_credentials(self.workspace_name, creds)
+            return None
+
+        except Exception as e:
+            return Exception(f"Failed to refresh token: {str(e)}")