diff --git a/cpa_xai/oauth_device.py b/cpa_xai/oauth_device.py new file mode 100644 index 0000000..77f3b05 --- /dev/null +++ b/cpa_xai/oauth_device.py @@ -0,0 +1,300 @@ +"""xAI OAuth device-code flow for Build/CLI credentials.""" + +import json +import time +import urllib.error +import urllib.parse +import urllib.request + +from .proxyutil import resolve_proxy + + +CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +ISSUER = "https://auth.x.ai" +DISCOVERY_URL = ISSUER + "/.well-known/openid-configuration" +SCOPE = "openid profile email offline_access grok-cli:access api:access" + + +class OAuthDeviceError(RuntimeError): + pass + + +class DeviceCodeSession(object): + def __init__( + self, + device_code, + user_code, + verification_uri, + verification_uri_complete, + expires_in, + interval, + token_endpoint, + raw, + ): + self.device_code = device_code + self.user_code = user_code + self.verification_uri = verification_uri + self.verification_uri_complete = verification_uri_complete + self.expires_in = int(expires_in or 1800) + self.interval = int(interval or 5) + self.token_endpoint = token_endpoint + self.raw = raw + + +class TokenResult(object): + def __init__(self, access_token, refresh_token, id_token, token_type, expires_in, raw): + self.access_token = access_token + self.refresh_token = refresh_token + self.id_token = id_token + self.token_type = token_type or "Bearer" + self.expires_in = int(expires_in or 0) + self.raw = raw + + +def _build_opener(proxy=None): + handlers = [] + resolved = resolve_proxy(proxy) + if resolved: + handlers.append(urllib.request.ProxyHandler({"http": resolved, "https": resolved})) + return urllib.request.build_opener(*handlers) if handlers else urllib.request.build_opener() + + +def _validate_endpoint(raw_url, field_name): + value = str(raw_url or "").strip() + if not value: + raise OAuthDeviceError("xAI discovery %s is empty" % field_name) + parsed = urllib.parse.urlparse(value) + if parsed.scheme != "https": + raise OAuthDeviceError("xAI discovery %s must use https: %s" % (field_name, value)) + host = (parsed.hostname or "").lower().strip() + if host != "x.ai" and not host.endswith(".x.ai"): + raise OAuthDeviceError("xAI discovery %s host is invalid: %s" % (field_name, host)) + return value + + +def discover(proxy=None, timeout=30.0): + opener = _build_opener(proxy) + request = urllib.request.Request( + DISCOVERY_URL, + method="GET", + headers={"Accept": "application/json", "User-Agent": "grok-register-cpa/1.0"}, + ) + try: + with opener.open(request, timeout=timeout) as response: + body = response.read().decode("utf-8", errors="replace") + status = int(getattr(response, "status", 200) or 200) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + raise OAuthDeviceError("xAI discovery failed HTTP %s: %s" % (exc.code, body)) + except Exception as exc: + raise OAuthDeviceError("xAI discovery request failed: %s" % exc) + if status != 200: + raise OAuthDeviceError("xAI discovery failed HTTP %s: %s" % (status, body)) + try: + payload = json.loads(body) + except Exception as exc: + raise OAuthDeviceError("xAI discovery parse failed: %s" % exc) + return { + "device_authorization_endpoint": _validate_endpoint( + payload.get("device_authorization_endpoint"), "device_authorization_endpoint" + ), + "token_endpoint": _validate_endpoint(payload.get("token_endpoint"), "token_endpoint"), + } + + +def _is_transient_net_error(exc): + if isinstance(exc, (TimeoutError, BrokenPipeError, ConnectionResetError, ConnectionAbortedError, ConnectionRefusedError)): + return True + if isinstance(exc, urllib.error.URLError): + reason = getattr(exc, "reason", None) + if isinstance(reason, BaseException) and _is_transient_net_error(reason): + return True + text = str(exc).lower() + return any( + needle in text + for needle in ( + "broken pipe", + "connection reset", + "connection aborted", + "timed out", + "timeout", + "temporarily unavailable", + "network is unreachable", + "name or service not known", + "unexpected_eof", + "eof occurred", + "ssl", + "handshake", + "remote end closed", + "bad gateway", + "connection refused", + ) + ) + try: + import ssl as _ssl + + if isinstance(exc, _ssl.SSLError): + return True + except Exception: + pass + if isinstance(exc, OSError): + if getattr(exc, "errno", None) in (32, 104, 110, 111, 113, 101): + return True + text = str(exc).lower() + return any(needle in text for needle in ("broken pipe", "timed out", "connection reset", "ssl")) + return False + + +def _post_form(url, form, timeout=30.0, proxy=None, retries=0, retry_sleep=1.5): + data = urllib.parse.urlencode(form).encode("utf-8") + request = urllib.request.Request( + url, + data=data, + method="POST", + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + "User-Agent": "grok-register-cpa/1.0", + }, + ) + last_error = None + for attempt in range(max(int(retries), 0) + 1): + opener = _build_opener(proxy) + try: + with opener.open(request, timeout=timeout) as response: + body = response.read().decode("utf-8", errors="replace") + status = int(getattr(response, "status", 200) or 200) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + status = int(exc.code) + except BaseException as exc: + last_error = exc + if not _is_transient_net_error(exc) or attempt >= int(retries): + raise + time.sleep(float(retry_sleep) * (attempt + 1)) + continue + try: + return status, json.loads(body) + except Exception: + return status, body + if last_error is not None: + raise last_error + raise OAuthDeviceError("form request failed without response") + + +def request_device_code(client_id=CLIENT_ID, scope=SCOPE, timeout=30.0, proxy=None): + discovery = discover(proxy=proxy, timeout=timeout) + device_endpoint = discovery["device_authorization_endpoint"] + token_endpoint = discovery["token_endpoint"] + status, payload = _post_form( + device_endpoint, + {"client_id": client_id, "scope": scope}, + timeout=timeout, + proxy=proxy, + retries=2, + retry_sleep=1.0, + ) + if status != 200 or not isinstance(payload, dict): + raise OAuthDeviceError("device code request failed HTTP %s: %r" % (status, payload)) + device_code = str(payload.get("device_code") or "").strip() + user_code = str(payload.get("user_code") or "").strip() + if not device_code or not user_code: + raise OAuthDeviceError("device code response missing fields: %r" % payload) + verification_uri = str(payload.get("verification_uri") or "https://accounts.x.ai/oauth2/device").strip() + verification_uri_complete = str( + payload.get("verification_uri_complete") or ("%s?user_code=%s" % (verification_uri, user_code)) + ).strip() + return DeviceCodeSession( + device_code=device_code, + user_code=user_code, + verification_uri=verification_uri, + verification_uri_complete=verification_uri_complete, + expires_in=int(payload.get("expires_in") or 1800), + interval=max(int(payload.get("interval") or 5), 1), + token_endpoint=token_endpoint, + raw=payload, + ) + + +def poll_device_token( + device_code, + token_endpoint, + client_id=CLIENT_ID, + interval=5, + expires_in=1800, + timeout=30.0, + log=None, + cancel=None, + proxy=None, +): + logger = log or (lambda message: None) + deadline = time.time() + max(int(expires_in) - 5, 30) + sleep_seconds = max(int(interval), 1) + net_streak = 0 + max_net_streak = 20 + while time.time() < deadline: + if cancel and cancel(): + raise OAuthDeviceError("cancelled") + try: + status, payload = _post_form( + token_endpoint, + { + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + "device_code": str(device_code).strip(), + "client_id": client_id, + }, + timeout=timeout, + proxy=proxy, + retries=2, + retry_sleep=1.0, + ) + net_streak = 0 + except BaseException as exc: + if not _is_transient_net_error(exc): + raise + net_streak += 1 + wait_seconds = min(sleep_seconds + min(net_streak, 5), 20) + logger("oauth poll network blip (%s/%s): %s — retry in %ss" % (net_streak, max_net_streak, exc, wait_seconds)) + if net_streak >= max_net_streak: + raise OAuthDeviceError("device auth aborted after %s network errors: %s" % (net_streak, exc)) + time.sleep(wait_seconds) + continue + if status == 200 and isinstance(payload, dict) and payload.get("access_token"): + access_token = str(payload.get("access_token") or "").strip() + refresh_token = str(payload.get("refresh_token") or "").strip() + if not refresh_token: + raise OAuthDeviceError("token response missing refresh_token") + return TokenResult( + access_token=access_token, + refresh_token=refresh_token, + id_token=(str(payload.get("id_token") or "").strip() or None), + token_type=str(payload.get("token_type") or "Bearer"), + expires_in=int(payload.get("expires_in") or 21600), + raw=payload, + ) + error_code = "" + error_description = "" + if isinstance(payload, dict): + error_code = str(payload.get("error") or "") + error_description = str(payload.get("error_description") or "") + if error_code in ("authorization_pending", "slow_down"): + if error_code == "slow_down": + sleep_seconds = min(sleep_seconds + 5, 30) + logger("oauth poll: %s (sleep %ss)" % (error_code, sleep_seconds)) + time.sleep(sleep_seconds) + continue + if error_code in ("expired_token", "access_denied"): + raise OAuthDeviceError("device auth failed: %s: %s" % (error_code, error_description)) + if status == 400 and error_code: + raise OAuthDeviceError("device auth token error: %s: %s" % (error_code, error_description or payload)) + if status >= 500 or not isinstance(payload, dict): + net_streak += 1 + wait_seconds = min(sleep_seconds + 2, 20) + logger("oauth poll soft HTTP %s: %r — retry in %ss" % (status, payload, wait_seconds)) + if net_streak >= max_net_streak: + raise OAuthDeviceError("device auth aborted after repeated soft HTTP failures status=%s" % status) + time.sleep(wait_seconds) + continue + logger("oauth poll unexpected HTTP %s: %r" % (status, payload)) + time.sleep(sleep_seconds) + raise OAuthDeviceError("device auth timed out waiting for user approval")