d10009d639
Open-source Grok free registration CLI, xai_enroller auth pipeline, local auth service, tests and docs.
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
import asyncio
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Protocol
|
|
from urllib.parse import urlencode
|
|
|
|
import httpx
|
|
|
|
from .models import OAuthCredential, SinkReceipt
|
|
|
|
|
|
class SinkError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class CredentialSink(Protocol):
|
|
async def store(self, credential: OAuthCredential) -> SinkReceipt: ...
|
|
|
|
|
|
def cpa_document(credential: OAuthCredential):
|
|
return {
|
|
"type": "xai",
|
|
"access_token": credential.access_token,
|
|
"refresh_token": credential.refresh_token,
|
|
"id_token": credential.id_token,
|
|
"token_type": credential.token_type,
|
|
"expires_in": credential.expires_in,
|
|
"expired": credential.expires_at,
|
|
"last_refresh": credential.last_refresh,
|
|
"sub": credential.subject,
|
|
"base_url": "https://api.x.ai/v1",
|
|
"token_endpoint": credential.token_endpoint,
|
|
"auth_kind": "oauth",
|
|
}
|
|
|
|
|
|
def credential_filename(credential: OAuthCredential, name_secret: bytes):
|
|
subject = credential.subject or credential.refresh_token
|
|
digest = hmac.new(name_secret, subject.encode(), hashlib.sha256).hexdigest()[:16]
|
|
return f"xai-{digest}.json"
|
|
|
|
|
|
class CPAAuthFileSink:
|
|
def __init__(self, base_url, management_secret, client: httpx.AsyncClient, name_secret=None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.management_secret = management_secret
|
|
self.client = client
|
|
self.name_secret = name_secret or management_secret.encode()
|
|
|
|
async def store(self, credential: OAuthCredential):
|
|
filename = credential_filename(credential, self.name_secret)
|
|
document = cpa_document(credential)
|
|
response = await self.client.post(
|
|
f"{self.base_url}/v0/management/auth-files?{urlencode({'name': filename})}",
|
|
headers={
|
|
"Authorization": f"Bearer {self.management_secret}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json=document,
|
|
follow_redirects=False,
|
|
)
|
|
if response.status_code // 100 != 2:
|
|
raise SinkError("CPA upload rejected")
|
|
return SinkReceipt(filename.removesuffix(".json"))
|
|
|
|
|
|
class LocalAuthFileSink:
|
|
"""Atomically persist canonical CPA-compatible documents locally."""
|
|
|
|
def __init__(self, directory, *, name_secret: bytes):
|
|
self.directory = Path(directory).expanduser()
|
|
self.name_secret = name_secret
|
|
|
|
async def store(self, credential: OAuthCredential):
|
|
return await asyncio.to_thread(self._store_sync, credential)
|
|
|
|
def _store_sync(self, credential: OAuthCredential):
|
|
self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
os.chmod(self.directory, 0o700)
|
|
filename = credential_filename(credential, self.name_secret)
|
|
destination = self.directory / filename
|
|
payload = json.dumps(cpa_document(credential), ensure_ascii=False, indent=2) + "\n"
|
|
fd, temporary_name = tempfile.mkstemp(
|
|
prefix=f".{filename}.", suffix=".tmp", dir=self.directory, text=True
|
|
)
|
|
try:
|
|
os.fchmod(fd, 0o600)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
|
stream.write(payload)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(temporary_name, destination)
|
|
os.chmod(destination, 0o600)
|
|
directory_fd = os.open(self.directory, os.O_RDONLY)
|
|
try:
|
|
os.fsync(directory_fd)
|
|
finally:
|
|
os.close(directory_fd)
|
|
finally:
|
|
if os.path.exists(temporary_name):
|
|
os.unlink(temporary_name)
|
|
return SinkReceipt(filename.removesuffix(".json"))
|