Merge pull request #158 from dreadnode/brian/eng-2922-feat-add-platform-download-capability-to-sdk

feat: add platform download capability to sdk

Author: briangreunke

.vscode/settings.json

@@ -15,4 +15,4 @@
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true
-}
+}

dreadnode/api/client.py

@@ -15,13 +15,15 @@
 
 from dreadnode.api.models import (
     AccessRefreshTokenResponse,
+    ContainerRegistryCredentials,
     DeviceCodeResponse,
     ExportFormat,
     GithubTokenResponse,
     MetricAggregationType,
     Project,
     RawRun,
     RawTask,
+    RegistryImageDetails,
     Run,
     RunSummary,
     StatusFilter,
@@ -710,3 +712,55 @@ def get_user_data_credentials(self) -> UserDataCredentials:
         """
         response = self._request("GET", "/user-data/credentials")
         return UserDataCredentials(**response.json())
+
+    # Container registry access
+
+    def get_container_registry_credentials(self) -> ContainerRegistryCredentials:
+        """
+        Retrieves container registry credentials for Docker image access.
+
+        Returns:
+            The container registry credentials object.
+        """
+        response = self.request("POST", "/platform/registry-token")
+        return ContainerRegistryCredentials(**response.json())
+
+    def get_platform_releases(
+        self, tag: str, services: list[str], cli_version: str | None
+    ) -> RegistryImageDetails:
+        """
+        Resolves the platform releases for the current project.
+
+        Returns:
+            The resolved platform releases as a ResolveReleasesResponse object.
+        """
+        payload = {
+            "tag": tag,
+            "services": services,
+            "cli_version": cli_version,
+        }
+        try:
+            response = self.request("POST", "/platform/get-releases", json_data=payload)
+
+        except RuntimeError as e:
+            if "403" in str(e):
+                raise RuntimeError("You do not have access to platform releases.") from e
+
+            if "404" in str(e):
+                if "Image not found" in str(e):
+                    raise RuntimeError("Image not found") from e
+
+                raise RuntimeError(
+                    f"Failed to get platform releases: {e}. The feature is likely disabled on this server"
+                ) from e
+            raise
+        return RegistryImageDetails(**response.json())
+
+    def get_platform_templates(self, tag: str) -> bytes:
+        """
+        Retrieves the available platform templates.
+        """
+        params = {"tag": tag}
+        response = self.request("GET", "/platform/templates/all", params=params)
+        zip_content: bytes = response.content
+        return zip_content

dreadnode/api/models.py

@@ -43,6 +43,33 @@ class UserDataCredentials(BaseModel):
     endpoint: str | None
 
 
+class ContainerRegistryCredentials(BaseModel):
+    registry: str
+    username: str
+    password: str
+    expires_at: datetime
+
+
+class PlatformImage(BaseModel):
+    service: str
+    uri: str
+    digest: str
+    tag: str
+
+    @property
+    def full_uri(self) -> str:
+        return f"{self.uri}@{self.digest}"
+
+    @property
+    def registry(self) -> str:
+        return self.uri.split("/")[0]
+
+
+class RegistryImageDetails(BaseModel):
+    tag: str
+    images: list[PlatformImage]
+
+
 # Auth
 
 

dreadnode/cli/main.py

@@ -20,6 +20,7 @@
     download_and_unzip_archive,
     validate_server_for_clone,
 )
+from dreadnode.cli.platform import cli as platform_cli
 from dreadnode.cli.profile import cli as profile_cli
 from dreadnode.constants import DEBUG, PLATFORM_BASE_URL
 from dreadnode.user_config import ServerConfig, UserConfig
@@ -28,8 +29,9 @@
 
 cli["--help"].group = "Meta"
 
-cli.command(profile_cli)
 cli.command(agent_cli)
+cli.command(platform_cli)
+cli.command(profile_cli)
 
 
 @cli.meta.default

dreadnode/cli/platform/__init__.py

@@ -0,0 +1,3 @@
+from dreadnode.cli.platform.cli import cli
+
+__all__ = ["cli"]

dreadnode/cli/platform/cli.py

@@ -0,0 +1,61 @@
+import cyclopts
+
+from dreadnode.cli.platform.configure import configure_platform
+from dreadnode.cli.platform.download import download_platform
+from dreadnode.cli.platform.login import log_into_registries
+from dreadnode.cli.platform.start import start_platform
+from dreadnode.cli.platform.stop import stop_platform
+from dreadnode.cli.platform.upgrade import upgrade_platform
+
+cli = cyclopts.App("platform", help="Run and manage the platform.", help_flags=[])
+
+
+@cli.command()
+def start(tag: str | None = None) -> None:
+    """Start the platform. Optionally, provide a tagged version to start.
+
+    Args:
+        tag: Optional image tag to use when starting the platform.
+    """
+    start_platform(tag=tag)
+
+
+@cli.command(name=["stop", "down"])
+def stop() -> None:
+    """Stop the running platform."""
+    stop_platform()
+
+
+@cli.command()
+def download(tag: str | None = None) -> None:
+    """Download platform files for a specific tag.
+
+    Args:
+        tag: Optional image tag to download.
+    """
+    download_platform(tag=tag)
+
+
+@cli.command()
+def upgrade() -> None:
+    """Upgrade the platform to the latest version."""
+    upgrade_platform()
+
+
+@cli.command()
+def refresh_registry_auth() -> None:
+    """Refresh container registry credentials for platform access.
+
+    Used for out of band Docker management.
+    """
+    log_into_registries()
+
+
+@cli.command()
+def configure(service: str = "api") -> None:
+    """Configure the platform for a specific service.
+
+    Args:
+        service: The name of the service to configure.
+    """
+    configure_platform(service=service)

dreadnode/cli/platform/configure.py

@@ -0,0 +1,21 @@
+from dreadnode.cli.platform.utils.env_mgmt import open_env_file
+from dreadnode.cli.platform.utils.printing import print_info
+from dreadnode.cli.platform.utils.versions import get_current_version, get_local_cache_dir
+
+
+def configure_platform(service: str = "api", tag: str | None = None) -> None:
+    """Configure the platform for a specific service.
+
+    Args:
+        service: The name of the service to configure.
+    """
+    if not tag:
+        current_version = get_current_version()
+        tag = current_version.tag if current_version else "latest"
+
+    print_info(f"Configuring {service} service...")
+    env_file = get_local_cache_dir() / tag / f".{service}.env"
+    open_env_file(env_file)
+    print_info(
+        f"Configuration for {service} service loaded. It will take effect the next time the service is started."
+    )

dreadnode/cli/platform/constants.py

@@ -0,0 +1,9 @@
+import typing as t
+
+API_SERVICE = "api"
+UI_SERVICE = "ui"
+SERVICES = [API_SERVICE, UI_SERVICE]
+VERSIONS_MANIFEST = "versions.json"
+
+SupportedArchitecture = t.Literal["amd64", "arm64"]
+SUPPORTED_ARCHITECTURES: list[SupportedArchitecture] = ["amd64", "arm64"]