kscale 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
kscale/__init__.py CHANGED
@@ -1,6 +1,6 @@
1
1
  """Defines the common interface for the K-Scale Python API."""
2
2
 
3
- __version__ = "0.1.2"
3
+ __version__ = "0.1.4"
4
4
 
5
5
  from pathlib import Path
6
6
 
kscale/web/api.py CHANGED
@@ -12,3 +12,7 @@ class WebAPI(APIBase):
12
12
  async def get_profile_info(self) -> ProfileResponse:
13
13
  client = await self.www_client()
14
14
  return await client.get_profile_info()
15
+
16
+ async def get_api_key(self, num_hours: int = 24) -> str:
17
+ client = await self.www_client()
18
+ return await client.get_api_key(num_hours)
kscale/web/cli/user.py CHANGED
@@ -3,6 +3,7 @@
3
3
  import logging
4
4
 
5
5
  import click
6
+ from tabulate import tabulate
6
7
 
7
8
  from kscale.utils.cli import coro
8
9
  from kscale.web.clients.user import UserClient
@@ -19,14 +20,33 @@ def cli() -> None:
19
20
  @cli.command()
20
21
  @coro
21
22
  async def me() -> None:
23
+ """Get information about the currently-authenticated user."""
22
24
  client = UserClient()
23
25
  profile = await client.get_profile_info()
24
- logger.info("Email: %s", profile.email)
25
- logger.info("Email verified: %s", profile.email_verified)
26
- logger.info("User ID: %s", profile.user.user_id)
27
- logger.info("Is admin: %s", profile.user.is_admin)
28
- logger.info("Can upload: %s", profile.user.can_upload)
29
- logger.info("Can test: %s", profile.user.can_test)
26
+ click.echo(
27
+ tabulate(
28
+ [
29
+ ["Email", profile.email],
30
+ ["Email verified", profile.email_verified],
31
+ ["User ID", profile.user.user_id],
32
+ ["Is admin", profile.user.is_admin],
33
+ ["Can upload", profile.user.can_upload],
34
+ ["Can test", profile.user.can_test],
35
+ ],
36
+ headers=["Key", "Value"],
37
+ tablefmt="simple",
38
+ )
39
+ )
40
+
41
+
42
+ @cli.command()
43
+ @coro
44
+ async def key() -> None:
45
+ """Get an API key for the currently-authenticated user."""
46
+ client = UserClient()
47
+ api_key = await client.get_api_key()
48
+ click.echo("API key:")
49
+ click.echo(click.style(api_key, fg="green"))
30
50
 
31
51
 
32
52
  if __name__ == "__main__":
@@ -3,6 +3,7 @@
3
3
  import asyncio
4
4
  import json
5
5
  import logging
6
+ import os
6
7
  import secrets
7
8
  import time
8
9
  import webbrowser
@@ -19,24 +20,35 @@ from pydantic import BaseModel
19
20
  from yarl import URL
20
21
 
21
22
  from kscale.web.gen.api import OICDInfo
22
- from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_root, get_cache_dir
23
+ from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_root, get_auth_dir
23
24
 
24
25
  logger = logging.getLogger(__name__)
25
26
 
26
27
  # This port matches the available port for the OAuth callback.
27
28
  OAUTH_PORT = 16821
28
29
 
30
+ # This is the name of the API key header for the K-Scale WWW API.
31
+ HEADER_NAME = "x-kscale-api-key"
32
+
29
33
 
30
34
  class OAuthCallback:
31
35
  def __init__(self) -> None:
36
+ self.token_type: str | None = None
32
37
  self.access_token: str | None = None
38
+ self.id_token: str | None = None
39
+ self.state: str | None = None
40
+ self.expires_in: str | None = None
33
41
  self.app = web.Application()
34
42
  self.app.router.add_get("/token", self.handle_token)
35
43
  self.app.router.add_get("/callback", self.handle_callback)
36
44
 
37
45
  async def handle_token(self, request: web.Request) -> web.Response:
38
46
  """Handle the token extraction."""
47
+ self.token_type = request.query.get("token_type")
39
48
  self.access_token = request.query.get("access_token")
49
+ self.id_token = request.query.get("id_token")
50
+ self.state = request.query.get("state")
51
+ self.expires_in = request.query.get("expires_in")
40
52
  return web.Response(text="OK")
41
53
 
42
54
  async def handle_callback(self, request: web.Request) -> web.Response:
@@ -45,62 +57,92 @@ class OAuthCallback:
45
57
  text="""
46
58
  <!DOCTYPE html>
47
59
  <html lang="en">
48
-
49
60
  <head>
50
61
  <meta charset="UTF-8">
51
- <meta http-equiv="X-UA-Compatible" content="IE=edge">
52
62
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
53
63
  <title>Authentication successful</title>
54
64
  <style>
55
65
  body {
66
+ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
56
67
  display: flex;
57
68
  justify-content: center;
58
69
  align-items: center;
59
70
  min-height: 100vh;
60
71
  margin: 0;
61
- text-align: center;
72
+ background: #f5f5f5;
73
+ color: #333;
74
+ }
75
+ .container {
76
+ background: white;
77
+ padding: 2rem;
78
+ border-radius: 8px;
79
+ box-shadow: 0 2px 4px rgba(0,0,0,0.1);
80
+ max-width: 600px;
81
+ width: 90%;
82
+ }
83
+ h1 {
84
+ color: #2c3e50;
85
+ margin-bottom: 1rem;
62
86
  }
63
- #content {
64
- padding: 20px;
87
+ .token-info {
88
+ background: #f8f9fa;
89
+ border: 1px solid #dee2e6;
90
+ border-radius: 4px;
91
+ padding: 1rem;
92
+ margin: 1rem 0;
93
+ word-break: break-all;
65
94
  }
66
- #closeNotification {
67
- display: none;
68
- padding: 10px 20px;
69
- margin-top: 20px;
70
- cursor: pointer;
71
- margin-left: auto;
72
- margin-right: auto;
95
+ .token-label {
96
+ font-weight: bold;
97
+ color: #6c757d;
98
+ margin-bottom: 0.5rem;
99
+ }
100
+ .success-icon {
101
+ color: #28a745;
102
+ font-size: 48px;
103
+ margin-bottom: 1rem;
73
104
  }
74
105
  </style>
75
106
  </head>
76
-
77
107
  <body>
78
- <div id="content">
108
+ <div class="container">
109
+ <div class="success-icon">✓</div>
79
110
  <h1>Authentication successful!</h1>
80
- <p>This window will close in <span id="countdown">3</span> seconds.</p>
81
- <p id="closeNotification" onclick="window.close()">Please close this window manually.</p>
111
+ <p>Your authentication tokens are shown below. You can now close this window.</p>
112
+
113
+ <div class="token-info">
114
+ <div class="token-label">Access Token:</div>
115
+ <div id="accessTokenDisplay"></div>
116
+ </div>
117
+
118
+ <div class="token-info">
119
+ <div class="token-label">ID Token:</div>
120
+ <div id="idTokenDisplay"></div>
121
+ </div>
82
122
  </div>
123
+
83
124
  <script>
84
125
  const params = new URLSearchParams(window.location.hash.substring(1));
85
- const token = params.get('access_token');
86
- if (token) {
87
- fetch('/token?access_token=' + token);
126
+ const tokenType = params.get('token_type');
127
+ const accessToken = params.get('access_token');
128
+ const idToken = params.get('id_token');
129
+ const state = params.get('state');
130
+ const expiresIn = params.get('expires_in');
131
+
132
+ // Display tokens
133
+ document.getElementById('accessTokenDisplay').textContent = accessToken || 'Not provided';
134
+ document.getElementById('idTokenDisplay').textContent = idToken || 'Not provided';
135
+
136
+ if (accessToken) {
137
+ const tokenUrl = new URL(window.location.href);
138
+ tokenUrl.pathname = '/token';
139
+ tokenUrl.searchParams.set('access_token', accessToken);
140
+ tokenUrl.searchParams.set('token_type', tokenType);
141
+ tokenUrl.searchParams.set('id_token', idToken);
142
+ tokenUrl.searchParams.set('state', state);
143
+ tokenUrl.searchParams.set('expires_in', expiresIn);
144
+ fetch(tokenUrl.toString());
88
145
  }
89
-
90
- let timeLeft = 3;
91
- const countdownElement = document.getElementById('countdown');
92
- const closeNotification = document.getElementById('closeNotification');
93
- const timer = setInterval(() => {
94
- timeLeft--;
95
- countdownElement.textContent = timeLeft;
96
- if (timeLeft <= 0) {
97
- clearInterval(timer);
98
- window.close();
99
- setTimeout(() => {
100
- closeNotification.style.display = 'block';
101
- }, 500);
102
- }
103
- }, 1000);
104
146
  </script>
105
147
  </body>
106
148
  </html>
@@ -124,7 +166,7 @@ class BaseClient:
124
166
 
125
167
  @alru_cache
126
168
  async def _get_oicd_info(self) -> OICDInfo:
127
- cache_path = get_cache_dir() / "oicd_info.json"
169
+ cache_path = get_auth_dir() / "oicd_info.json"
128
170
  if self.use_cache and cache_path.exists():
129
171
  with open(cache_path, "r") as f:
130
172
  return OICDInfo(**json.load(f))
@@ -142,7 +184,7 @@ class BaseClient:
142
184
  Returns:
143
185
  The OpenID Connect server configuration.
144
186
  """
145
- cache_path = get_cache_dir() / "oicd_metadata.json"
187
+ cache_path = get_auth_dir() / "oicd_metadata.json"
146
188
  if self.use_cache and cache_path.exists():
147
189
  with open(cache_path, "r") as f:
148
190
  return json.load(f)
@@ -167,8 +209,23 @@ class BaseClient:
167
209
  oicd_info = await self._get_oicd_info()
168
210
  metadata = await self._get_oicd_metadata()
169
211
  auth_endpoint = metadata["authorization_endpoint"]
170
- state = secrets.token_urlsafe(32)
171
- nonce = secrets.token_urlsafe(32)
212
+
213
+ # Use the cached state and nonce if available, otherwise generate.
214
+ state_file = get_auth_dir() / "oauth_state.json"
215
+ state: str | None = None
216
+ nonce: str | None = None
217
+ if state_file.exists():
218
+ with open(state_file, "r") as f:
219
+ state_data = json.load(f)
220
+ state = state_data.get("state")
221
+ nonce = state_data.get("nonce")
222
+ if state is None:
223
+ state = secrets.token_urlsafe(32)
224
+ if nonce is None:
225
+ nonce = secrets.token_urlsafe(32)
226
+
227
+ # Change /oauth2/authorize to /login to use the login endpoint.
228
+ auth_endpoint = auth_endpoint.replace("/oauth2/authorize", "/login")
172
229
 
173
230
  auth_url = str(
174
231
  URL(auth_endpoint).with_query(
@@ -208,6 +265,11 @@ class BaseClient:
208
265
  raise TimeoutError("Authentication timed out after 30 seconds")
209
266
  await asyncio.sleep(0.1)
210
267
 
268
+ # Save the state and nonce to the cache.
269
+ state = callback_handler.state
270
+ state_file.parent.mkdir(parents=True, exist_ok=True)
271
+ state_file.write_text(json.dumps({"state": state, "nonce": nonce}))
272
+
211
273
  return callback_handler.access_token
212
274
  finally:
213
275
  await runner.cleanup()
@@ -243,7 +305,7 @@ class BaseClient:
243
305
  Returns:
244
306
  A bearer token to use with the K-Scale WWW API.
245
307
  """
246
- cache_path = get_cache_dir() / "bearer_token.txt"
308
+ cache_path = get_auth_dir() / "bearer_token.txt"
247
309
  if self.use_cache and cache_path.exists():
248
310
  token = cache_path.read_text()
249
311
  if not await self._is_token_expired(token):
@@ -257,9 +319,16 @@ class BaseClient:
257
319
  async def get_client(self, *, auth: bool = True) -> httpx.AsyncClient:
258
320
  client = self._client if auth else self._client_no_auth
259
321
  if client is None:
322
+ headers: dict[str, str] = {}
323
+ if auth:
324
+ if "KSCALE_API_KEY" in os.environ:
325
+ headers[HEADER_NAME] = os.environ["KSCALE_API_KEY"]
326
+ else:
327
+ headers["Authorization"] = f"Bearer {await self.get_bearer_token()}"
328
+
260
329
  client = httpx.AsyncClient(
261
330
  base_url=self.base_url,
262
- headers={"Authorization": f"Bearer {await self.get_bearer_token()}"} if auth else None,
331
+ headers=headers,
263
332
  timeout=httpx.Timeout(30.0),
264
333
  )
265
334
  if auth:
@@ -3,13 +3,18 @@
3
3
  import hashlib
4
4
  import json
5
5
  import logging
6
+ import tarfile
6
7
  from pathlib import Path
7
8
 
8
9
  import httpx
9
10
 
10
11
  from kscale.web.clients.base import BaseClient
11
- from kscale.web.gen.api import RobotClass, RobotDownloadURDFResponse, RobotUploadURDFResponse
12
- from kscale.web.utils import get_cache_dir, should_refresh_file
12
+ from kscale.web.gen.api import (
13
+ RobotClass,
14
+ RobotDownloadURDFResponse,
15
+ RobotUploadURDFResponse,
16
+ )
17
+ from kscale.web.utils import get_robots_dir, should_refresh_file
13
18
 
14
19
  logger = logging.getLogger(__name__)
15
20
 
@@ -94,7 +99,7 @@ class RobotClassClient(BaseClient):
94
99
  return response
95
100
 
96
101
  async def download_robot_class_urdf(self, class_name: str, *, cache: bool = True) -> Path:
97
- cache_path = get_cache_dir() / class_name / "robot.tgz"
102
+ cache_path = get_robots_dir() / class_name / "robot.tgz"
98
103
  if cache and cache_path.exists() and not should_refresh_file(cache_path):
99
104
  return cache_path
100
105
  data = await self._request("GET", f"/robots/urdf/{class_name}", auth=True)
@@ -128,7 +133,10 @@ class RobotClassClient(BaseClient):
128
133
  if hash_value_hex != expected_hash:
129
134
  raise ValueError(f"MD5 hash mismatch: {hash_value_hex} != {expected_hash}")
130
135
 
131
- # Updates the info file.
136
+ logger.info("Unpacking downloaded file")
137
+ with tarfile.open(cache_path, "r:gz") as tar:
138
+ tar.extractall(path=cache_path.parent)
139
+
132
140
  logger.info("Updating downloaded file information")
133
141
  info = {"md5_hash": hash_value_hex}
134
142
  with open(cache_path_info, "w") as f:
@@ -8,3 +8,7 @@ class UserClient(BaseClient):
8
8
  async def get_profile_info(self) -> ProfileResponse:
9
9
  data = await self._request("GET", "/auth/profile", auth=True)
10
10
  return ProfileResponse(**data)
11
+
12
+ async def get_api_key(self, num_hours: int = 24) -> str:
13
+ data = await self._request("POST", "/auth/key", auth=True, data={"num_hours": num_hours})
14
+ return data["api_key"]
kscale/web/utils.py CHANGED
@@ -13,24 +13,26 @@ DEFAULT_UPLOAD_TIMEOUT = 300.0 # 5 minutes
13
13
 
14
14
 
15
15
  @functools.lru_cache
16
- def get_cache_dir() -> Path:
16
+ def get_kscale_dir() -> Path:
17
17
  """Returns the cache directory for artifacts."""
18
18
  return Path(Settings.load().www.cache_dir).expanduser().resolve()
19
19
 
20
20
 
21
+ def get_auth_dir() -> Path:
22
+ """Returns the directory for authentication artifacts."""
23
+ return get_kscale_dir() / "auth"
24
+
25
+
26
+ def get_robots_dir() -> Path:
27
+ """Returns the directory for robot artifacts."""
28
+ return get_kscale_dir() / "robots"
29
+
30
+
21
31
  def should_refresh_file(file: Path) -> bool:
22
32
  """Returns whether the file should be refreshed."""
23
33
  return file.exists() and file.stat().st_mtime < time.time() - Settings.load().www.refresh_interval_minutes * 60
24
34
 
25
35
 
26
- @functools.lru_cache
27
- def get_artifact_dir(artifact_id: str) -> Path:
28
- """Returns the directory for a specific artifact."""
29
- cache_dir = get_cache_dir() / artifact_id
30
- cache_dir.mkdir(parents=True, exist_ok=True)
31
- return cache_dir
32
-
33
-
34
36
  @functools.lru_cache
35
37
  def get_api_root() -> str:
36
38
  """Returns the root URL for the K-Scale WWW API."""
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: kscale
3
- Version: 0.1.2
3
+ Version: 0.1.4
4
4
  Summary: The kscale project
5
5
  Home-page: https://github.com/kscalelabs/kscale
6
6
  Author: Benjamin Bolte
@@ -1,4 +1,4 @@
1
- kscale/__init__.py,sha256=Do_oyhXSMNLURbvLDeTx4nJqgYgPhuREBynKDoeJyX4,172
1
+ kscale/__init__.py,sha256=yhQJ7pWCzruLSpCCQHN706gdrMsRyf9UFGe9NRP-QPM,172
2
2
  kscale/api.py,sha256=jmiuFurTN_Gj_-k-6asqxw8wp-_bgJUXgMPFgJ4lqHA,230
3
3
  kscale/cli.py,sha256=PMHLKR5UwdbbReVmqHXpJ-K9-mGHv_0I7KQkwxmFcUA,881
4
4
  kscale/conf.py,sha256=x9jJ8pmSweUsLScN741B21SySCrnRZioV-1ZkhmERbI,1585
@@ -13,24 +13,24 @@ kscale/utils/api_base.py,sha256=Kk_WtRDdJHmOg6NtHmVxVrcfARSUkhfr29ypLch_pO0,112
13
13
  kscale/utils/checksum.py,sha256=jt6QmmQND9zrOEnUtOfZpLYROhgto4Gh3OpdUWk4tZA,1093
14
14
  kscale/utils/cli.py,sha256=JoaY5x5SdUx97KmMM9j5AjRRUqqrTlJ9qVckZptEsYA,827
15
15
  kscale/web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
16
- kscale/web/api.py,sha256=YioIdruq7LCKSBf9SvOGkv914W36_zBmpTzsJqKc0wE,439
17
- kscale/web/utils.py,sha256=l_m8GETCThhnXbHfXq53j1b-E__uZ5GuUbsQ8HR0WA0,1037
16
+ kscale/web/api.py,sha256=oyW0XLfX96RPe1xNgdf8ejfATdLlNlP0CL1lP0FN1nM,593
17
+ kscale/web/utils.py,sha256=kZcOYw5oa1vDe2wNn8ZaHRGhTElIJ8tSq-vVhtY9dzM,1047
18
18
  kscale/web/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
19
19
  kscale/web/cli/robot.py,sha256=rI-A4_0uvJPeA71Apl4Z3mV5fIfWkgmzT9JRmJYxz3A,3307
20
20
  kscale/web/cli/robot_class.py,sha256=ymC5phUqofvOXv5P6f51b9lMK5eRDaavvnzS0x9rDbU,3574
21
21
  kscale/web/cli/token.py,sha256=1rFC8MYKtqbNsQa2KIqwW1tqpaMtFaxuNsallwejXTU,787
22
- kscale/web/cli/user.py,sha256=qO0z2K5uA48hEiOOYEzv6BO2nOlCpITTDZFuiNl6d34,817
22
+ kscale/web/cli/user.py,sha256=aaJJCL1P5lfhK6ZC9OwOHXKA-I3MWqVZ_k7TYnx33CY,1303
23
23
  kscale/web/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
24
- kscale/web/clients/base.py,sha256=NRqRH2JTXjJLUMDM93txKgwFu9bQRYhCYLYy7vE76Ik,11462
24
+ kscale/web/clients/base.py,sha256=eHdHH6uYLLN7lKRU78JiyOxtzflvYwQ9f7oxYlDwmK4,14594
25
25
  kscale/web/clients/client.py,sha256=QjBicdHQYNoUG9XRjAYmGu3THae9DzWa_hQox3OO1Gw,214
26
26
  kscale/web/clients/robot.py,sha256=HMfJnkDxaJ_o7X2vdYYS9iob1JRoBG2qiGmQpCQZpAk,1485
27
- kscale/web/clients/robot_class.py,sha256=vFXlvwQBmt7zEOixAI-mByBEkXCMZIW5OZjiESMavps,5050
28
- kscale/web/clients/user.py,sha256=9iv8J-ROm_yBIwi-0oqldReLkNBFktdHRv3UCOxBzjY,377
27
+ kscale/web/clients/robot_class.py,sha256=O_6lAKAcdNGFVtojuiQlgPCVP5MQlYPZG5Z-awwauho,5206
28
+ kscale/web/clients/user.py,sha256=jsa1_s6qXRM-AGBbHlPhd1NierUtynjY9tVAPNr6_Os,568
29
29
  kscale/web/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
30
30
  kscale/web/gen/api.py,sha256=SovcII36JFgK9jd2CXlLPMjiUROGB4vEnapOsYMUrkU,2188
31
- kscale-0.1.2.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
32
- kscale-0.1.2.dist-info/METADATA,sha256=YEcA0zqZwXvW9BnfZn51PrTXcwemCxMuolC1jB7NG-g,2340
33
- kscale-0.1.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
34
- kscale-0.1.2.dist-info/entry_points.txt,sha256=N_0pCpPnwGDYVzOeuaSOrbJkS5L3lS9d8CxpJF1f8UI,62
35
- kscale-0.1.2.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
36
- kscale-0.1.2.dist-info/RECORD,,
31
+ kscale-0.1.4.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
32
+ kscale-0.1.4.dist-info/METADATA,sha256=BIuP-bpb_ed4i0R1LeQ1UAU5UUst3eU3G4E1FkrrT2k,2340
33
+ kscale-0.1.4.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
34
+ kscale-0.1.4.dist-info/entry_points.txt,sha256=N_0pCpPnwGDYVzOeuaSOrbJkS5L3lS9d8CxpJF1f8UI,62
35
+ kscale-0.1.4.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
36
+ kscale-0.1.4.dist-info/RECORD,,
File without changes