antioch-py 2.0.6__py3-none-any.whl → 3.0.12__py3-none-any.whl

common/core/auth.py

@@ -7,25 +7,18 @@ from pathlib import Path
 import requests
 from pydantic import BaseModel
 
-from common.constants import get_auth_dir
-
-# Authentication routes
-AUTH_DOMAIN = os.environ.get("AUTH_DOMAIN", "https://staging.auth.antioch.com")
-AUTH_TOKEN_URL = f"{AUTH_DOMAIN}/oauth/token"
-DEVICE_CODE_URL = f"{AUTH_DOMAIN}/oauth/device/code"
-
-# Authentication constants
-AUTH_CLIENT_ID = "x0aOquV43Xe76ehqAm6Zir80O0MWpqTV"
-ALGORITHMS = ["RS256"]
-AUDIENCE = "https://sessions.antioch.com"
-AUTH_SCOPE = "openid profile email"
-AUTH_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code"
-AUTH_TIMEOUT_SECONDS = 120
-
-# Authentication claims
-AUTH_ORG_ID_CLAIM = "https://antioch.com/org_id"
-AUTH_ORG_NAME_CLAIM = "https://antioch.com/org_name"
-AUTH_ORGANIZATIONS_CLAIM = "https://antioch.com/organizations"
+from common.constants import (
+    AUDIENCE,
+    AUTH_CLIENT_ID,
+    AUTH_GRANT_TYPE,
+    AUTH_ORG_ID_CLAIM,
+    AUTH_ORG_NAME_CLAIM,
+    AUTH_SCOPE,
+    AUTH_TIMEOUT_SECONDS,
+    AUTH_TOKEN_URL,
+    DEVICE_CODE_URL,
+    get_auth_dir,
+)
 
 
 class AuthError(Exception):
@@ -36,38 +29,51 @@ class AuthError(Exception):
 
 class Organization(BaseModel):
     """
-    Organization information.
+    Organization information extracted from JWT token.
     """
 
     org_id: str
     org_name: str
 
 
+class UserInfo(BaseModel):
+    """
+    User information extracted from JWT token.
+    """
+
+    user_id: str
+    name: str | None = None
+    email: str | None = None
+
+
 class AuthHandler:
     """
-    Client for handling authentication.
+    Client for handling authentication via OAuth2 device code flow.
 
-    Auth is used for:
-    - Pulling artifacts from the remote Ark registry
-    - Pulling assets from the remote asset registry
+    Manages authentication tokens and organization context for interacting
+    with Antioch services.
     """
 
     def __init__(self):
         """
         Initialize the auth handler.
+
+        Automatically loads any existing token from disk.
         """
 
         self._token: str | None = None
-        self._user_id: str | None = None
-        self._current_org: Organization | None = None
-        self._available_orgs: list[Organization] = []
+        self._org: Organization | None = None
+        self._user: UserInfo | None = None
         self._load_local_token()
 
     def login(self) -> None:
         """
-        Authenticate the user via device code flow.
+        Authenticate the user via OAuth2 device code flow.
 
-        :raises AuthError: If authentication fails.
+        Initiates the device code flow, prompts the user to authenticate
+        in their browser, and saves the token to disk on success.
+
+        :raises AuthError: If authentication fails or times out.
         """
 
         if self.is_authenticated():
@@ -97,26 +103,27 @@ class AuthHandler:
             "client_id": AUTH_CLIENT_ID,
         }
 
-        authenticated = False
         start_time = time.time()
-        while not authenticated:
+        while True:
             token_response = requests.post(AUTH_TOKEN_URL, data=token_payload)
             token_data = token_response.json()
             if token_response.status_code == 200:
                 print("Authenticated!")
-                authenticated = True
-            elif token_data["error"] not in ("authorization_pending", "slow_down"):
+                self._token = token_data["access_token"]
+                break
+
+            if token_data["error"] not in ("authorization_pending", "slow_down"):
                 print(token_data["error_description"])
                 raise AuthError("Error authenticating the user") from Exception(token_data)
-            else:
-                if time.time() - start_time > AUTH_TIMEOUT_SECONDS:
-                    raise AuthError("Timeout waiting for authentication")
-                time.sleep(device_code_data["interval"])
 
-        # Save token
-        self._token = token_data["access_token"]
+            if time.time() - start_time > AUTH_TIMEOUT_SECONDS:
+                raise AuthError("Timeout waiting for authentication")
+
+            time.sleep(device_code_data["interval"])
+
         if self._token is None:
             raise AuthError("No token received")
+
         self._validate_token_claims(self._token)
         self.save_token()
 
@@ -124,69 +131,44 @@ class AuthHandler:
         """
         Check if the user is authenticated.
 
-        :return: True if authenticated, False otherwise.
+        :return: True if authenticated with a valid token, False otherwise.
         """
 
-        return self._current_org is not None
+        return self._org is not None
 
-    def select_organization(self, org_id: str):
+    def get_org(self) -> Organization:
         """
-        Select the organization to use for the session.
+        Get the current organization.
 
-        :param org_id: The ID of the organization to select.
+        :return: The current organization.
         :raises AuthError: If the user is not authenticated.
         """
 
-        if not self.is_authenticated():
+        if not self.is_authenticated() or self._org is None:
             raise AuthError("Not authenticated. Please login first")
+        return self._org
 
-        for org in self._available_orgs:
-            if org.org_id == org_id:
-                self._current_org = org
-                return
-
-        raise AuthError(f"Organization '{org_id}' is not in your available organizations")
-
-    def get_current_org(self) -> Organization | None:
+    def get_user_info(self) -> UserInfo | None:
         """
-        Get the current organization.
+        Get the current user information.
 
-        :return: The current organization.
+        :return: The current user info, or None if not available.
         :raises AuthError: If the user is not authenticated.
         """
 
         if not self.is_authenticated():
             raise AuthError("Not authenticated. Please login first")
+        return self._user
 
-        return self._current_org
-
-    def get_user_id(self) -> str | None:
-        """
-        Get the user ID.
-
-        :return: The user ID.
+    def get_token(self) -> str:
         """
+        Get the current authentication token.
 
-        return self._user_id
-
-    def get_available_orgs(self) -> list[Organization]:
-        """
-        Get the available organizations.
-
-        :return: The available organizations.
-        """
-
-        return self._available_orgs
-
-    def get_token(self) -> str | None:
-        """
-        Get the token.
-
-        :return: The token.
+        :return: The JWT access token.
         :raises AuthError: If the user is not authenticated.
         """
 
-        if not self.is_authenticated():
+        if not self.is_authenticated() or self._token is None:
             raise AuthError("Not authenticated. Please login first")
         return self._token
 
@@ -194,22 +176,16 @@ class AuthHandler:
         """
         Save the authentication token and organization data to disk.
 
+        Creates the token file with restrictive permissions (0600).
+
         :raises AuthError: If not authenticated.
         """
 
         if not self.is_authenticated():
             raise AuthError("Not authenticated. Please login first")
 
-        stored_data = {
-            "token": self._token,
-            "current_org": self._current_org.model_dump() if self._current_org else None,
-            "available_orgs": [org.model_dump() for org in self._available_orgs],
-        }
-
+        stored_data = {"token": self._token, "org": self._org.model_dump() if self._org else None}
         token_path = self._get_token_path()
-
-        # Create file with restrictive permissions (owner read/write only)
-        # Use os.open to atomically create file with 0o600 permissions
         fd = os.open(token_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
         with os.fdopen(fd, "w") as f:
             json.dump(stored_data, f, indent=2)
@@ -225,9 +201,9 @@ class AuthHandler:
 
     def _load_local_token(self) -> None:
         """
-        Load the authentication token and organization data from disk.
+        Load the authentication token from disk.
 
-        Silently returns if no token exists or if loading fails. Clears invalid tokens.
+        Silently returns if no token exists. Clears invalid or expired tokens.
         """
 
         token_path = self._get_token_path()
@@ -244,17 +220,16 @@ class AuthHandler:
 
             # Validate and extract all claims from token
             self._validate_token_claims(self._token)
+            if stored_data.get("org"):
+                self._org = Organization(**stored_data["org"])
         except Exception as e:
             print(f"Error loading local token: {e}")
-
-            # Clear invalid or expired tokens
             self._token = None
             self.clear_token()
-            return
 
-    def _validate_token_claims(self, token: str):
+    def _validate_token_claims(self, token: str) -> None:
         """
-        Validate the token and extract all claims including user ID and organization information.
+        Validate the token and extract organization and user information.
 
         :param token: The JWT token to validate.
         :raises AuthError: If the token is invalid, expired, or missing required claims.
@@ -264,7 +239,7 @@ class AuthHandler:
         if len(parts) != 3:
             raise AuthError("Invalid token format")
 
-        # Decode the payload (middle part)
+        # Decode the payload
         payload_encoded = parts[1]
         padding = len(payload_encoded) % 4
         if padding:
@@ -277,29 +252,27 @@ class AuthHandler:
         if exp and time.time() > exp:
             raise AuthError("Token has expired")
 
-        # Extract user ID
-        self._user_id = payload.get("sub")
-        if self._user_id is None:
-            raise AuthError("User ID not found in token claims")
-
-        # Extract current organization
-        self._current_org = Organization(
-            org_id=payload.get(AUTH_ORG_ID_CLAIM),
-            org_name=payload.get(AUTH_ORG_NAME_CLAIM),
-        )
-
-        # Extract available organizations
-        # Note: Auth0 returns organizations with "id" and "name" keys, not "org_id" and "org_name"
-        self._available_orgs = [
-            Organization(org_id=org.get("id") or org.get("org_id"), org_name=org.get("name") or org.get("org_name"))
-            for org in payload.get(AUTH_ORGANIZATIONS_CLAIM, [])
-        ]
+        # Extract organization
+        org_id = payload.get(AUTH_ORG_ID_CLAIM)
+        org_name = payload.get(AUTH_ORG_NAME_CLAIM)
+        if not org_id or not org_name:
+            raise AuthError("Organization information not found in token claims")
+        self._org = Organization(org_id=org_id, org_name=org_name)
+
+        # Extract user info (optional claims)
+        user_id = payload.get("sub")
+        if user_id:
+            self._user = UserInfo(
+                user_id=user_id,
+                name=payload.get("name") or payload.get("nickname"),
+                email=payload.get("email"),
+            )
 
     def _get_token_path(self) -> Path:
         """
         Get the token file path.
 
-        :return: Path to the token file.
+        :return: Path to the token.json file.
         """
 
         return get_auth_dir() / "token.json"

common/core/container.py

@@ -0,0 +1,261 @@
+import contextlib
+import time
+from enum import Enum
+
+import docker
+from docker.errors import APIError
+from docker.models.containers import Container
+
+from common.ark import Ark as ArkDefinition, Environment
+from common.ark.module import ModuleImage, ModuleReady, ModuleStart
+from common.constants import ANTIOCH_API_URL
+from common.core.auth import AuthHandler
+from common.core.rome import RomeClient
+from common.utils.comms import CommsSession
+from common.utils.time import now_us
+
+# Container naming prefix for all Antioch module containers
+CONTAINER_PREFIX = "antioch-module-"
+
+# Synchronization paths for module coordination
+ARK_MODULE_READY_PATH = "_ark/module_ready"
+ARK_MODULE_START_PATH = "_ark/module_start"
+
+
+class ContainerSource(str, Enum):
+    """
+    Source location for container images.
+    """
+
+    LOCAL = "Local"
+    REMOTE = "Remote"
+
+
+class ContainerManagerError(Exception):
+    """
+    Raised when container management operations fail.
+    """
+
+    pass
+
+
+class ContainerManager:
+    """
+    Manages Docker containers for Ark modules.
+
+    Handles launching, coordination, and cleanup of module containers.
+    Uses host networking for Zenoh communication.
+    """
+
+    def __init__(self) -> None:
+        """
+        Create a new container manager.
+
+        Initializes Docker client and Zenoh communication session.
+        """
+
+        self._comms = CommsSession()
+        self._client = docker.from_env()
+        self._containers: dict[str, Container] = {}
+
+    def launch_ark(
+        self,
+        ark: ArkDefinition,
+        source: ContainerSource = ContainerSource.LOCAL,
+        environment: Environment = Environment.SIM,
+        debug: bool = False,
+        timeout: float = 30.0,
+    ) -> int:
+        """
+        Launch all module containers for an Ark.
+
+        Stops any existing module containers first to ensure idempotent behavior.
+
+        :param ark: Ark definition to launch.
+        :param source: Container image source (local or remote).
+        :param environment: Environment to run in (sim or real).
+        :param debug: Enable debug mode.
+        :param timeout: Timeout in seconds for modules to become ready.
+        :return: Global start time in microseconds.
+        :raises ContainerManagerError: If environment is incompatible or launch fails.
+        """
+
+        # Validate environment compatibility
+        if ark.capability == Environment.SIM and environment == Environment.REAL:
+            raise ContainerManagerError(f"Ark '{ark.name}' has sim capability but requested for real")
+        if ark.capability == Environment.REAL and environment == Environment.SIM:
+            raise ContainerManagerError(f"Ark '{ark.name}' has real capability but requested for sim")
+
+        # Stop all existing module containers (idempotent)
+        self._stop_all()
+
+        # Get GAR credentials if pulling from remote
+        gar_auth = self._get_gar_auth() if source == ContainerSource.REMOTE else None
+
+        # Build container configs
+        configs: list[tuple[str, str, str]] = []
+        for module in ark.modules:
+            image = self._get_image(module.image, environment)
+            if image is None:
+                raise ContainerManagerError(f"No image for module '{module.name}' in {environment}")
+            if gar_auth is not None:
+                image = f"{gar_auth['registry_host']}/{gar_auth['repository']}/{image}"
+            container_name = f"{CONTAINER_PREFIX}{ark.name.replace('_', '-')}-{module.name.replace('_', '-')}"
+            configs.append((module.name, container_name, image))
+
+        # Pull images if remote
+        if gar_auth is not None:
+            self._pull_images([c[2] for c in configs], gar_auth)
+
+        # Set up ready subscriber before launching
+        ready_sub = self._comms.declare_async_subscriber(ARK_MODULE_READY_PATH)
+
+        # Launch containers
+        ark_json = ark.model_dump_json()
+        for module_name, container_name, image in configs:
+            self._launch(module_name, container_name, image, ark_json, environment, debug)
+
+        # Wait for all modules to be ready
+        pending = {m.name for m in ark.modules}
+        start = time.time()
+        while pending:
+            if time.time() - start > timeout:
+                raise ContainerManagerError(f"Timeout waiting for modules: {', '.join(sorted(pending))}")
+            msg = ready_sub.recv_timeout(ModuleReady, timeout=0.1)
+            if msg is not None:
+                print(f"Module ready: {msg.module_name}")
+                pending.discard(msg.module_name)
+
+        # Broadcast global start time (2s in future for sync)
+        global_start_us = ((now_us() // 1_000_000) + 2) * 1_000_000
+        self._comms.declare_publisher(ARK_MODULE_START_PATH).publish(ModuleStart(global_start_time_us=global_start_us))
+        return global_start_us
+
+    def stop(self, timeout: float = 10.0) -> None:
+        """
+        Stop all module containers.
+
+        :param timeout: Timeout in seconds for container stop operation.
+        """
+
+        self._stop_all(timeout)
+
+    def close(self, timeout: float = 10.0) -> None:
+        """
+        Close the container manager and clean up resources.
+
+        Stops all module containers and closes the Zenoh session.
+
+        :param timeout: Timeout in seconds for container stop operation.
+        """
+
+        self._stop_all(timeout)
+        self._comms.close()
+
+    def _launch(
+        self,
+        module_name: str,
+        container_name: str,
+        image: str,
+        ark_json: str,
+        environment: Environment,
+        debug: bool,
+    ) -> None:
+        """
+        Launch a single module container.
+
+        :param module_name: Name of the module.
+        :param container_name: Docker container name.
+        :param image: Docker image to use.
+        :param ark_json: Serialized Ark definition.
+        :param environment: Environment (sim or real).
+        :param debug: Enable debug mode.
+        :raises ContainerManagerError: If container launch fails.
+        """
+
+        try:
+            container = self._client.containers.run(
+                image=image,
+                name=container_name,
+                environment={
+                    "_MODULE_NAME": module_name,
+                    "_ARK": ark_json,
+                    "_ENVIRONMENT": str(environment.value),
+                    "_DEBUG": str(debug).lower(),
+                },
+                network_mode="host",
+                ipc_mode="host",
+                detach=True,
+                remove=False,
+            )
+            self._containers[container_name] = container
+            print(f"Launched container: {container_name}")
+        except APIError as e:
+            raise ContainerManagerError(f"Failed to launch '{container_name}': {e}") from e
+
+    def _stop_all(self, timeout: float = 10.0) -> None:
+        """
+        Stop all Antioch module containers.
+
+        Finds all containers with the antioch-module- prefix and stops them.
+
+        :param timeout: Timeout in seconds for stop operation.
+        """
+
+        with contextlib.suppress(APIError):
+            for container in self._client.containers.list(all=True):
+                if container.name and container.name.startswith(CONTAINER_PREFIX):
+                    print(f"Stopping container: {container.name}")
+                    self._stop_container(container, timeout)
+
+        self._containers.clear()
+
+    def _stop_container(self, container: Container, timeout: float) -> None:
+        """
+        Stop and remove a single container.
+
+        :param container: Docker container to stop.
+        :param timeout: Timeout in seconds for stop operation.
+        """
+
+        with contextlib.suppress(APIError):
+            container.stop(timeout=int(timeout))
+        with contextlib.suppress(APIError):
+            container.remove(force=True)
+
+    def _get_image(self, image: str | ModuleImage, environment: Environment) -> str | None:
+        """
+        Get the image name for the given environment.
+
+        :param image: Image specification (string or ModuleImage).
+        :param environment: Environment to get image for.
+        :return: Image name or None if not available.
+        """
+
+        if isinstance(image, str):
+            return image
+        return image.sim if environment == Environment.SIM else image.real
+
+    def _get_gar_auth(self) -> dict:
+        """
+        Get Google Artifact Registry authentication credentials.
+
+        :return: Dictionary containing registry host, repository, and access token.
+        """
+
+        auth = AuthHandler()
+        rome = RomeClient(ANTIOCH_API_URL, auth.get_token())
+        return rome.get_gar_token()
+
+    def _pull_images(self, images: list[str], gar_auth: dict) -> None:
+        """
+        Pull container images from the registry.
+
+        :param images: List of image names to pull.
+        :param gar_auth: GAR authentication credentials.
+        """
+
+        auth_config = {"username": "oauth2accesstoken", "password": gar_auth["access_token"]}
+        for image in set(images):
+            print(f"Pulling image: {image}")
+            self._client.images.pull(image, auth_config=auth_config)