from __future__ import annotations import json import os import shutil from collections.abc import Iterator from pathlib import Path from account_registry import Account, AccountRef, AccountRegistry, StagedAccount from grok_oauth import ( GrokOAuthSession, poll_device_code, start_device_code, write_auth_file, ) from provider_services import LifecycleEvent PROTECTED_SLUG = "roy-grok" class GrokAuthOperationError(RuntimeError): pass class GrokAuthOperation: def __init__(self, registry: AccountRegistry) -> None: self.registry = registry def add( self, alias: str, *, login_hint: str | None = None, slug: str | None = None, ) -> Iterator[LifecycleEvent]: del login_hint staged: StagedAccount | None = None pending: Account | None = None try: staged = self.registry.stage_add(alias, slug=slug) pending = self._pending_account(staged) except Exception as exc: yield LifecycleEvent( kind="failure", message=f"Adding Grok account {alias} failed: {exc}", error=exc, ) return yield LifecycleEvent(kind="started", account=pending) try: device = start_device_code() uri = str(device.get("verification_uri_complete") or device.get("verification_uri") or "") user_code = str(device.get("user_code") or "") yield LifecycleEvent( kind="prompt_ready", account=pending, message=f"Grok device login: open {uri} code {user_code}", prompt=_Prompt(url=uri, raw_text=f"{uri}\n{user_code}"), session=None, ) tokens = poll_device_code(device) access = str(tokens.get("access_token") or "") refresh = str(tokens.get("refresh_token") or "") if not access or not refresh: raise GrokAuthOperationError("device login missing tokens") auth_path = staged.account_home / "auth.json" write_auth_file( auth_path, access=access, refresh=refresh, expires_in=float(tokens["expires_in"]) if isinstance(tokens.get("expires_in"), (int, float)) else None, ) identity = self._write_identity(staged.account_home, auth_path) collision = self._find_collision(identity.get("account_id"), identity.get("email")) if collision is not None: staged.rollback() yield LifecycleEvent( kind="collision", account=pending, collision=collision, message=f"Grok identity already registered as {collision.alias}", ) return account = staged.commit() yield LifecycleEvent(kind="success", account=account) except Exception as exc: if staged is not None and not staged._closed: staged.rollback() yield LifecycleEvent( kind="failure", account=pending, message=f"Adding Grok account failed: {exc}", error=exc, ) def reauthenticate(self, account: Account) -> Iterator[LifecycleEvent]: if account.slug == PROTECTED_SLUG and os.environ.get("SYSTRAY_ALLOW_ROY_GROK_MUTATION") != "1": yield LifecycleEvent( kind="failure", account=account, message="roy-grok is protected; set SYSTRAY_ALLOW_ROY_GROK_MUTATION=1 to re-auth", ) return yield LifecycleEvent(kind="started", account=account) try: device = start_device_code() uri = str(device.get("verification_uri_complete") or device.get("verification_uri") or "") user_code = str(device.get("user_code") or "") yield LifecycleEvent( kind="prompt_ready", account=account, message=f"Grok re-auth: open {uri} code {user_code}", prompt=_Prompt(url=uri, raw_text=f"{uri}\n{user_code}"), session=None, ) tokens = poll_device_code(device) access = str(tokens.get("access_token") or "") refresh = str(tokens.get("refresh_token") or "") if not access or not refresh: raise GrokAuthOperationError("device login missing tokens") write_auth_file( account.account_home / "auth.json", access=access, refresh=refresh, expires_in=float(tokens["expires_in"]) if isinstance(tokens.get("expires_in"), (int, float)) else None, ) self._write_identity(account.account_home, account.account_home / "auth.json") yield LifecycleEvent(kind="success", account=account) except Exception as exc: yield LifecycleEvent( kind="failure", account=account, message=f"Grok re-auth failed: {exc}", error=exc, ) def import_from_legacy_home( self, *, slug: str = PROTECTED_SLUG, alias: str = PROTECTED_SLUG, legacy_auth: Path | None = None, ) -> Account: """Copy-only import of ~/.grok/auth.json. Never hits the token endpoint.""" source = legacy_auth or (Path.home() / ".grok" / "auth.json") if not source.is_file(): raise GrokAuthOperationError(f"missing Grok auth at {source}") existing = next((a for a in self.registry.list() if a.slug == slug), None) if existing is not None: return existing account_home = self.registry.add_dir(slug, alias) dest = account_home / "auth.json" shutil.copy2(source, dest) os.chmod(dest, 0o600) try: self._write_identity(account_home, dest, allow_refresh=False) except Exception: # Identity is best-effort; credentials are what matter. pass account = next(a for a in self.registry.list() if a.slug == slug) self.registry.set_default(account) return account def _write_identity( self, account_home: Path, auth_path: Path, *, allow_refresh: bool = True, ) -> dict[str, str | None]: session = GrokOAuthSession(auth_path) user = session.get_user(allow_refresh=allow_refresh) email = user.get("email") if isinstance(user.get("email"), str) else None account_id = ( user.get("principalId") if isinstance(user.get("principalId"), str) else user.get("userId") if isinstance(user.get("userId"), str) else None ) plan = "grok" identity = { "email": email, "org_id": account_id, "subscription_type": plan, } path = account_home / "account_identity.json" path.write_text(json.dumps(identity, indent=2) + "\n", encoding="utf-8") os.chmod(path, 0o600) return {"email": email, "account_id": account_id, "plan": plan} def _find_collision(self, account_id: str | None, email: str | None) -> Account | None: for account in self.registry.list(): if account_id and account.account_id == account_id: return account if email and account.email and account.email.lower() == email.lower(): return account return None def _pending_account(self, staged: StagedAccount) -> Account: return Account( ref=AccountRef(self.registry.kind.value, staged.slug), alias=staged.alias, account_home=staged.account_home, email=None, plan=None, account_id=None, ) class _Prompt: def __init__(self, url: str, raw_text: str) -> None: self.url = url self.raw_text = raw_text