cradle.sdk.auth.device
1import contextlib 2import json 3import logging 4import time 5from datetime import UTC, datetime, timedelta 6from typing import Annotated, Literal, Self 7from urllib.parse import urlencode, urljoin 8 9import httpx 10import jwt 11import keyring 12import keyring.errors 13from keyring.backends.chainer import ChainerBackend 14from pydantic import BaseModel, Field, TypeAdapter, ValidationError 15from typing_extensions import Generator 16 17from cradle.sdk.exceptions import ClientError 18 19logger = logging.getLogger(__name__) 20 21 22class ClientAuthError(Exception): 23 def __init__(self, message: str): 24 self.message = message 25 super().__init__(f"Authentication Error: {message}") 26 27 28class ReauthenticationRequiredError(ClientAuthError): 29 def __init__(self, logout_url): 30 self.logout_url = logout_url 31 super().__init__("Reauthentication required") 32 33 34class DeviceAuthResponse(BaseModel): 35 device_code: str 36 user_code: str 37 verification_uri: str 38 verification_uri_complete: str 39 expires_in: int 40 interval: int 41 42 43class AuthenticateSuccessResponse(BaseModel): 44 access_token: str 45 refresh_token: str 46 organization_id: str | None = None 47 authentication_method: str 48 user: dict 49 50 51class AuthenticateAuthenticationErrorResponse(BaseModel): 52 error: str 53 error_description: str 54 55 56class UnknownErrorResponse(BaseModel): 57 response: dict 58 59 60AuthenticateErrorResponse = AuthenticateAuthenticationErrorResponse | UnknownErrorResponse 61 62 63class UnauthorizedResponse(BaseModel): 64 authorized: Literal[False] = False 65 can_reauth: bool 66 org_id: str | None 67 68 69class AuthorizedResponse(BaseModel): 70 authorized: Literal[True] = True 71 72 73CheckAuthResponse = Annotated[UnauthorizedResponse | AuthorizedResponse, Field(discriminator="authorized")] 74CheckAuthResponseAdapter = TypeAdapter(CheckAuthResponse) 75 76 77class _DeviceAuthStrategy(BaseModel): 78 client_id: str 79 app_url: str 80 auth_api_url: str 81 82 83class DeviceAuth(httpx.Auth): 84 def __init__( 85 self, 86 client_id: str, 87 cradle_app_url: str, 88 cradle_api_base_url: str, 89 auth_base_url: str, 90 workspace: str | None = None, 91 use_keyring: bool = True, 92 user_agent: str | None = None, 93 ): 94 self.client_id = client_id 95 self.cradle_app_url = cradle_app_url 96 self.use_keyring = use_keyring 97 self._keyring_servicename: str = __name__ 98 self.workspace = workspace 99 self.cradle_api_base_url = cradle_api_base_url 100 self.token: AuthenticateSuccessResponse | None = None 101 102 self.checked_for_reauth = False 103 self.force_org_id = None 104 105 self.auth_base_url = auth_base_url 106 self.user_agent = user_agent 107 108 # Try to load token from keyring on initialization 109 if self.use_keyring: 110 self._load_token_from_keyring() 111 112 @classmethod 113 def from_strategy( 114 cls, 115 client: httpx.Client, 116 base_url: str, 117 workspace: str | None, 118 use_keyring: bool = True, 119 ) -> Self: 120 """Create device auth with parameters looked up in the API.""" 121 api_strategy_url = f"{base_url}/auth:apiStrategy" 122 123 try: 124 response = client.get(api_strategy_url) 125 response.raise_for_status() 126 strategy = _DeviceAuthStrategy.model_validate(response.json()) 127 except Exception as exc: 128 raise RuntimeError(f"Failed to retrieve auth strategy from {api_strategy_url}") from exc 129 130 return cls( 131 workspace=workspace, 132 cradle_api_base_url=base_url, 133 client_id=strategy.client_id, 134 cradle_app_url=strategy.app_url, 135 auth_base_url="https://auth.cradle.bio", 136 use_keyring=use_keyring, 137 user_agent=client.headers.get("User-Agent"), 138 ) 139 140 @property 141 def access_token(self) -> str | None: 142 return self.token.access_token if self.token else None 143 144 @property 145 def refresh_token(self) -> str | None: 146 return self.token.refresh_token if self.token else None 147 148 @property 149 def session_id(self) -> str | None: 150 return self.jwt_payload.get("sid") if self.jwt_payload else None 151 152 @property 153 def authorized_workspace_id(self) -> str | None: 154 return self.jwt_payload.get("urn:cradle:workspace_id") if self.jwt_payload else None 155 156 @property 157 def jwks_url(self) -> str: 158 return f"{self.auth_base_url}/sso/jwks/{self.client_id}" 159 160 @property 161 def jwt_payload(self) -> dict | None: 162 if self.token is None: 163 return None 164 165 # Per https://github.com/jpadilla/pyjwt/issues/939 - we should not be verifying the `iat` time here 166 # as it could be slightly in the future. 167 return jwt.decode(self.token.access_token, options={"verify_signature": False, "verify_iat": False}) 168 169 @property 170 def expires_at(self) -> datetime | None: 171 if self.jwt_payload is None: 172 return None 173 return datetime.fromtimestamp(self.jwt_payload["exp"], tz=UTC) 174 175 @property 176 def org_id(self) -> str | None: 177 if self.jwt_payload is None: 178 return None 179 return self.jwt_payload.get("org_id") 180 181 @property 182 def logout_url(self) -> str | None: 183 session_id = self.session_id 184 if session_id is None: 185 return None 186 187 logout_path = "/user_management/sessions/logout" 188 return_to = f"{self.cradle_app_url}/_/sdk/post-logout" 189 query = urlencode({"session_id": session_id, "return_to": return_to}) 190 return urljoin(self.auth_base_url, logout_path) + f"?{query}" 191 192 def _needs_refresh(self) -> bool: 193 if self.access_token is None or self.expires_at is None: 194 return False 195 196 if self.force_org_id is not None and self.org_id != self.force_org_id: 197 return True 198 199 # access_token is valid until expires_at (which I think is currently issue time + 5 min) 200 # We refresh the access_token 10 seconds before expiration to avoid race conditions 201 return datetime.now(tz=UTC) > self.expires_at - timedelta(seconds=10) 202 203 def _save_token_to_keyring(self, token: AuthenticateSuccessResponse) -> None: 204 """Save token data to keyring.""" 205 if not self.use_keyring: 206 return 207 208 # The backends built into keyring are secure. 209 # They will only be in the list of chainer backends if they are "viable" 210 # If this list is empty, it likely means we have no secure keyrings available. 211 has_keyring_backend = any(x.__class__.__module__.startswith("keyring.") for x in ChainerBackend.backends) 212 # The backends built into keyrings.alt are insecure. 213 has_alt_backend = any(x.__class__.__module__.startswith("keyrings.alt") for x in ChainerBackend.backends) 214 215 if not has_keyring_backend and has_alt_backend: 216 logger.warning( 217 "Keyring is likely using a backend from `keyrings.alt`. This is an *insecure* way to save your authentication for a limited period of time." 218 ) 219 220 # Store the token response, don't worry if saving to keyring fails - maybe this system doesn't support it 221 # 222 # Some systems don't support keyring, such as a headless linux machine (because there is no GUI to pop up the 223 # "do you want to allow this application to access your keyring?" dialog). This can be a problem when people run 224 # code on a remote machine. 225 # 226 # In those cases, a user can install `keyrings.alt` into their notebook or environment, and a new backend called 227 # `keyrings.alt.file.PlaintextKeyring` will be available. Because this is insecure, we don't do it by default. 228 # However, all a user needs to do is install `keyrings.alt` into their notebook or environment, and it will automatically 229 # be added to the resolution order of keyring backends. 230 # 231 # If the system does not support any secure keyrings, and no alternate backends are installed, we will get a NoKeyringError. 232 # 233 # If there is a keyring that does not support set_password in the resolution order - such as keyrings.gauth.GooglePythonAuth 234 # aka `keyrings.google-artifactregistry-auth`, we will get a NotImplementedError. If there is another keyring below it in the 235 # order, like PlaintextKeyring that does, then the ChainerBackend will fall through to the PlaintextKeyring. 236 try: 237 keyring.set_password(self._keyring_servicename, self.client_id, token.model_dump_json()) 238 except (NotImplementedError, keyring.errors.NoKeyringError): 239 logger.warning("""No keyring available. You will need to re-authenticate every time you run this code. 240 241To save your authentication for a limited period of time in an *insecure* way, install python package `keyrings.alt` into your notebook or environment and re-run this code. 242""") 243 244 def _load_token_from_keyring(self) -> None: 245 """Load token data from keyring.""" 246 if not self.use_keyring: 247 return 248 249 with contextlib.suppress(keyring.errors.KeyringError): 250 token_json = keyring.get_password(self._keyring_servicename, self.client_id) 251 if token_json is not None: 252 self._update_token(AuthenticateSuccessResponse.model_validate(json.loads(token_json))) 253 254 def _clear_token_from_keyring(self) -> None: 255 """Clear token data from keyring.""" 256 if not self.use_keyring: 257 return 258 259 # Don't worry if deleting from keyring fails - maybe this system doesn't support it 260 with contextlib.suppress(keyring.errors.KeyringError): 261 keyring.delete_password(self._keyring_servicename, self.client_id) 262 263 def _update_token(self, success_response: AuthenticateSuccessResponse) -> None: 264 self.token = success_response 265 # Save to keyring for persistence 266 self._save_token_to_keyring(success_response) 267 268 @staticmethod 269 def parse_response(response: httpx.Response) -> AuthenticateSuccessResponse | AuthenticateErrorResponse: 270 if response.is_success: 271 return AuthenticateSuccessResponse.model_validate(response.json()) 272 try: 273 return AuthenticateAuthenticationErrorResponse.model_validate(response.json()) 274 except ValidationError: 275 return UnknownErrorResponse(response=response.json()) 276 277 def _authenticate( 278 self, data: dict[str, str] 279 ) -> Generator[httpx.Request, httpx.Response, AuthenticateSuccessResponse | AuthenticateErrorResponse]: 280 response = yield httpx.Request( 281 method="POST", 282 url=f"{self.auth_base_url}/user_management/authenticate", 283 headers={"Content-Type": "application/x-www-form-urlencoded", "User-Agent": self.user_agent or "unknown"}, 284 data={**data, "client_id": self.client_id}, 285 ) 286 response.read() 287 288 return self.parse_response(response) # noqa: B901 289 290 def _refresh_token(self) -> Generator[httpx.Request, httpx.Response, None]: 291 if self.refresh_token is None: 292 raise ClientAuthError("No refresh token available") 293 294 refresh_data = { 295 "grant_type": "refresh_token", 296 "refresh_token": self.refresh_token, 297 } 298 if self.force_org_id is not None: 299 refresh_data["organization_id"] = self.force_org_id 300 301 response = yield from self._authenticate(refresh_data) 302 303 match response: 304 case AuthenticateAuthenticationErrorResponse(error="invalid_grant"): 305 # This might mean the refresh token has expired, so let's clear keyring and re-run the device flow 306 self._clear_token_from_keyring() 307 self.token = None 308 yield from self._start_device_flow() 309 case AuthenticateAuthenticationErrorResponse(error="sso_required"): 310 # This means the token cannot be refreshed for a different org that requires SSO, 311 # so let's clear keyring and re-run the device flow 312 self._clear_token_from_keyring() 313 self.token = None 314 yield from self._start_device_flow() 315 case AuthenticateAuthenticationErrorResponse(error="mfa_enrollment"): 316 # It might not be possible to refresh a token to a different organization if the new 317 # organization has a different security policy in place, e.g. regarding MFA. 318 # In this case we need to log out and re-authenticate to the new workspace. 319 # Unfortunately this requires the user to manually click a logout link. 320 logout_url = self.logout_url 321 if logout_url is None: 322 raise ClientAuthError("No session ID available to log out.") 323 324 print( 325 "Access to the requested workspace requires re-authentication. " 326 "Please log out under the following link and then retry the request.\n\n" 327 f"{logout_url}" 328 ) 329 330 raise ReauthenticationRequiredError(logout_url) 331 case AuthenticateAuthenticationErrorResponse(): 332 raise ClientAuthError(response.error_description) 333 case UnknownErrorResponse(): 334 raise ClientAuthError(json.dumps(response.response)) 335 case AuthenticateSuccessResponse(): 336 self._update_token(response) 337 338 def _poll_for_tokens( 339 self, device_code: str, expires_in: int = 300, interval: int = 5 340 ) -> Generator[httpx.Request, httpx.Response, AuthenticateSuccessResponse]: 341 """Poll for authentication tokens using device code flow. 342 343 Args: 344 device_code: The device code from the initial auth request 345 expires_in: Timeout in seconds (default 300) 346 interval: Polling interval in seconds (default 5) 347 348 Returns: 349 AuthenticateSuccessResponse: access and refresh tokens 350 351 Raises: 352 Exception: If polling for access/refresh tokens fails or times out 353 """ 354 start_time = time.monotonic() 355 356 while True: 357 # Check timeout 358 if time.monotonic() - start_time > expires_in: 359 raise ClientAuthError("Authentication timed out") 360 361 try: 362 refresh_data = { 363 "device_code": device_code, 364 "grant_type": "urn:ietf:params:oauth:grant-type:device_code", 365 } 366 if self.force_org_id is not None: 367 refresh_data["organization_id"] = self.force_org_id 368 response = yield from self._authenticate(refresh_data) 369 370 except httpx.TimeoutException as e: 371 raise ClientAuthError("Authentication timed out") from e 372 373 match response: 374 case AuthenticateSuccessResponse(): 375 return response # noqa: B901 376 case AuthenticateAuthenticationErrorResponse(error="authorization_pending"): 377 time.sleep(interval) 378 case AuthenticateAuthenticationErrorResponse(error="slow_down"): 379 interval += 1 380 case AuthenticateAuthenticationErrorResponse(): 381 raise ClientAuthError(response.error_description) 382 case UnknownErrorResponse(): 383 raise ClientAuthError(json.dumps(response.response)) 384 385 def _start_device_flow(self) -> Generator[httpx.Request, httpx.Response, None]: 386 device_auth_http_response: httpx.Response = yield httpx.Request( 387 method="POST", 388 url=f"{self.auth_base_url}/user_management/authorize/device", 389 json={ 390 "client_id": self.client_id, 391 }, 392 headers={"user-agent": self.user_agent or "unknown"}, 393 ) 394 device_auth_http_response.read() 395 device_auth_http_response.raise_for_status() 396 device_auth_response = DeviceAuthResponse.model_validate(device_auth_http_response.json()) 397 398 print(f""" 399Please click the following URL to complete authentication 400 401{device_auth_response.verification_uri_complete} 402 403and verify that the code matches: {device_auth_response.user_code} 404""") 405 406 # Poll for tokens using the device code 407 token_data = yield from self._poll_for_tokens( 408 device_code=device_auth_response.device_code, 409 expires_in=device_auth_response.expires_in, 410 interval=device_auth_response.interval, 411 ) 412 self._update_token(token_data) 413 414 def _check_auth(self) -> Generator[httpx.Request, httpx.Response, None | CheckAuthResponse]: 415 if self.cradle_api_base_url is None or self.workspace is None or self.access_token is None: 416 return None 417 418 needs_reauth_response: httpx.Response = yield httpx.Request( 419 method="GET", 420 url=f"{self.cradle_api_base_url}/auth:checkAuth", 421 params={"workspace": self.workspace}, 422 headers={"Authorization": f"Bearer ca_{self.access_token}", "User-Agent": self.user_agent or "unknown"}, 423 ) 424 425 needs_reauth_response.read() 426 needs_reauth_response.raise_for_status() 427 return CheckAuthResponseAdapter.validate_python(needs_reauth_response.json()) # noqa: B901 428 429 def auth_flow(self, request): 430 if not self.access_token: 431 yield from self._start_device_flow() 432 433 # If we loaded the token from keyring, the access_token is almost certainly expired, 434 # so we need to refresh it *before* hitting the checkAuth endpoint, which is authenticated 435 # otherwise we get a 403. 436 if self._needs_refresh(): 437 yield from self._refresh_token() 438 439 # Only check for reauth once - the requested workspace that this client is constructed with 440 # will not change, so we won't need to check if we need to reauth to a different workspace/org 441 # more than once 442 if not self.checked_for_reauth: 443 check_auth_response = yield from self._check_auth() 444 # can_reauth=True means that we are a member of the workspace we want to access, but need to reauthorize 445 # our access_token to the correct org 446 if isinstance(check_auth_response, UnauthorizedResponse) and check_auth_response.can_reauth: 447 # This will force _refresh_token to request access to the correct org 448 self.force_org_id = check_auth_response.org_id 449 450 self.checked_for_reauth = True 451 452 # After checking if we need to reauth, this will fire if we need to re-auth the token to a different org 453 # from the one it was issued for. 454 if self._needs_refresh(): 455 yield from self._refresh_token() 456 457 request.headers["Authorization"] = f"Bearer ca_{self.access_token}" 458 yield request 459 460 def suggest_logout(self, workspace: str) -> None: # noqa: ARG002 - might use workspace in the future 461 """Print logout URL.""" 462 logout_url = self.logout_url 463 if logout_url is None: 464 raise ClientAuthError("No session ID available to log out.") 465 466 # Don't clear token in case this is caused by fetching a workspace immediately after 467 # creating it. 468 # 469 # Don't clear keyring in case the user decides to recreate the client with a different 470 # workspace name. 471 472 # Reset this check in case we just created a workspace and now need to auth to it 473 self.checked_for_reauth = False 474 475 message = ( 476 f'Authorized workspace "{self.authorized_workspace_id}" does not match requested workspace "{self.workspace}"\n' 477 f"If you believe this is an error, you can log out of this session at {logout_url} and then try the request again." 478 ) 479 raise ClientError(403, message, [])
logger =
<Logger cradle.sdk.auth.device (WARNING)>
class
ClientAuthError(builtins.Exception):
23class ClientAuthError(Exception): 24 def __init__(self, message: str): 25 self.message = message 26 super().__init__(f"Authentication Error: {message}")
Common base class for all non-exit exceptions.
29class ReauthenticationRequiredError(ClientAuthError): 30 def __init__(self, logout_url): 31 self.logout_url = logout_url 32 super().__init__("Reauthentication required")
Common base class for all non-exit exceptions.
Inherited Members
class
DeviceAuthResponse(pydantic.main.BaseModel):
class
AuthenticateSuccessResponse(pydantic.main.BaseModel):
class
AuthenticateAuthenticationErrorResponse(pydantic.main.BaseModel):
class
UnknownErrorResponse(pydantic.main.BaseModel):
AuthenticateErrorResponse =
AuthenticateAuthenticationErrorResponse | UnknownErrorResponse
class
AuthorizedResponse(pydantic.main.BaseModel):
CheckAuthResponse =
typing.Annotated[UnauthorizedResponse | AuthorizedResponse, FieldInfo(annotation=NoneType, required=True, discriminator='authorized')]
CheckAuthResponseAdapter =
TypeAdapter(Annotated[Union[UnauthorizedResponse, AuthorizedResponse], FieldInfo(annotation=NoneType, required=True, discriminator='authorized')])
class
DeviceAuth(httpx.Auth):
84class DeviceAuth(httpx.Auth): 85 def __init__( 86 self, 87 client_id: str, 88 cradle_app_url: str, 89 cradle_api_base_url: str, 90 auth_base_url: str, 91 workspace: str | None = None, 92 use_keyring: bool = True, 93 user_agent: str | None = None, 94 ): 95 self.client_id = client_id 96 self.cradle_app_url = cradle_app_url 97 self.use_keyring = use_keyring 98 self._keyring_servicename: str = __name__ 99 self.workspace = workspace 100 self.cradle_api_base_url = cradle_api_base_url 101 self.token: AuthenticateSuccessResponse | None = None 102 103 self.checked_for_reauth = False 104 self.force_org_id = None 105 106 self.auth_base_url = auth_base_url 107 self.user_agent = user_agent 108 109 # Try to load token from keyring on initialization 110 if self.use_keyring: 111 self._load_token_from_keyring() 112 113 @classmethod 114 def from_strategy( 115 cls, 116 client: httpx.Client, 117 base_url: str, 118 workspace: str | None, 119 use_keyring: bool = True, 120 ) -> Self: 121 """Create device auth with parameters looked up in the API.""" 122 api_strategy_url = f"{base_url}/auth:apiStrategy" 123 124 try: 125 response = client.get(api_strategy_url) 126 response.raise_for_status() 127 strategy = _DeviceAuthStrategy.model_validate(response.json()) 128 except Exception as exc: 129 raise RuntimeError(f"Failed to retrieve auth strategy from {api_strategy_url}") from exc 130 131 return cls( 132 workspace=workspace, 133 cradle_api_base_url=base_url, 134 client_id=strategy.client_id, 135 cradle_app_url=strategy.app_url, 136 auth_base_url="https://auth.cradle.bio", 137 use_keyring=use_keyring, 138 user_agent=client.headers.get("User-Agent"), 139 ) 140 141 @property 142 def access_token(self) -> str | None: 143 return self.token.access_token if self.token else None 144 145 @property 146 def refresh_token(self) -> str | None: 147 return self.token.refresh_token if self.token else None 148 149 @property 150 def session_id(self) -> str | None: 151 return self.jwt_payload.get("sid") if self.jwt_payload else None 152 153 @property 154 def authorized_workspace_id(self) -> str | None: 155 return self.jwt_payload.get("urn:cradle:workspace_id") if self.jwt_payload else None 156 157 @property 158 def jwks_url(self) -> str: 159 return f"{self.auth_base_url}/sso/jwks/{self.client_id}" 160 161 @property 162 def jwt_payload(self) -> dict | None: 163 if self.token is None: 164 return None 165 166 # Per https://github.com/jpadilla/pyjwt/issues/939 - we should not be verifying the `iat` time here 167 # as it could be slightly in the future. 168 return jwt.decode(self.token.access_token, options={"verify_signature": False, "verify_iat": False}) 169 170 @property 171 def expires_at(self) -> datetime | None: 172 if self.jwt_payload is None: 173 return None 174 return datetime.fromtimestamp(self.jwt_payload["exp"], tz=UTC) 175 176 @property 177 def org_id(self) -> str | None: 178 if self.jwt_payload is None: 179 return None 180 return self.jwt_payload.get("org_id") 181 182 @property 183 def logout_url(self) -> str | None: 184 session_id = self.session_id 185 if session_id is None: 186 return None 187 188 logout_path = "/user_management/sessions/logout" 189 return_to = f"{self.cradle_app_url}/_/sdk/post-logout" 190 query = urlencode({"session_id": session_id, "return_to": return_to}) 191 return urljoin(self.auth_base_url, logout_path) + f"?{query}" 192 193 def _needs_refresh(self) -> bool: 194 if self.access_token is None or self.expires_at is None: 195 return False 196 197 if self.force_org_id is not None and self.org_id != self.force_org_id: 198 return True 199 200 # access_token is valid until expires_at (which I think is currently issue time + 5 min) 201 # We refresh the access_token 10 seconds before expiration to avoid race conditions 202 return datetime.now(tz=UTC) > self.expires_at - timedelta(seconds=10) 203 204 def _save_token_to_keyring(self, token: AuthenticateSuccessResponse) -> None: 205 """Save token data to keyring.""" 206 if not self.use_keyring: 207 return 208 209 # The backends built into keyring are secure. 210 # They will only be in the list of chainer backends if they are "viable" 211 # If this list is empty, it likely means we have no secure keyrings available. 212 has_keyring_backend = any(x.__class__.__module__.startswith("keyring.") for x in ChainerBackend.backends) 213 # The backends built into keyrings.alt are insecure. 214 has_alt_backend = any(x.__class__.__module__.startswith("keyrings.alt") for x in ChainerBackend.backends) 215 216 if not has_keyring_backend and has_alt_backend: 217 logger.warning( 218 "Keyring is likely using a backend from `keyrings.alt`. This is an *insecure* way to save your authentication for a limited period of time." 219 ) 220 221 # Store the token response, don't worry if saving to keyring fails - maybe this system doesn't support it 222 # 223 # Some systems don't support keyring, such as a headless linux machine (because there is no GUI to pop up the 224 # "do you want to allow this application to access your keyring?" dialog). This can be a problem when people run 225 # code on a remote machine. 226 # 227 # In those cases, a user can install `keyrings.alt` into their notebook or environment, and a new backend called 228 # `keyrings.alt.file.PlaintextKeyring` will be available. Because this is insecure, we don't do it by default. 229 # However, all a user needs to do is install `keyrings.alt` into their notebook or environment, and it will automatically 230 # be added to the resolution order of keyring backends. 231 # 232 # If the system does not support any secure keyrings, and no alternate backends are installed, we will get a NoKeyringError. 233 # 234 # If there is a keyring that does not support set_password in the resolution order - such as keyrings.gauth.GooglePythonAuth 235 # aka `keyrings.google-artifactregistry-auth`, we will get a NotImplementedError. If there is another keyring below it in the 236 # order, like PlaintextKeyring that does, then the ChainerBackend will fall through to the PlaintextKeyring. 237 try: 238 keyring.set_password(self._keyring_servicename, self.client_id, token.model_dump_json()) 239 except (NotImplementedError, keyring.errors.NoKeyringError): 240 logger.warning("""No keyring available. You will need to re-authenticate every time you run this code. 241 242To save your authentication for a limited period of time in an *insecure* way, install python package `keyrings.alt` into your notebook or environment and re-run this code. 243""") 244 245 def _load_token_from_keyring(self) -> None: 246 """Load token data from keyring.""" 247 if not self.use_keyring: 248 return 249 250 with contextlib.suppress(keyring.errors.KeyringError): 251 token_json = keyring.get_password(self._keyring_servicename, self.client_id) 252 if token_json is not None: 253 self._update_token(AuthenticateSuccessResponse.model_validate(json.loads(token_json))) 254 255 def _clear_token_from_keyring(self) -> None: 256 """Clear token data from keyring.""" 257 if not self.use_keyring: 258 return 259 260 # Don't worry if deleting from keyring fails - maybe this system doesn't support it 261 with contextlib.suppress(keyring.errors.KeyringError): 262 keyring.delete_password(self._keyring_servicename, self.client_id) 263 264 def _update_token(self, success_response: AuthenticateSuccessResponse) -> None: 265 self.token = success_response 266 # Save to keyring for persistence 267 self._save_token_to_keyring(success_response) 268 269 @staticmethod 270 def parse_response(response: httpx.Response) -> AuthenticateSuccessResponse | AuthenticateErrorResponse: 271 if response.is_success: 272 return AuthenticateSuccessResponse.model_validate(response.json()) 273 try: 274 return AuthenticateAuthenticationErrorResponse.model_validate(response.json()) 275 except ValidationError: 276 return UnknownErrorResponse(response=response.json()) 277 278 def _authenticate( 279 self, data: dict[str, str] 280 ) -> Generator[httpx.Request, httpx.Response, AuthenticateSuccessResponse | AuthenticateErrorResponse]: 281 response = yield httpx.Request( 282 method="POST", 283 url=f"{self.auth_base_url}/user_management/authenticate", 284 headers={"Content-Type": "application/x-www-form-urlencoded", "User-Agent": self.user_agent or "unknown"}, 285 data={**data, "client_id": self.client_id}, 286 ) 287 response.read() 288 289 return self.parse_response(response) # noqa: B901 290 291 def _refresh_token(self) -> Generator[httpx.Request, httpx.Response, None]: 292 if self.refresh_token is None: 293 raise ClientAuthError("No refresh token available") 294 295 refresh_data = { 296 "grant_type": "refresh_token", 297 "refresh_token": self.refresh_token, 298 } 299 if self.force_org_id is not None: 300 refresh_data["organization_id"] = self.force_org_id 301 302 response = yield from self._authenticate(refresh_data) 303 304 match response: 305 case AuthenticateAuthenticationErrorResponse(error="invalid_grant"): 306 # This might mean the refresh token has expired, so let's clear keyring and re-run the device flow 307 self._clear_token_from_keyring() 308 self.token = None 309 yield from self._start_device_flow() 310 case AuthenticateAuthenticationErrorResponse(error="sso_required"): 311 # This means the token cannot be refreshed for a different org that requires SSO, 312 # so let's clear keyring and re-run the device flow 313 self._clear_token_from_keyring() 314 self.token = None 315 yield from self._start_device_flow() 316 case AuthenticateAuthenticationErrorResponse(error="mfa_enrollment"): 317 # It might not be possible to refresh a token to a different organization if the new 318 # organization has a different security policy in place, e.g. regarding MFA. 319 # In this case we need to log out and re-authenticate to the new workspace. 320 # Unfortunately this requires the user to manually click a logout link. 321 logout_url = self.logout_url 322 if logout_url is None: 323 raise ClientAuthError("No session ID available to log out.") 324 325 print( 326 "Access to the requested workspace requires re-authentication. " 327 "Please log out under the following link and then retry the request.\n\n" 328 f"{logout_url}" 329 ) 330 331 raise ReauthenticationRequiredError(logout_url) 332 case AuthenticateAuthenticationErrorResponse(): 333 raise ClientAuthError(response.error_description) 334 case UnknownErrorResponse(): 335 raise ClientAuthError(json.dumps(response.response)) 336 case AuthenticateSuccessResponse(): 337 self._update_token(response) 338 339 def _poll_for_tokens( 340 self, device_code: str, expires_in: int = 300, interval: int = 5 341 ) -> Generator[httpx.Request, httpx.Response, AuthenticateSuccessResponse]: 342 """Poll for authentication tokens using device code flow. 343 344 Args: 345 device_code: The device code from the initial auth request 346 expires_in: Timeout in seconds (default 300) 347 interval: Polling interval in seconds (default 5) 348 349 Returns: 350 AuthenticateSuccessResponse: access and refresh tokens 351 352 Raises: 353 Exception: If polling for access/refresh tokens fails or times out 354 """ 355 start_time = time.monotonic() 356 357 while True: 358 # Check timeout 359 if time.monotonic() - start_time > expires_in: 360 raise ClientAuthError("Authentication timed out") 361 362 try: 363 refresh_data = { 364 "device_code": device_code, 365 "grant_type": "urn:ietf:params:oauth:grant-type:device_code", 366 } 367 if self.force_org_id is not None: 368 refresh_data["organization_id"] = self.force_org_id 369 response = yield from self._authenticate(refresh_data) 370 371 except httpx.TimeoutException as e: 372 raise ClientAuthError("Authentication timed out") from e 373 374 match response: 375 case AuthenticateSuccessResponse(): 376 return response # noqa: B901 377 case AuthenticateAuthenticationErrorResponse(error="authorization_pending"): 378 time.sleep(interval) 379 case AuthenticateAuthenticationErrorResponse(error="slow_down"): 380 interval += 1 381 case AuthenticateAuthenticationErrorResponse(): 382 raise ClientAuthError(response.error_description) 383 case UnknownErrorResponse(): 384 raise ClientAuthError(json.dumps(response.response)) 385 386 def _start_device_flow(self) -> Generator[httpx.Request, httpx.Response, None]: 387 device_auth_http_response: httpx.Response = yield httpx.Request( 388 method="POST", 389 url=f"{self.auth_base_url}/user_management/authorize/device", 390 json={ 391 "client_id": self.client_id, 392 }, 393 headers={"user-agent": self.user_agent or "unknown"}, 394 ) 395 device_auth_http_response.read() 396 device_auth_http_response.raise_for_status() 397 device_auth_response = DeviceAuthResponse.model_validate(device_auth_http_response.json()) 398 399 print(f""" 400Please click the following URL to complete authentication 401 402{device_auth_response.verification_uri_complete} 403 404and verify that the code matches: {device_auth_response.user_code} 405""") 406 407 # Poll for tokens using the device code 408 token_data = yield from self._poll_for_tokens( 409 device_code=device_auth_response.device_code, 410 expires_in=device_auth_response.expires_in, 411 interval=device_auth_response.interval, 412 ) 413 self._update_token(token_data) 414 415 def _check_auth(self) -> Generator[httpx.Request, httpx.Response, None | CheckAuthResponse]: 416 if self.cradle_api_base_url is None or self.workspace is None or self.access_token is None: 417 return None 418 419 needs_reauth_response: httpx.Response = yield httpx.Request( 420 method="GET", 421 url=f"{self.cradle_api_base_url}/auth:checkAuth", 422 params={"workspace": self.workspace}, 423 headers={"Authorization": f"Bearer ca_{self.access_token}", "User-Agent": self.user_agent or "unknown"}, 424 ) 425 426 needs_reauth_response.read() 427 needs_reauth_response.raise_for_status() 428 return CheckAuthResponseAdapter.validate_python(needs_reauth_response.json()) # noqa: B901 429 430 def auth_flow(self, request): 431 if not self.access_token: 432 yield from self._start_device_flow() 433 434 # If we loaded the token from keyring, the access_token is almost certainly expired, 435 # so we need to refresh it *before* hitting the checkAuth endpoint, which is authenticated 436 # otherwise we get a 403. 437 if self._needs_refresh(): 438 yield from self._refresh_token() 439 440 # Only check for reauth once - the requested workspace that this client is constructed with 441 # will not change, so we won't need to check if we need to reauth to a different workspace/org 442 # more than once 443 if not self.checked_for_reauth: 444 check_auth_response = yield from self._check_auth() 445 # can_reauth=True means that we are a member of the workspace we want to access, but need to reauthorize 446 # our access_token to the correct org 447 if isinstance(check_auth_response, UnauthorizedResponse) and check_auth_response.can_reauth: 448 # This will force _refresh_token to request access to the correct org 449 self.force_org_id = check_auth_response.org_id 450 451 self.checked_for_reauth = True 452 453 # After checking if we need to reauth, this will fire if we need to re-auth the token to a different org 454 # from the one it was issued for. 455 if self._needs_refresh(): 456 yield from self._refresh_token() 457 458 request.headers["Authorization"] = f"Bearer ca_{self.access_token}" 459 yield request 460 461 def suggest_logout(self, workspace: str) -> None: # noqa: ARG002 - might use workspace in the future 462 """Print logout URL.""" 463 logout_url = self.logout_url 464 if logout_url is None: 465 raise ClientAuthError("No session ID available to log out.") 466 467 # Don't clear token in case this is caused by fetching a workspace immediately after 468 # creating it. 469 # 470 # Don't clear keyring in case the user decides to recreate the client with a different 471 # workspace name. 472 473 # Reset this check in case we just created a workspace and now need to auth to it 474 self.checked_for_reauth = False 475 476 message = ( 477 f'Authorized workspace "{self.authorized_workspace_id}" does not match requested workspace "{self.workspace}"\n' 478 f"If you believe this is an error, you can log out of this session at {logout_url} and then try the request again." 479 ) 480 raise ClientError(403, message, [])
Base class for all authentication schemes.
To implement a custom authentication scheme, subclass Auth and override
the .auth_flow() method.
If the authentication scheme does I/O such as disk access or network calls, or uses
synchronization primitives such as locks, you should override .sync_auth_flow()
and/or .async_auth_flow() instead of .auth_flow() to provide specialized
implementations that will be used by Client and AsyncClient respectively.
DeviceAuth( client_id: str, cradle_app_url: str, cradle_api_base_url: str, auth_base_url: str, workspace: str | None = None, use_keyring: bool = True, user_agent: str | None = None)
85 def __init__( 86 self, 87 client_id: str, 88 cradle_app_url: str, 89 cradle_api_base_url: str, 90 auth_base_url: str, 91 workspace: str | None = None, 92 use_keyring: bool = True, 93 user_agent: str | None = None, 94 ): 95 self.client_id = client_id 96 self.cradle_app_url = cradle_app_url 97 self.use_keyring = use_keyring 98 self._keyring_servicename: str = __name__ 99 self.workspace = workspace 100 self.cradle_api_base_url = cradle_api_base_url 101 self.token: AuthenticateSuccessResponse | None = None 102 103 self.checked_for_reauth = False 104 self.force_org_id = None 105 106 self.auth_base_url = auth_base_url 107 self.user_agent = user_agent 108 109 # Try to load token from keyring on initialization 110 if self.use_keyring: 111 self._load_token_from_keyring()
token: AuthenticateSuccessResponse | None
@classmethod
def
from_strategy( cls, client: httpx.Client, base_url: str, workspace: str | None, use_keyring: bool = True) -> Self:
113 @classmethod 114 def from_strategy( 115 cls, 116 client: httpx.Client, 117 base_url: str, 118 workspace: str | None, 119 use_keyring: bool = True, 120 ) -> Self: 121 """Create device auth with parameters looked up in the API.""" 122 api_strategy_url = f"{base_url}/auth:apiStrategy" 123 124 try: 125 response = client.get(api_strategy_url) 126 response.raise_for_status() 127 strategy = _DeviceAuthStrategy.model_validate(response.json()) 128 except Exception as exc: 129 raise RuntimeError(f"Failed to retrieve auth strategy from {api_strategy_url}") from exc 130 131 return cls( 132 workspace=workspace, 133 cradle_api_base_url=base_url, 134 client_id=strategy.client_id, 135 cradle_app_url=strategy.app_url, 136 auth_base_url="https://auth.cradle.bio", 137 use_keyring=use_keyring, 138 user_agent=client.headers.get("User-Agent"), 139 )
Create device auth with parameters looked up in the API.
jwt_payload: dict | None
161 @property 162 def jwt_payload(self) -> dict | None: 163 if self.token is None: 164 return None 165 166 # Per https://github.com/jpadilla/pyjwt/issues/939 - we should not be verifying the `iat` time here 167 # as it could be slightly in the future. 168 return jwt.decode(self.token.access_token, options={"verify_signature": False, "verify_iat": False})
logout_url: str | None
182 @property 183 def logout_url(self) -> str | None: 184 session_id = self.session_id 185 if session_id is None: 186 return None 187 188 logout_path = "/user_management/sessions/logout" 189 return_to = f"{self.cradle_app_url}/_/sdk/post-logout" 190 query = urlencode({"session_id": session_id, "return_to": return_to}) 191 return urljoin(self.auth_base_url, logout_path) + f"?{query}"
@staticmethod
def
parse_response( response: httpx.Response) -> AuthenticateSuccessResponse | AuthenticateAuthenticationErrorResponse | UnknownErrorResponse:
269 @staticmethod 270 def parse_response(response: httpx.Response) -> AuthenticateSuccessResponse | AuthenticateErrorResponse: 271 if response.is_success: 272 return AuthenticateSuccessResponse.model_validate(response.json()) 273 try: 274 return AuthenticateAuthenticationErrorResponse.model_validate(response.json()) 275 except ValidationError: 276 return UnknownErrorResponse(response=response.json())
def
auth_flow(self, request):
430 def auth_flow(self, request): 431 if not self.access_token: 432 yield from self._start_device_flow() 433 434 # If we loaded the token from keyring, the access_token is almost certainly expired, 435 # so we need to refresh it *before* hitting the checkAuth endpoint, which is authenticated 436 # otherwise we get a 403. 437 if self._needs_refresh(): 438 yield from self._refresh_token() 439 440 # Only check for reauth once - the requested workspace that this client is constructed with 441 # will not change, so we won't need to check if we need to reauth to a different workspace/org 442 # more than once 443 if not self.checked_for_reauth: 444 check_auth_response = yield from self._check_auth() 445 # can_reauth=True means that we are a member of the workspace we want to access, but need to reauthorize 446 # our access_token to the correct org 447 if isinstance(check_auth_response, UnauthorizedResponse) and check_auth_response.can_reauth: 448 # This will force _refresh_token to request access to the correct org 449 self.force_org_id = check_auth_response.org_id 450 451 self.checked_for_reauth = True 452 453 # After checking if we need to reauth, this will fire if we need to re-auth the token to a different org 454 # from the one it was issued for. 455 if self._needs_refresh(): 456 yield from self._refresh_token() 457 458 request.headers["Authorization"] = f"Bearer ca_{self.access_token}" 459 yield request
Execute the authentication flow.
To dispatch a request, yield it:
yield request
The client will .send() the response back into the flow generator. You can
access it like so:
response = yield request
A return (or reaching the end of the generator) will result in the
client returning the last response obtained from the server.
You can dispatch as many requests as is necessary.
def
suggest_logout(self, workspace: str) -> None:
461 def suggest_logout(self, workspace: str) -> None: # noqa: ARG002 - might use workspace in the future 462 """Print logout URL.""" 463 logout_url = self.logout_url 464 if logout_url is None: 465 raise ClientAuthError("No session ID available to log out.") 466 467 # Don't clear token in case this is caused by fetching a workspace immediately after 468 # creating it. 469 # 470 # Don't clear keyring in case the user decides to recreate the client with a different 471 # workspace name. 472 473 # Reset this check in case we just created a workspace and now need to auth to it 474 self.checked_for_reauth = False 475 476 message = ( 477 f'Authorized workspace "{self.authorized_workspace_id}" does not match requested workspace "{self.workspace}"\n' 478 f"If you believe this is an error, you can log out of this session at {logout_url} and then try the request again." 479 ) 480 raise ClientError(403, message, [])
Print logout URL.