oidc token endpoint (#25)

* oidc token endpoint

* format

* ci stuff

* format

* format

* lint

* rust

.github/workflows/publish.yml

@@ -127,47 +127,3 @@ jobs:
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           packages-dir: final_dist/
-
-  publish-rust:
-    name: Build and publish Rust package
-    timeout-minutes: 10
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-
-      - name: Set up Rust
-        uses: actions-rs/toolchain@v1
-        with:
-          toolchain: stable
-
-      - name: Install dependencies
-        run: |
-          sudo apt-get update
-          sudo apt-get install -y libudev-dev pkg-config
-
-      - name: Install protoc
-        uses: arduino/setup-protoc@v3
-
-      - name: Cache Cargo registry
-        uses: actions/cache@v2
-        with:
-          path: ~/.cargo/registry
-          key: ${{ runner.os }}-cargo-registry
-          restore-keys: |
-            ${{ runner.os }}-cargo-registry
-
-      - name: Cache Cargo index
-        uses: actions/cache@v2
-        with:
-          path: ~/.cargo/git
-          key: ${{ runner.os }}-cargo-index
-          restore-keys: |
-            ${{ runner.os }}-cargo-index
-
-      - name: Publish K-Rec package to crates.io
-        env:
-          CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
-        run: |
-          cargo publish -p kscale

.github/workflows/test.yml

@@ -52,7 +52,7 @@ jobs:
 
       - name: Install package
         run: |
-          pip install --upgrade --upgrade-strategy eager -e '.'
+          pip install --upgrade --upgrade-strategy eager -e '.[dev]'
 
       - name: Run Python linting
         run: |
@@ -62,14 +62,6 @@ jobs:
         run: |
           make test
 
-      - name: Check Rust lint
-        run: |
-          cargo fmt --check
-
-      - name: Run Rust tests
-        run: |
-          cargo test
-
       - name: Save cache
         uses: actions/cache/save@v3
         if: github.ref == 'refs/heads/master'

Makefile

@@ -9,6 +9,7 @@ py-files := $(shell find . -type f -name '*.py' ! -path "./.venv/*")
 format:
 	@black $(py-files)
 	@ruff format $(py-files)
+	@isort $(py-files)
 .PHONY: format
 
 static-checks:

README.md

@@ -1,9 +1,3 @@
-<p align="center">
-  <picture>
-    <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
-  </picture>
-</p>
-
 <div align="center">
 
 [![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)

kscale/cli.py

@@ -6,6 +6,7 @@
 from kscale.web.kernels import cli as kernel_images_cli
 from kscale.web.krec import cli as krec_cli
 from kscale.web.pybullet import cli as pybullet_cli
+from kscale.web.token import cli as token_cli
 from kscale.web.urdf import cli as urdf_cli
 
 
@@ -15,6 +16,7 @@ def cli() -> None:
     pass
 
 
+cli.add_command(token_cli, "token")
 cli.add_command(urdf_cli, "urdf")
 cli.add_command(pybullet_cli, "pybullet")
 cli.add_command(kernel_images_cli, "kernel")

kscale/requirements.txt

@@ -5,9 +5,11 @@ omegaconf
 email_validator
 
 # HTTP requests
+aiohttp
 httpx
 requests
 pydantic
+yarl
 
 # CLI
 click

kscale/web/kernels.py

@@ -12,7 +12,12 @@
 from kscale.utils.checksum import FileChecksum
 from kscale.utils.cli import coro
 from kscale.web.gen.api import SingleArtifactResponse
-from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_artifact_dir, get_cache_dir
+from kscale.web.utils import (
+    DEFAULT_UPLOAD_TIMEOUT,
+    get_api_key,
+    get_artifact_dir,
+    get_cache_dir,
+)
 from kscale.web.www_client import KScaleWWWClient
 
 httpx_logger = logging.getLogger("httpx")

kscale/web/token.py

@@ -0,0 +1,142 @@
+"""Command-line interface example for getting a bearer token from OpenID Connect."""
+
+import asyncio
+import logging
+import secrets
+import time
+import webbrowser
+
+import aiohttp
+import click
+from aiohttp import web
+from yarl import URL
+
+from kscale.utils.cli import coro
+
+logger = logging.getLogger(__name__)
+
+SERVER_METADATA_URL = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_dqtJl1Iew/.well-known/openid-configuration"
+
+
+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>
+                </head>
+
+                <body>
+                    <h1>Authentication successful!</h1>
+                    <p>This window will close in <span id="countdown">3</span> seconds.</p>
+                    <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 timer = setInterval(() => {
+                            timeLeft--;
+                            countdownElement.textContent = timeLeft;
+                            if (timeLeft <= 0) {
+                                clearInterval(timer);
+                                window.close();
+                            }
+                        }, 1000);
+                    </script>
+                </body>
+
+                </html>
+            """,
+            content_type="text/html",
+        )
+
+
+async def get_bearer_token(redirect_uri: str = "http://localhost:8080/callback") -> str:
+    """Get a bearer token using the OAuth2 implicit flow."""
+    async with aiohttp.ClientSession() as session:
+        # Get the OpenID Connect metadata
+        async with session.get(SERVER_METADATA_URL) as response:
+            metadata = await response.json()
+
+        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": redirect_uri,
+                    "state": state,
+                    "nonce": nonce,
+                    "scope": "openid profile email",
+                    "client_id": "5lu9h7nhtf6dvlunpodjr9qil5",
+                }
+            )
+        )
+
+        # Start local server to receive callback
+        callback_handler = OAuthCallback()
+        runner = web.AppRunner(callback_handler.app)
+        await runner.setup()
+        site = web.TCPSite(runner, "localhost", 8080)
+        await site.start()
+
+        # 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()
+
+
+@click.group()
+def cli() -> None:
+    """K-Scale OpenID Connect CLI tool."""
+    pass
+
+
+@cli.command()
+@coro
+async def get() -> None:
+    """Get a bearer token from OpenID Connect."""
+    logging.getLogger("aiohttp.access").setLevel(logging.WARNING)
+    try:
+        token = await get_bearer_token()
+        logger.info("Bearer token: %s", token)
+    except Exception:
+        logger.exception("Error getting bearer token")
+
+
+if __name__ == "__main__":
+    cli()