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

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