auth_config.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""CLI account + token storage (~/.aquin/config.json). One token per Aquin account."""
from __future__ import annotations
 
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_CONFIG_PATH = Path.home() / ".aquin" / "config.json"
_BASE_URL_DEFAULT = "https://api.aquin.app"
# Re-check saved tokens against cloud at most this often (seconds).
_WHOAMI_CACHE_TTL_S = 15 * 60
_MIN_TOKEN_LEN = 20  # aq- + 16 hex bytes = 19; allow a little headroom for future formats
 
 
def config_path() -> Path:
    return _CONFIG_PATH
 
 
def _base_url() -> str:
    return os.environ.get("AQUIN_BASE_URL", _BASE_URL_DEFAULT).rstrip("/")
 
 
def safe_print(msg: str, *, file=None) -> None:
    """Print without crashing on Windows cp1252 consoles (no fancy unicode)."""
    stream = file or sys.stdout
    text = (
        str(msg)
        .replace("\u2192", "->")
        .replace("\u2014", "-")
        .replace("\u2013", "-")
        .replace("\u2026", "...")
    )
    try:
        print(text, file=stream)
    except UnicodeEncodeError:
        enc = getattr(stream, "encoding", None) or "utf-8"
        print(text.encode(enc, errors="replace").decode(enc, errors="replace"), file=stream)
 
 
def normalize_cli_token(raw: str) -> str:
    """Strip whitespace / accidental 'Bearer ' prefix."""
    key = (raw or "").strip()
    if key.lower().startswith("bearer "):
        key = key[7:].strip()
    return key
 
 
def validate_cli_token_format(raw: str) -> str | None:
    """
    Return an error message if the token is obviously not an Aquin CLI token.
    Returns None when the shape looks acceptable (still must pass whoami).
    """
    key = normalize_cli_token(raw)
    if not key:
        return "No token entered."
    if not key.startswith("aq-"):
        return (
            "Invalid CLI token: must start with aq- "
            "(copy it from aquin.app Profile -> Account -> CLI token)."
        )
    if len(key) < _MIN_TOKEN_LEN:
        return "Invalid CLI token: too short."
    if any(ch.isspace() for ch in key):
        return "Invalid CLI token: must not contain spaces."
    return None
 
 
def load_config() -> dict[str, Any]:
    if not _CONFIG_PATH.exists():
        return {"accounts": {}}
    try:
        data = json.loads(_CONFIG_PATH.read_text(encoding="utf-8"))
        if not isinstance(data, dict):
            return {"accounts": {}}
    except Exception:
        return {"accounts": {}}
    return _migrate_legacy(data)
 
 
def save_config(cfg: dict[str, Any]) -> None:
    _CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    _CONFIG_PATH.write_text(json.dumps(cfg, indent=2), encoding="utf-8")
 
 
def _migrate_legacy(cfg: dict[str, Any]) -> dict[str, Any]:
    accounts = cfg.get("accounts")
    if isinstance(accounts, dict) and accounts:
        cfg.setdefault("active_account", "")
        return cfg
    key = str(cfg.get("api_key") or "").strip()
    accounts_dict: dict[str, Any] = {}
    if key:
        label = str(cfg.get("active_account") or "saved").strip() or "saved"
        accounts_dict[label] = {
            "api_key": key,
            "name": cfg.get("account_name"),
            "email": cfg.get("account_email") if "@" in str(cfg.get("account_email") or "") else None,
        }
        cfg["active_account"] = label
    cfg["accounts"] = accounts_dict
    cfg.pop("api_key", None)
    return cfg
 
 
def _account_id(email: str | None, name: str | None, user_id: str | None) -> str:
    if email and "@" in email:
        return email.strip().lower()
    if user_id:
        return user_id
    if name:
        return name.strip().lower().replace(" ", "-")
    return "account"
 
 
def active_account_id(cfg: dict[str, Any] | None = None) -> str:
    cfg = cfg or load_config()
    return str(cfg.get("active_account") or "").strip()
 
 
def list_accounts(cfg: dict[str, Any] | None = None) -> list[dict[str, str]]:
    cfg = cfg or load_config()
    active = active_account_id(cfg)
    rows: list[dict[str, str]] = []
    for aid, meta in (cfg.get("accounts") or {}).items():
        if not isinstance(meta, dict):
            continue
        label = str(meta.get("email") or meta.get("name") or aid)
        rows.append({
            "id": str(aid),
            "label": label,
            "name": str(meta.get("name") or ""),
            "email": str(meta.get("email") or ""),
            "active": "yes" if aid == active else "",
        })
    return rows
 
 
def resolve_api_key(*, allow_missing: bool = False, explicit: str | None = None) -> str:
    if explicit and explicit.strip():
        return normalize_cli_token(explicit)
    env = os.environ.get("AQUIN_API_KEY", "").strip()
    if env:
        return normalize_cli_token(env)
    cfg = load_config()
    aid = active_account_id(cfg)
    accounts = cfg.get("accounts") or {}
    if aid and isinstance(accounts.get(aid), dict):
        key = str(accounts[aid].get("api_key") or "").strip()
        if key:
            return normalize_cli_token(key)
    if accounts:
        first = next(iter(accounts.values()))
        if isinstance(first, dict):
            key = str(first.get("api_key") or "").strip()
            if key:
                return normalize_cli_token(key)
    if not allow_missing:
        safe_print("Not logged in. Run:  aquin login", file=sys.stderr)
        sys.exit(1)
    return ""
 
 
def fetch_whoami(api_key: str, base_url: str | None = None) -> dict[str, Any]:
    import httpx
 
    key = normalize_cli_token(api_key)
    fmt_err = validate_cli_token_format(key)
    if fmt_err:
        return {"ok": False, "error": "invalid_format", "hint": fmt_err}
 
    url = (base_url or _base_url()).rstrip("/")
    try:
        resp = httpx.get(
            f"{url}/api/sdk/whoami",
            headers={"Authorization": f"Bearer {key}"},
            timeout=15.0,
            verify=False,
        )
    except Exception as exc:
        return {
            "ok": False,
            "error": "network",
            "detail": str(exc),
            "hint": "Could not reach Aquin to verify the token.",
        }
 
    if resp.status_code == 401:
        try:
            body = resp.json()
        except Exception:
            body = {}
        hint = body.get("hint") or (
            "Invalid CLI token. Copy a fresh aq- token from aquin.app "
            "(Profile -> Account -> CLI token), then run aquin login --force."
        )
        return {
            "ok": False,
            "error": body.get("error") or "Unauthorized",
            "revoked_at": body.get("revoked_at"),
            "hint": str(hint).replace("\u2192", "->"),
        }
 
    if resp.status_code != 200:
        return {
            "ok": False,
            "error": f"HTTP {resp.status_code}",
            "hint": f"Token verification failed (HTTP {resp.status_code}).",
        }
 
    try:
        data = resp.json()
    except Exception:
        return {
            "ok": False,
            "error": "bad_response",
            "hint": "Token verification returned an invalid response.",
        }
 
    if not isinstance(data, dict):
        return {
            "ok": False,
            "error": "bad_response",
            "hint": "Token verification returned an invalid response.",
        }
 
    user_id = data.get("user_id")
    email = data.get("email")
    name = data.get("name")
    if not user_id and not email and not name:
        return {
            "ok": False,
            "error": "empty_identity",
            "hint": "Token verification succeeded but returned no account identity.",
        }
 
    return {
        "ok": True,
        "user_id": user_id,
        "email": email,
        "name": name,
        "avatar_url": data.get("avatar_url"),
    }
 
 
def format_account_label(whoami: dict[str, Any]) -> str:
    email = str(whoami.get("email") or "").strip()
    name = str(whoami.get("name") or "").strip()
    if name and email:
        return f"{name} ({email})"
    return email or name or "?"
 
 
def cached_account_label(cfg: dict[str, Any] | None = None) -> str:
    """Account label from saved login metadata - no network."""
    cfg = cfg or load_config()
    aid = active_account_id(cfg)
    accounts = cfg.get("accounts") or {}
    meta = accounts.get(aid) if aid else None
    if not isinstance(meta, dict) and accounts:
        meta = next(iter(accounts.values()), None)
    if not isinstance(meta, dict):
        return "?"
    name = str(meta.get("name") or "").strip()
    email = str(meta.get("email") or "").strip()
    if name and email:
        return f"{name} ({email})"
    return email or name or str(aid or "?")
 
 
def format_revoked_ago(revoked_at: str) -> str:
    try:
        raw = revoked_at.replace("Z", "+00:00")
        dt = datetime.fromisoformat(raw)
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        delta = datetime.now(timezone.utc) - dt.astimezone(timezone.utc)
        secs = int(max(0, delta.total_seconds()))
        if secs < 60:
            return "just now"
        mins = secs // 60
        if mins < 60:
            return f"{mins} minute{'s' if mins != 1 else ''} ago"
        hours = mins // 60
        if hours < 48:
            return f"{hours} hour{'s' if hours != 1 else ''} ago"
        days = hours // 24
        return f"{days} day{'s' if days != 1 else ''} ago"
    except Exception:
        return "recently"
 
 
def _cache_whoami_ok(api_key: str, whoami: dict[str, Any]) -> None:
    """Remember a successful whoami so we do not hit the network every command."""
    cfg = load_config()
    aid = active_account_id(cfg)
    accounts = dict(cfg.get("accounts") or {})
    meta = accounts.get(aid) if aid else None
    if not isinstance(meta, dict):
        return
    if normalize_cli_token(str(meta.get("api_key") or "")) != normalize_cli_token(api_key):
        return
    meta = dict(meta)
    meta["verified_at"] = time.time()
    if whoami.get("email"):
        meta["email"] = whoami.get("email")
    if whoami.get("name"):
        meta["name"] = whoami.get("name")
    if whoami.get("user_id"):
        meta["user_id"] = whoami.get("user_id")
    if whoami.get("avatar_url"):
        meta["avatar_url"] = whoami.get("avatar_url")
    accounts[aid] = meta
    cfg["accounts"] = accounts
    save_config(cfg)
 
 
def _cached_whoami_fresh(api_key: str, *, ttl_s: float = _WHOAMI_CACHE_TTL_S) -> bool:
    cfg = load_config()
    aid = active_account_id(cfg)
    accounts = cfg.get("accounts") or {}
    meta = accounts.get(aid) if aid else None
    if not isinstance(meta, dict):
        return False
    if normalize_cli_token(str(meta.get("api_key") or "")) != normalize_cli_token(api_key):
        return False
    try:
        verified_at = float(meta.get("verified_at") or 0)
    except (TypeError, ValueError):
        return False
    return verified_at > 0 and (time.time() - verified_at) < ttl_s
 
 
def require_valid_login(*, force_refresh: bool = False) -> str:
    """
    Ensure a CLI token is present AND valid with Aquin cloud.
    Exits the process on failure. Returns the verified api key.
    """
    key = resolve_api_key(allow_missing=False)
    fmt_err = validate_cli_token_format(key)
    if fmt_err:
        safe_print(fmt_err, file=sys.stderr)
        safe_print("Run:  aquin login --force", file=sys.stderr)
        sys.exit(1)
 
    if not force_refresh and _cached_whoami_fresh(key):
        return key
 
    whoami = fetch_whoami(key)
    if whoami.get("ok"):
        _cache_whoami_ok(key, whoami)
        return key
 
    if whoami.get("revoked_at"):
        ago = format_revoked_ago(str(whoami["revoked_at"]))
        safe_print(f"Your CLI token was revoked {ago}.", file=sys.stderr)
        safe_print("Regenerate at aquin.app (Profile -> Account -> CLI token), then:", file=sys.stderr)
        safe_print("  aquin login --force", file=sys.stderr)
        sys.exit(1)
 
    if whoami.get("error") == "network":
        # Offline: allow local GPU work if we previously verified this token.
        if _cached_whoami_fresh(key, ttl_s=7 * 24 * 3600):
            return key
        safe_print(
            whoami.get("hint") or "Could not reach Aquin to verify your CLI token.",
            file=sys.stderr,
        )
        safe_print("Check your network, or run:  aquin login --force", file=sys.stderr)
        sys.exit(1)
 
    print_token_rejected(key)
    return key  # unreachable
 
 
def print_token_rejected(api_key: str | None = None) -> None:
    """User-facing auth failure with revoked-at hint when available."""
    safe_print("Cannot connect to Aquin: your CLI token was rejected.", file=sys.stderr)
    if api_key:
        status = fetch_whoami(api_key)
        if status.get("revoked_at"):
            ago = format_revoked_ago(str(status["revoked_at"]))
            safe_print(f"This token was revoked {ago}.", file=sys.stderr)
        elif status.get("hint"):
            safe_print(str(status["hint"]), file=sys.stderr)
    safe_print("", file=sys.stderr)
    safe_print("Run:  aquin login --force   (paste a fresh aq- token)", file=sys.stderr)
    safe_print("      aquin switch          (use another saved account)", file=sys.stderr)
    safe_print("      aquin.app/profile     (Profile -> Account -> CLI token)", file=sys.stderr)
    sys.exit(1)
 
 
def save_login(api_key: str, *, whoami: dict[str, Any] | None = None) -> str:
    key = normalize_cli_token(api_key)
    fmt_err = validate_cli_token_format(key)
    if fmt_err:
        raise ValueError(fmt_err)
 
    info = whoami if whoami and whoami.get("ok") else fetch_whoami(key)
    if not info.get("ok"):
        if info.get("revoked_at"):
            ago = format_revoked_ago(str(info["revoked_at"]))
            raise ValueError(f"Token was revoked {ago}. Regenerate at aquin.app and try again.")
        hint = info.get("hint") or "Invalid CLI token."
        raise ValueError(str(hint).replace("\u2192", "->"))
 
    aid = _account_id(
        str(info.get("email") or ""),
        str(info.get("name") or ""),
        str(info.get("user_id") or ""),
    )
    cfg = load_config()
    accounts = dict(cfg.get("accounts") or {})
    accounts[aid] = {
        "api_key": key,
        "email": info.get("email"),
        "name": info.get("name"),
        "user_id": info.get("user_id"),
        "avatar_url": info.get("avatar_url"),
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "verified_at": time.time(),
    }
    cfg["accounts"] = accounts
    cfg["active_account"] = aid
    save_config(cfg)
    return aid
 
 
def switch_account(account_id: str) -> dict[str, str]:
    cfg = load_config()
    accounts = cfg.get("accounts") or {}
    aid = account_id.strip().lower()
    if aid not in accounts:
        for k in accounts:
            meta = accounts[k]
            if not isinstance(meta, dict):
                continue
            email = str(meta.get("email") or "").lower()
            if email == aid or k.lower() == aid:
                aid = k
                break
        else:
            raise ValueError(
                f"No CLI token saved for '{account_id}' on this machine.\n"
                f"Add that account's token:  aquin switch --add"
            )
 
    cfg["active_account"] = aid
    save_config(cfg)
    meta = accounts[aid]
    return {
        "id": aid,
        "label": format_account_label(meta),
        "api_key": str(meta.get("api_key") or ""),
    }
 
 
def logout_local(*, account_id: str | None = None) -> None:
    cfg = load_config()
    accounts = dict(cfg.get("accounts") or {})
    if account_id:
        aid = account_id.strip().lower()
        keys = [
            k
            for k in accounts
            if k.lower() == aid or str((accounts[k] or {}).get("email") or "").lower() == aid
        ]
        for k in keys:
            accounts.pop(k, None)
        if active_account_id(cfg) in keys:
            cfg.pop("active_account", None)
    else:
        aid = active_account_id(cfg)
        if aid:
            accounts.pop(aid, None)
        cfg.pop("active_account", None)
    cfg["accounts"] = accounts
    if accounts and not cfg.get("active_account"):
        cfg["active_account"] = next(iter(accounts))
    save_config(cfg)