PyPI - kscale - Versions diffs - 0.0.11__py3-none-any.whl → 0.1.0__py3-none-any.whl - Mend

kscale 0.0.11py3-none-any.whl → 0.1.0py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (38) hide show

kscale/__init__.py +2 -2
kscale/api.py +3 -6
kscale/cli.py +32 -0
kscale/conf.py +6 -3
kscale/requirements.txt +17 -1
kscale/utils/checksum.py +41 -0
kscale/utils/cli.py +28 -0
kscale/web/api.py +14 -0
kscale/web/cli/robot.py +100 -0
kscale/web/cli/robot_class.py +113 -0
kscale/web/cli/token.py +33 -0
kscale/web/cli/user.py +33 -0
kscale/web/clients/__init__.py +0 -0
kscale/web/clients/base.py +314 -0
kscale/web/clients/client.py +11 -0
kscale/web/clients/robot.py +39 -0
kscale/web/clients/robot_class.py +114 -0
kscale/web/clients/user.py +10 -0
kscale/web/gen/__init__.py +0 -0
kscale/web/gen/api.py +73 -0
kscale/web/utils.py +31 -0
{kscale-0.0.11.dist-info → kscale-0.1.0.dist-info}/METADATA +29 -48
kscale-0.1.0.dist-info/RECORD +36 -0
{kscale-0.0.11.dist-info → kscale-0.1.0.dist-info}/WHEEL +1 -1
kscale-0.1.0.dist-info/entry_points.txt +3 -0
kscale/store/api.py +0 -64
kscale/store/cli.py +0 -35
kscale/store/client.py +0 -82
kscale/store/gen/api.py +0 -397
kscale/store/pybullet.py +0 -180
kscale/store/urdf.py +0 -193
kscale/store/utils.py +0 -33
kscale-0.0.11.dist-info/RECORD +0 -26
kscale-0.0.11.dist-info/entry_points.txt +0 -2
/kscale/{store → web}/__init__.py +0 -0
/kscale/{store/gen → web/cli}/__init__.py +0 -0
{kscale-0.0.11.dist-info → kscale-0.1.0.dist-info}/LICENSE +0 -0
{kscale-0.0.11.dist-info → kscale-0.1.0.dist-info}/top_level.txt +0 -0

kscale/web/clients/base.py ADDED Viewed

@@ -0,0 +1,314 @@
+"""Defines a base client for the K-Scale WWW API client."""
+import asyncio
+import json
+import logging
+import secrets
+import time
+import webbrowser
+from types import TracebackType
+from typing import Any, Self, Type
+from urllib.parse import urljoin
+import aiohttp
+import httpx
+from aiohttp import web
+from async_lru import alru_cache
+from jwt import ExpiredSignatureError, PyJWKClient, decode as jwt_decode
+from pydantic import BaseModel
+from yarl import URL
+from kscale.web.gen.api import OICDInfo
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_root, get_cache_dir
+logger = logging.getLogger(__name__)
+# This port matches the available port for the OAuth callback.
+OAUTH_PORT = 16821
+class OAuthCallback:
+    def __init__(self) -> None:
+        self.access_token: str | None = None
+        self.app = web.Application()
+        self.app.router.add_get("/token", self.handle_token)
+        self.app.router.add_get("/callback", self.handle_callback)
+    async def handle_token(self, request: web.Request) -> web.Response:
+        """Handle the token extraction."""
+        self.access_token = request.query.get("access_token")
+        return web.Response(text="OK")
+    async def handle_callback(self, request: web.Request) -> web.Response:
+        """Handle the OAuth callback with token in URL fragment."""
+        return web.Response(
+            text="""
+                <!DOCTYPE html>
+                <html lang="en">
+                <head>
+                    <meta charset="UTF-8">
+                    <meta http-equiv="X-UA-Compatible" content="IE=edge">
+                    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+                    <title>Authentication successful</title>
+                    <style>
+                        body {
+                            display: flex;
+                            justify-content: center;
+                            align-items: center;
+                            min-height: 100vh;
+                            margin: 0;
+                            text-align: center;
+                        }
+                        #content {
+                            padding: 20px;
+                        }
+                        #closeNotification {
+                            display: none;
+                            padding: 10px 20px;
+                            margin-top: 20px;
+                            cursor: pointer;
+                            margin-left: auto;
+                            margin-right: auto;
+                        }
+                    </style>
+                </head>
+                <body>
+                    <div id="content">
+                        <h1>Authentication successful!</h1>
+                        <p>This window will close in <span id="countdown">3</span> seconds.</p>
+                        <p id="closeNotification" onclick="window.close()">Please close this window manually.</p>
+                    </div>
+                    <script>
+                        const params = new URLSearchParams(window.location.hash.substring(1));
+                        const token = params.get('access_token');
+                        if (token) {
+                            fetch('/token?access_token=' + token);
+                        }
+                        let timeLeft = 3;
+                        const countdownElement = document.getElementById('countdown');
+                        const closeNotification = document.getElementById('closeNotification');
+                        const timer = setInterval(() => {
+                            timeLeft--;
+                            countdownElement.textContent = timeLeft;
+                            if (timeLeft <= 0) {
+                                clearInterval(timer);
+                                window.close();
+                                setTimeout(() => {
+                                    closeNotification.style.display = 'block';
+                                }, 500);
+                            }
+                        }, 1000);
+                    </script>
+                </body>
+                </html>
+            """,
+            content_type="text/html",
+        )
+class BaseClient:
+    def __init__(
+        self,
+        base_url: str | None = None,
+        upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT,
+        use_cache: bool = True,
+    ) -> None:
+        self.base_url = get_api_root() if base_url is None else base_url
+        self.upload_timeout = upload_timeout
+        self.use_cache = use_cache
+        self._client: httpx.AsyncClient | None = None
+        self._client_no_auth: httpx.AsyncClient | None = None
+    @alru_cache
+    async def _get_oicd_info(self) -> OICDInfo:
+        cache_path = get_cache_dir() / "oicd_info.json"
+        if self.use_cache and cache_path.exists():
+            with open(cache_path, "r") as f:
+                return OICDInfo(**json.load(f))
+        data = await self._request("GET", "/auth/oicd", auth=False)
+        if self.use_cache:
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(cache_path, "w") as f:
+                json.dump(data, f)
+        return OICDInfo(**data)
+    @alru_cache
+    async def _get_oicd_metadata(self) -> dict:
+        """Returns the OpenID Connect server configuration.
+        Returns:
+            The OpenID Connect server configuration.
+        """
+        cache_path = get_cache_dir() / "oicd_metadata.json"
+        if self.use_cache and cache_path.exists():
+            with open(cache_path, "r") as f:
+                return json.load(f)
+        oicd_info = await self._get_oicd_info()
+        oicd_config_url = f"{oicd_info.authority}/.well-known/openid-configuration"
+        async with aiohttp.ClientSession() as session:
+            async with session.get(oicd_config_url) as response:
+                metadata = await response.json()
+        if self.use_cache:
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(cache_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+            logger.info("Cached OpenID Connect metadata to %s", cache_path)
+        return metadata
+    async def _get_bearer_token(self) -> str:
+        """Get a bearer token using the OAuth2 implicit flow.
+        Returns:
+            A bearer token to use with the K-Scale WWW API.
+        """
+        oicd_info = await self._get_oicd_info()
+        metadata = await self._get_oicd_metadata()
+        auth_endpoint = metadata["authorization_endpoint"]
+        state = secrets.token_urlsafe(32)
+        nonce = secrets.token_urlsafe(32)
+        auth_url = str(
+            URL(auth_endpoint).with_query(
+                {
+                    "response_type": "token",
+                    "redirect_uri": f"http://localhost:{OAUTH_PORT}/callback",
+                    "state": state,
+                    "nonce": nonce,
+                    "scope": "openid profile email",
+                    "client_id": oicd_info.client_id,
+                }
+            )
+        )
+        # Start local server to receive callback
+        callback_handler = OAuthCallback()
+        runner = web.AppRunner(callback_handler.app)
+        await runner.setup()
+        site = web.TCPSite(runner, "localhost", OAUTH_PORT)
+        try:
+            await site.start()
+        except OSError as e:
+            raise OSError(
+                f"The command line interface requires access to local port {OAUTH_PORT} in order to authenticate with "
+                "OpenID Connect. Please ensure that no other application is using this port."
+            ) from e
+        # Open browser for user authentication
+        webbrowser.open(auth_url)
+        # Wait for the callback with timeout
+        try:
+            start_time = time.time()
+            while callback_handler.access_token is None:
+                if time.time() - start_time > 30:
+                    raise TimeoutError("Authentication timed out after 30 seconds")
+                await asyncio.sleep(0.1)
+            return callback_handler.access_token
+        finally:
+            await runner.cleanup()
+    @alru_cache
+    async def _get_jwk_client(self) -> PyJWKClient:
+        """Returns a JWK client for the OpenID Connect server."""
+        oicd_info = await self._get_oicd_info()
+        jwks_uri = f"{oicd_info.authority}/.well-known/jwks.json"
+        return PyJWKClient(uri=jwks_uri)
+    async def _is_token_expired(self, token: str) -> bool:
+        """Check if a token is expired."""
+        jwk_client = await self._get_jwk_client()
+        signing_key = jwk_client.get_signing_key_from_jwt(token)
+        try:
+            claims = jwt_decode(
+                token,
+                signing_key.key,
+                algorithms=["RS256"],
+                options={"verify_aud": False},
+            )
+        except ExpiredSignatureError:
+            return True
+        return claims["exp"] < time.time()
+    @alru_cache
+    async def get_bearer_token(self) -> str:
+        """Get a bearer token from OpenID Connect.
+        Returns:
+            A bearer token to use with the K-Scale WWW API.
+        """
+        cache_path = get_cache_dir() / "bearer_token.txt"
+        if self.use_cache and cache_path.exists():
+            token = cache_path.read_text()
+            if not await self._is_token_expired(token):
+                return token
+        token = await self._get_bearer_token()
+        if self.use_cache:
+            cache_path.write_text(token)
+            cache_path.chmod(0o600)
+        return token
+    async def get_client(self, *, auth: bool = True) -> httpx.AsyncClient:
+        client = self._client if auth else self._client_no_auth
+        if client is None:
+            client = httpx.AsyncClient(
+                base_url=self.base_url,
+                headers={"Authorization": f"Bearer {await self.get_bearer_token()}"} if auth else None,
+                timeout=httpx.Timeout(30.0),
+            )
+            if auth:
+                self._client = client
+            else:
+                self._client_no_auth = client
+        return client
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        *,
+        auth: bool = True,
+        params: dict[str, Any] | None = None,
+        data: BaseModel | dict[str, Any] | None = None,
+        files: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        url = urljoin(self.base_url, endpoint)
+        kwargs: dict[str, Any] = {"params": params}
+        if data:
+            if isinstance(data, BaseModel):
+                kwargs["json"] = data.model_dump(exclude_unset=True)
+            else:
+                kwargs["json"] = data
+        if files:
+            kwargs["files"] = files
+        client = await self.get_client(auth=auth)
+        response = await client.request(method, url, **kwargs)
+        if response.is_error:
+            logger.error("Error response from K-Scale: %s", response.text)
+        response.raise_for_status()
+        return response.json()
+    async def close(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None
+    async def __aenter__(self) -> Self:
+        return self
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()

kscale/web/clients/client.py ADDED Viewed

@@ -0,0 +1,11 @@
+"""Defines a unified client for the K-Scale WWW API."""
+from kscale.web.clients.base import BaseClient
+from kscale.web.clients.user import UserClient
+class WWWClient(
+    UserClient,
+    BaseClient,
+):
+    pass

kscale/web/clients/robot.py ADDED Viewed

@@ -0,0 +1,39 @@
+"""Defines the client for interacting with the K-Scale robot endpoints."""
+from kscale.web.clients.base import BaseClient
+from kscale.web.gen.api import Robot, RobotResponse
+class RobotClient(BaseClient):
+    async def get_all_robots(self) -> list[Robot]:
+        data = await self._request("GET", "/robot/", auth=True)
+        return [Robot.model_validate(item) for item in data]
+    async def get_user_robots(self, user_id: str = "me") -> list[Robot]:
+        data = await self._request("GET", f"/robot/user/{user_id}", auth=True)
+        return [Robot.model_validate(item) for item in data]
+    async def add_robot(
+        self,
+        robot_name: str,
+        class_name: str,
+        description: str | None = None,
+    ) -> RobotResponse:
+        params = {"class_name": class_name}
+        if description is not None:
+            params["description"] = description
+        data = await self._request(
+            "PUT",
+            f"/robot/{robot_name}",
+            params=params,
+            auth=True,
+        )
+        return RobotResponse.model_validate(data)
+    async def get_robot_by_id(self, robot_id: str) -> RobotResponse:
+        data = await self._request("GET", f"/robot/id/{robot_id}", auth=True)
+        return RobotResponse.model_validate(data)
+    async def get_robot_by_name(self, robot_name: str) -> RobotResponse:
+        data = await self._request("GET", f"/robot/name/{robot_name}", auth=True)
+        return RobotResponse.model_validate(data)

kscale/web/clients/robot_class.py ADDED Viewed

@@ -0,0 +1,114 @@
+"""Defines the client for interacting with the K-Scale robot class endpoints."""
+import hashlib
+import logging
+from pathlib import Path
+import httpx
+from kscale.web.clients.base import BaseClient
+from kscale.web.gen.api import RobotClass, RobotDownloadURDFResponse, RobotUploadURDFResponse
+from kscale.web.utils import get_cache_dir
+logger = logging.getLogger(__name__)
+class RobotClassClient(BaseClient):
+    async def get_robot_classes(self) -> list[RobotClass]:
+        data = await self._request(
+            "GET",
+            "/robots/",
+            auth=True,
+        )
+        return [RobotClass.model_validate(item) for item in data]
+    async def create_robot_class(self, class_name: str, description: str | None = None) -> RobotClass:
+        params = {}
+        if description is not None:
+            params["description"] = description
+        data = await self._request(
+            "PUT",
+            f"/robots/{class_name}",
+            params=params,
+            auth=True,
+        )
+        return RobotClass.model_validate(data)
+    async def update_robot_class(
+        self,
+        class_name: str,
+        new_class_name: str | None = None,
+        new_description: str | None = None,
+    ) -> RobotClass:
+        params = {}
+        if new_class_name is not None:
+            params["new_class_name"] = new_class_name
+        if new_description is not None:
+            params["new_description"] = new_description
+        if not params:
+            raise ValueError("No parameters to update")
+        data = await self._request(
+            "POST",
+            f"/robots/{class_name}",
+            params=params,
+            auth=True,
+        )
+        return RobotClass.model_validate(data)
+    async def delete_robot_class(self, class_name: str) -> None:
+        await self._request("DELETE", f"/robots/{class_name}", auth=True)
+    async def upload_robot_class_urdf(self, class_name: str, urdf_file: str | Path) -> RobotUploadURDFResponse:
+        if not (urdf_file := Path(urdf_file)).exists():
+            raise FileNotFoundError(f"URDF file not found: {urdf_file}")
+        # Gets the content type from the file extension.
+        ext = urdf_file.suffix.lower()
+        match ext:
+            case ".tgz":
+                content_type = "application/x-compressed-tar"
+            case _:
+                raise ValueError(f"Unsupported file extension: {ext}")
+        data = await self._request(
+            "PUT",
+            f"/robots/urdf/{class_name}",
+            params={"filename": urdf_file.name, "content_type": content_type},
+            auth=True,
+        )
+        response = RobotUploadURDFResponse.model_validate(data)
+        async with httpx.AsyncClient() as client:
+            async with client.stream(
+                "PUT",
+                response.url,
+                content=urdf_file.read_bytes(),
+                headers={"Content-Type": response.content_type},
+            ) as r:
+                r.raise_for_status()
+        return response
+    async def download_robot_class_urdf(self, class_name: str, *, cache: bool = True) -> Path:
+        cache_path = get_cache_dir() / class_name / "robot.tgz"
+        if cache and cache_path.exists():
+            return cache_path
+        data = await self._request("GET", f"/robots/urdf/{class_name}", auth=True)
+        response = RobotDownloadURDFResponse.model_validate(data)
+        expected_hash = response.md5_hash
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.info("Downloading URDF file from %s", response.url)
+        async with httpx.AsyncClient() as client:
+            with open(cache_path, "wb") as file:
+                hash_value = hashlib.md5()
+                async with client.stream("GET", response.url) as r:
+                    r.raise_for_status()
+                    async for chunk in r.aiter_bytes():
+                        file.write(chunk)
+                        hash_value.update(chunk)
+        logger.info("Checking MD5 hash of downloaded file")
+        hash_value_hex = f'"{hash_value.hexdigest()}"'
+        if hash_value_hex != expected_hash:
+            raise ValueError(f"MD5 hash mismatch: {hash_value_hex} != {expected_hash}")
+        return cache_path

kscale/web/clients/user.py ADDED Viewed

@@ -0,0 +1,10 @@
+"""Defines the client for interacting with the K-Scale authentication endpoints."""
+from kscale.web.clients.base import BaseClient
+from kscale.web.gen.api import ProfileResponse
+class UserClient(BaseClient):
+    async def get_profile_info(self) -> ProfileResponse:
+        data = await self._request("GET", "/auth/profile", auth=True)
+        return ProfileResponse(**data)

kscale/web/gen/__init__.py ADDED Viewed

File without changes

kscale/web/gen/api.py ADDED Viewed

@@ -0,0 +1,73 @@
+"""Auto-generated by generate.sh script."""
+# generated by datamodel-codegen:
+#   filename:  openapi.json
+#   timestamp: 2025-01-15T22:35:42+00:00
+from __future__ import annotations
+from typing import List, Optional, Union
+from pydantic import BaseModel, Field
+class OICDInfo(BaseModel):
+    authority: str = Field(..., title="Authority")
+    client_id: str = Field(..., title="Client Id")
+class Robot(BaseModel):
+    id: str = Field(..., title="Id")
+    robot_name: str = Field(..., title="Robot Name")
+    description: str = Field(..., title="Description")
+    user_id: str = Field(..., title="User Id")
+    class_id: str = Field(..., title="Class Id")
+class RobotClass(BaseModel):
+    id: str = Field(..., title="Id")
+    class_name: str = Field(..., title="Class Name")
+    description: str = Field(..., title="Description")
+    user_id: str = Field(..., title="User Id")
+class RobotDownloadURDFResponse(BaseModel):
+    url: str = Field(..., title="Url")
+    md5_hash: str = Field(..., title="Md5 Hash")
+class RobotResponse(BaseModel):
+    id: str = Field(..., title="Id")
+    robot_name: str = Field(..., title="Robot Name")
+    description: str = Field(..., title="Description")
+    user_id: str = Field(..., title="User Id")
+    class_name: str = Field(..., title="Class Name")
+class RobotUploadURDFResponse(BaseModel):
+    url: str = Field(..., title="Url")
+    filename: str = Field(..., title="Filename")
+    content_type: str = Field(..., title="Content Type")
+class UserResponse(BaseModel):
+    user_id: str = Field(..., title="User Id")
+    is_admin: bool = Field(..., title="Is Admin")
+    can_upload: bool = Field(..., title="Can Upload")
+    can_test: bool = Field(..., title="Can Test")
+class ValidationError(BaseModel):
+    loc: List[Union[str, int]] = Field(..., title="Location")
+    msg: str = Field(..., title="Message")
+    type: str = Field(..., title="Error Type")
+class HTTPValidationError(BaseModel):
+    detail: Optional[List[ValidationError]] = Field(None, title="Detail")
+class ProfileResponse(BaseModel):
+    email: str = Field(..., title="Email")
+    email_verified: bool = Field(..., title="Email Verified")
+    user: UserResponse

kscale/web/utils.py ADDED Viewed

@@ -0,0 +1,31 @@
+"""Utility functions for interacting with the K-Scale WWW API."""
+import functools
+import logging
+from pathlib import Path
+from kscale.conf import Settings
+logger = logging.getLogger(__name__)
+DEFAULT_UPLOAD_TIMEOUT = 300.0  # 5 minutes
+@functools.lru_cache
+def get_cache_dir() -> Path:
+    """Returns the cache directory for artifacts."""
+    return Path(Settings.load().www.cache_dir).expanduser().resolve()
+@functools.lru_cache
+def get_artifact_dir(artifact_id: str) -> Path:
+    """Returns the directory for a specific artifact."""
+    cache_dir = get_cache_dir() / artifact_id
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+@functools.lru_cache
+def get_api_root() -> str:
+    """Returns the root URL for the K-Scale WWW API."""
+    return Settings.load().www.api_root

kscale 0.0.11__py3-none-any.whl → 0.1.0__py3-none-any.whl

kscale 0.0.11py3-none-any.whl → 0.1.0py3-none-any.whl