329 lines
12 KiB
Python
329 lines
12 KiB
Python
"""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, cancel=None, retries=2):
|
|
request = urllib.request.Request(
|
|
DISCOVERY_URL,
|
|
method="GET",
|
|
headers={"Accept": "application/json", "User-Agent": "grok-register-cpa/1.0"},
|
|
)
|
|
last_error = None
|
|
for attempt in range(max(int(retries), 0) + 1):
|
|
_check_cancel(cancel)
|
|
opener = _build_opener(proxy)
|
|
try:
|
|
with opener.open(request, timeout=float(timeout)) as response:
|
|
body = response.read().decode("utf-8", errors="replace")
|
|
status = int(getattr(response, "status", 200) or 200)
|
|
_check_cancel(cancel)
|
|
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:
|
|
last_error = exc
|
|
if not _is_transient_net_error(exc) or attempt >= int(retries):
|
|
raise OAuthDeviceError("xAI discovery request failed: %s" % exc)
|
|
_sleep_with_cancel(1.0 * (attempt + 1), cancel)
|
|
continue
|
|
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"),
|
|
}
|
|
raise OAuthDeviceError("xAI discovery failed: %s" % last_error)
|
|
|
|
|
|
def _check_cancel(cancel):
|
|
if cancel and cancel():
|
|
raise OAuthDeviceError("cancelled")
|
|
|
|
|
|
def _sleep_with_cancel(seconds, cancel=None):
|
|
deadline = time.time() + max(float(seconds), 0.0)
|
|
while time.time() < deadline:
|
|
_check_cancel(cancel)
|
|
time.sleep(min(0.2, max(deadline - time.time(), 0.0)))
|
|
_check_cancel(cancel)
|
|
|
|
|
|
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, cancel=None):
|
|
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):
|
|
_check_cancel(cancel)
|
|
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)
|
|
_check_cancel(cancel)
|
|
except urllib.error.HTTPError as exc:
|
|
body = exc.read().decode("utf-8", errors="replace")
|
|
status = int(exc.code)
|
|
except Exception as exc:
|
|
last_error = exc
|
|
if not _is_transient_net_error(exc) or attempt >= int(retries):
|
|
raise
|
|
_sleep_with_cancel(float(retry_sleep) * (attempt + 1), cancel)
|
|
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=15.0, proxy=None, cancel=None, retries=2):
|
|
discovery = discover(proxy=proxy, timeout=timeout, cancel=cancel, retries=retries)
|
|
_check_cancel(cancel)
|
|
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=retries,
|
|
retry_sleep=1.0,
|
|
cancel=cancel,
|
|
)
|
|
_check_cancel(cancel)
|
|
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=float(timeout),
|
|
proxy=proxy,
|
|
retries=0,
|
|
retry_sleep=1.0,
|
|
cancel=cancel,
|
|
)
|
|
net_streak = 0
|
|
except Exception 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))
|
|
_sleep_with_cancel(wait_seconds, cancel)
|
|
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))
|
|
_sleep_with_cancel(sleep_seconds, cancel)
|
|
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)
|
|
_sleep_with_cancel(wait_seconds, cancel)
|
|
continue
|
|
logger("oauth poll unexpected HTTP %s: %r" % (status, payload))
|
|
_sleep_with_cancel(sleep_seconds, cancel)
|
|
raise OAuthDeviceError("device auth timed out waiting for user approval")
|