Initial commit: grok-free-register-oss

Open-source Grok free registration CLI, xai_enroller auth pipeline,
local auth service, tests and docs.

xai_enroller/sinks.py

@@ -0,0 +1,106 @@
+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"))