cradle.sdk.auth.m2m

  1import json
  2import time
  3from collections.abc import Generator
  4from datetime import UTC, datetime, timedelta
  5from pathlib import Path
  6
  7import httpx
  8import jwt
  9from pydantic import BaseModel
 10
 11
 12class ClientAuthError(Exception):
 13    def __init__(self, message: str):
 14        self.message = message
 15        super().__init__(f"M2M Authentication Error: {message}")
 16
 17
 18class M2MTokenResponse(BaseModel):
 19    access_token: str
 20    token_type: str = "Bearer"
 21    expires_in: int
 22    refresh_token: str | None = None
 23    scope: str | None = None
 24
 25
 26class M2MAuth(httpx.Auth):
 27    """WorkOS Machine-to-Machine authentication using OAuth2 client credentials flow."""
 28
 29    def __init__(
 30        self,
 31        client_id: str,
 32        client_secret: str,
 33        auth_api_url: str | None = None,
 34        scope: str = "openid profile",
 35    ):
 36        self.client_id = client_id
 37        self.client_secret = client_secret
 38        self.auth_api_url = auth_api_url.rstrip("/") if auth_api_url is not None else None
 39        self.scope = scope
 40        self.token: M2MTokenResponse | None = None
 41        self._token_acquired_at: float | None = None
 42
 43    @staticmethod
 44    def from_json_file(file_path: Path) -> "M2MAuth":
 45        json_data = json.loads(file_path.read_text("utf-8"))
 46        return M2MAuth.from_json(json_data)
 47
 48    @staticmethod
 49    def from_json(json: dict) -> "M2MAuth":
 50        client_id = json.get("client_id")
 51        client_secret = json.get("client_secret")
 52        if not client_id or not client_secret:
 53            raise ClientAuthError("Client ID or client secret is not set")
 54        return M2MAuth(client_id=client_id, client_secret=client_secret)
 55
 56    @property
 57    def access_token(self) -> str | None:
 58        return self.token.access_token if self.token else None
 59
 60    @property
 61    def refresh_token(self) -> str | None:
 62        return self.token.refresh_token if self.token else None
 63
 64    @property
 65    def expires_at(self) -> datetime | None:
 66        if self.token is None or self._token_acquired_at is None:
 67            return None
 68        # Per https://github.com/jpadilla/pyjwt/issues/939 - we should not be verifying the `iat` time here
 69        # as it could be slightly in the future.
 70        jwt_decoded = jwt.decode(self.token.access_token, options={"verify_signature": False, "verify_iat": False})
 71        return datetime.fromtimestamp(jwt_decoded["exp"], tz=UTC)
 72
 73    def _needs_refresh(self) -> bool:
 74        if self.access_token is None or self.expires_at is None:
 75            return True
 76        # Refresh token 30 seconds before expiration to avoid race conditions
 77        return datetime.now(tz=UTC) > self.expires_at - timedelta(seconds=30)
 78
 79    def _get_access_token(self) -> M2MTokenResponse:
 80        """Get access token using client credentials flow."""
 81        if self.auth_api_url is None:
 82            raise ClientAuthError("No auth API URL available")
 83
 84        data = {
 85            "grant_type": "client_credentials",
 86            "client_id": self.client_id,
 87            "client_secret": self.client_secret,
 88            "scope": self.scope,
 89        }
 90
 91        response = httpx.post(
 92            self.auth_api_url + "/oauth2/token",
 93            headers={"Content-Type": "application/x-www-form-urlencoded"},
 94            data=data,
 95        )
 96
 97        if not response.is_success:
 98            error_data = (
 99                response.json() if response.headers.get("content-type", "").startswith("application/json") else {}
100            )
101            error_msg = error_data.get("error_description", f"HTTP {response.status_code}")
102            raise ClientAuthError(f"Failed to obtain access token: {error_msg}")
103
104        token_data = response.json()
105        return M2MTokenResponse.model_validate(token_data)
106
107    def _refresh_access_token(self) -> M2MTokenResponse:
108        """Refresh access token using refresh token if available, otherwise get new token."""
109
110        if self.auth_api_url is None:
111            raise ClientAuthError("No auth API URL available")
112
113        if self.refresh_token:
114            data = {
115                "grant_type": "refresh_token",
116                "client_id": self.client_id,
117                "client_secret": self.client_secret,
118                "refresh_token": self.refresh_token,
119                "scope": self.scope,
120            }
121
122            response = httpx.post(
123                self.auth_api_url + "/oauth2/token",
124                headers={"Content-Type": "application/x-www-form-urlencoded"},
125                data=data,
126            )
127
128            if response.is_success:
129                token_data = response.json()
130                return M2MTokenResponse.model_validate(token_data)
131
132        # If refresh failed or no refresh token, get a new access token
133        return self._get_access_token()
134
135    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
136        """Auth flow implementation for httpx.Auth interface."""
137
138        if self.auth_api_url is None:
139            # Get the base api url from the request, so we can fetch the signin api url from it.
140            base_url = request.url.netloc.decode("utf-8")
141            try:
142                response = yield httpx.Request(
143                    url=f"https://{base_url}/auth:apiStrategy",
144                    method="GET",
145                    headers={
146                        "User-Agent": request.headers.get("User-Agent"),
147                    },
148                )
149                response.read()
150                response.raise_for_status()
151                response_data = response.json()
152                if "auth_api_url" not in response_data:
153                    raise ClientAuthError("Failed to retrieve auth API URL")
154                self.auth_api_url = response_data["auth_api_url"].rstrip("/")
155            except httpx.HTTPError as e:
156                raise ClientAuthError(f"Failed to discover auth API URL: {e}") from e
157        if self._needs_refresh():
158            try:
159                if self.token is None:
160                    self.token = self._get_access_token()
161                else:
162                    self.token = self._refresh_access_token()
163                self._token_acquired_at = time.time()
164            except (httpx.HTTPError, ValueError) as e:
165                raise ClientAuthError(f"Failed to authenticate: {e}") from e
166
167        if self.token is None:
168            raise ClientAuthError("No valid token available")
169
170        request.headers["Authorization"] = f"Bearer m2m_{self.token.access_token}"
171        yield request