298 lines
13 KiB
Python
Executable File
298 lines
13 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Grok token broker — hold a refreshable xAI session credential WITHOUT touching the CLI's.
|
|
|
|
⭐ THE PROBLEM. The Grok Build coding plan (`https://cli-chat-proxy.grok.com/v1`, serving
|
|
grok-4.6) authenticates with an OIDC **session** token, not an API key: `api_key: null`,
|
|
`env_key: null` in the CLI's own model cache. It expires roughly every six hours. LiteLLM and
|
|
every other gateway on this fleet hold a STATIC credential, so a naive alias works until the
|
|
session lapses and then fails closed.
|
|
|
|
⚠⚠ THE HAZARD THAT SHAPES THIS ENTIRE DESIGN, AND WHY IT IS NOT A SIMPLE REFRESH LOOP.
|
|
`https://auth.x.ai/oauth2/token` supports the `refresh_token` grant. But **a refresh may ROTATE
|
|
the refresh token**, and many OIDC providers invalidate the old one server-side the instant a
|
|
new one is issued. The operator's interactive `grok` CLI holds its refresh token in
|
|
`~/.grok/auth.json`. If xAI rotates-and-invalidates, then a broker refreshing against the SAME
|
|
credential silently kills the CLI login — and it does so **even though the broker never writes
|
|
to that file.** Not-writing is necessary and NOT sufficient; the damage would be server-side.
|
|
|
|
heid's `groa_http_dispatch.py` refuses to refresh at all for this reason, which is the correct
|
|
call absent an answer. This broker's job is to get the answer first and then act on it.
|
|
|
|
⛔ SO THE REFRESH LOOP IS NOT ARMED UNTIL `probe-rotation` HAS RUN AND SAID IT IS SAFE.
|
|
`serve` and `refresh` both refuse without a stored probe verdict. That is deliberate: this is a
|
|
tool whose whole value is contingent on one unknown, and running it before measuring the unknown
|
|
is how you find out by breaking the operator's login.
|
|
|
|
## The probe, and what each outcome means
|
|
|
|
`probe-rotation` performs ONE refresh using a COPY of the credential, then re-presents the
|
|
ORIGINAL refresh token:
|
|
|
|
original still valid -> NON-ROTATING (or rotation without invalidation). The broker can
|
|
coexist with the CLI on one credential. Loop may be armed.
|
|
original now rejected -> ROTATING-AND-INVALIDATING. The broker MUST NOT share the CLI's
|
|
credential. It needs its own login, or the idea is dead. Refuses
|
|
to arm, and says so.
|
|
|
|
⚠ The probe itself consumes one refresh. If the answer is "rotating", the CLI's token is
|
|
already dead by the time you read the verdict — there is no way to ask this question without
|
|
spending the credential, and pretending otherwise would be worse. The remedy is a `grok` CLI
|
|
re-login, which is the same remedy heid's exit-3 path already names. **Run it when a broken CLI
|
|
login is a two-minute annoyance, not mid-panel.**
|
|
|
|
## What it never does
|
|
|
|
- Never writes `~/.grok/auth.json`. Reads it once, on `seed`, and never again.
|
|
- Never prints a token to stdout except under `token --raw`, which is for a consumer to capture.
|
|
- Never logs a token value. Expiry, scope and subject only.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
CLI_AUTH = pathlib.Path.home() / ".grok" / "auth.json"
|
|
STORE = pathlib.Path(os.environ.get("GROK_BROKER_HOME",
|
|
pathlib.Path.home() / ".config" / "grok-token-broker"))
|
|
CRED = STORE / "credential.json"
|
|
VERDICT = STORE / "rotation-verdict.json"
|
|
TOKEN_URL = "https://auth.x.ai/oauth2/token"
|
|
API_BASE = "https://cli-chat-proxy.grok.com/v1"
|
|
REFRESH_MARGIN_S = 1800 # refresh when under 30 min remain
|
|
|
|
|
|
def _claims(jwt: str) -> dict:
|
|
pl = jwt.split(".")[1]
|
|
pl += "=" * (-len(pl) % 4)
|
|
return json.loads(base64.urlsafe_b64decode(pl))
|
|
|
|
|
|
def _remaining(jwt: str) -> int:
|
|
try:
|
|
return int(_claims(jwt)["exp"] - time.time())
|
|
except Exception:
|
|
return -1
|
|
|
|
|
|
def _post_form(url: str, fields: dict, timeout: int = 30):
|
|
data = urllib.parse.urlencode(fields).encode()
|
|
req = urllib.request.Request(url, data=data,
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"})
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=timeout) as r:
|
|
return r.status, json.loads(r.read())
|
|
except urllib.error.HTTPError as e:
|
|
body = e.read().decode("utf-8", "replace")[:400]
|
|
return e.code, {"_error": body}
|
|
except Exception as e: # noqa: BLE001
|
|
return 0, {"_error": f"{type(e).__name__}: {e}"}
|
|
|
|
|
|
def _read_cli_credential() -> dict:
|
|
"""Pull access token, refresh token and client id out of the CLI's auth file.
|
|
|
|
Shape-tolerant on purpose: this file belongs to another tool and its layout is not ours to
|
|
depend on. Anything not found is reported, never guessed.
|
|
"""
|
|
if not CLI_AUTH.exists():
|
|
raise SystemExit(f"REFUSING: no CLI credential at {CLI_AUTH}")
|
|
blob = json.loads(CLI_AUTH.read_text(encoding="utf-8"))
|
|
found: dict = {}
|
|
|
|
def walk(o, key_hint=""):
|
|
if isinstance(o, dict):
|
|
for k, v in o.items():
|
|
if isinstance(v, str):
|
|
if k == "key" and v.count(".") == 2 and "access" not in found:
|
|
found["access"] = v
|
|
elif "refresh" in k.lower() and len(v) > 20:
|
|
found.setdefault("refresh", v)
|
|
walk(v, k)
|
|
if "auth.x.ai::" in key_hint:
|
|
found.setdefault("client_id", key_hint.split("::", 1)[1])
|
|
elif isinstance(o, list):
|
|
for v in o:
|
|
walk(v, key_hint)
|
|
|
|
for k, v in blob.items():
|
|
walk(v, k)
|
|
missing = [f for f in ("access", "refresh") if f not in found]
|
|
if missing:
|
|
raise SystemExit(f"REFUSING: could not locate {missing} in {CLI_AUTH}. "
|
|
"Its layout changed; inspect it by hand rather than letting this guess.")
|
|
return found
|
|
|
|
|
|
def cmd_seed(_a) -> int:
|
|
STORE.mkdir(parents=True, exist_ok=True)
|
|
os.chmod(STORE, 0o700)
|
|
c = _read_cli_credential()
|
|
CRED.write_text(json.dumps({**c, "seeded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
|
"source": str(CLI_AUTH)}, indent=2), encoding="utf-8")
|
|
os.chmod(CRED, 0o600)
|
|
cl = _claims(c["access"])
|
|
print(f" seeded -> {CRED} (0600)")
|
|
print(f" scope : {cl.get('scope')}")
|
|
print(f" expires : {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(cl['exp']))} "
|
|
f"({_remaining(c['access']) // 60} min left)")
|
|
print(f" client : {c.get('client_id', '<not found>')}")
|
|
print("\n ⛔ The refresh loop is NOT armed. Run `probe-rotation` next, and read its warning "
|
|
"first — the probe spends one refresh and may end the CLI's session.")
|
|
return 0
|
|
|
|
|
|
def cmd_probe(a) -> int:
|
|
if not CRED.exists():
|
|
raise SystemExit("REFUSING: run `seed` first")
|
|
c = json.loads(CRED.read_text(encoding="utf-8"))
|
|
if not a.i_accept_this_may_end_the_cli_session:
|
|
print("REFUSING without --i-accept-this-may-end-the-cli-session.\n\n"
|
|
" This probe performs a real refresh. If xAI rotates AND invalidates, the\n"
|
|
" operator's `grok` CLI login dies at that moment and needs an interactive\n"
|
|
" re-login. There is no way to answer the question without spending the\n"
|
|
" credential. Run it when that is a two-minute annoyance, not mid-panel.",
|
|
file=sys.stderr)
|
|
return 2
|
|
|
|
old_refresh = c["refresh"]
|
|
fields = {"grant_type": "refresh_token", "refresh_token": old_refresh}
|
|
if c.get("client_id"):
|
|
fields["client_id"] = c["client_id"]
|
|
|
|
print(" [1/2] refreshing with the stored refresh token ...")
|
|
st, body = _post_form(TOKEN_URL, fields)
|
|
if st != 200:
|
|
print(f" refresh FAILED http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr)
|
|
print(" Verdict NOT recorded. The loop stays disarmed.", file=sys.stderr)
|
|
return 4
|
|
new_refresh = body.get("refresh_token")
|
|
rotated = bool(new_refresh and new_refresh != old_refresh)
|
|
print(f" ok. refresh_token {'ROTATED' if rotated else 'unchanged'}")
|
|
|
|
if not rotated:
|
|
verdict, safe = "non-rotating", True
|
|
note = "Provider returned the same refresh token; the CLI's copy is untouched."
|
|
else:
|
|
print(" [2/2] re-presenting the ORIGINAL refresh token to see if it still works ...")
|
|
st2, _ = _post_form(TOKEN_URL, fields)
|
|
if st2 == 200:
|
|
verdict, safe = "rotating-old-still-valid", True
|
|
note = ("Rotation happens but the old refresh token still works, so the CLI's copy "
|
|
"survives. Coexistence is safe.")
|
|
else:
|
|
verdict, safe = "rotating-and-invalidating", False
|
|
note = (f"The original refresh token is now rejected (http={st2}). The CLI's login "
|
|
"is ALREADY DEAD and needs an interactive `grok` re-login. A broker MUST "
|
|
"NOT share this credential — it needs its own login.")
|
|
print(f" original token -> http={st2}")
|
|
|
|
c["access"] = body.get("access_token", c["access"])
|
|
if new_refresh:
|
|
c["refresh"] = new_refresh
|
|
CRED.write_text(json.dumps(c, indent=2), encoding="utf-8")
|
|
os.chmod(CRED, 0o600)
|
|
|
|
VERDICT.write_text(json.dumps({
|
|
"verdict": verdict, "safe_to_arm": safe, "note": note,
|
|
"probed_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
|
"token_url": TOKEN_URL, "api_base": API_BASE,
|
|
}, indent=2), encoding="utf-8")
|
|
os.chmod(VERDICT, 0o600)
|
|
print(f"\n VERDICT: {verdict} (safe_to_arm={safe})\n {note}")
|
|
if not safe:
|
|
print("\n ⚠ TELL THE OPERATOR: the `grok` CLI on this box needs a re-login.", file=sys.stderr)
|
|
return 0 if safe else 5
|
|
|
|
|
|
def _ensure_armed() -> dict:
|
|
if not VERDICT.exists():
|
|
raise SystemExit("REFUSING: no rotation verdict. Run `probe-rotation` first — this tool "
|
|
"will not refresh a credential whose rotation behaviour is unmeasured.")
|
|
v = json.loads(VERDICT.read_text(encoding="utf-8"))
|
|
if not v.get("safe_to_arm"):
|
|
raise SystemExit(f"REFUSING: verdict is '{v['verdict']}'. {v['note']}")
|
|
return v
|
|
|
|
|
|
def cmd_refresh(_a) -> int:
|
|
_ensure_armed()
|
|
c = json.loads(CRED.read_text(encoding="utf-8"))
|
|
fields = {"grant_type": "refresh_token", "refresh_token": c["refresh"]}
|
|
if c.get("client_id"):
|
|
fields["client_id"] = c["client_id"]
|
|
st, body = _post_form(TOKEN_URL, fields)
|
|
if st != 200:
|
|
print(f"refresh failed http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr)
|
|
return 4
|
|
c["access"] = body["access_token"]
|
|
if body.get("refresh_token"):
|
|
c["refresh"] = body["refresh_token"]
|
|
c["refreshed_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
|
|
CRED.write_text(json.dumps(c, indent=2), encoding="utf-8")
|
|
os.chmod(CRED, 0o600)
|
|
print(f" refreshed; {_remaining(c['access']) // 60} min remaining")
|
|
return 0
|
|
|
|
|
|
def cmd_token(a) -> int:
|
|
"""Print a currently-valid access token, refreshing first if it is close to expiry."""
|
|
if not CRED.exists():
|
|
raise SystemExit("REFUSING: run `seed` first")
|
|
c = json.loads(CRED.read_text(encoding="utf-8"))
|
|
if _remaining(c["access"]) < REFRESH_MARGIN_S:
|
|
_ensure_armed()
|
|
if cmd_refresh(a) != 0:
|
|
return 4
|
|
c = json.loads(CRED.read_text(encoding="utf-8"))
|
|
if a.raw:
|
|
sys.stdout.write(c["access"])
|
|
else:
|
|
print(f" valid for {_remaining(c['access']) // 60} more minutes "
|
|
f"(use --raw to emit the token itself)")
|
|
return 0
|
|
|
|
|
|
def cmd_status(_a) -> int:
|
|
print(f" store : {STORE}")
|
|
if not CRED.exists():
|
|
print(" state : NOT SEEDED")
|
|
return 0
|
|
c = json.loads(CRED.read_text(encoding="utf-8"))
|
|
print(f" seeded : {c.get('seeded_at')} from {c.get('source')}")
|
|
print(f" access : {_remaining(c['access']) // 60} min remaining")
|
|
if VERDICT.exists():
|
|
v = json.loads(VERDICT.read_text(encoding="utf-8"))
|
|
print(f" rotation: {v['verdict']} safe_to_arm={v['safe_to_arm']} ({v['probed_at']})")
|
|
print(f" armed : {'YES' if v['safe_to_arm'] else 'NO — ' + v['verdict']}")
|
|
else:
|
|
print(" rotation: UNMEASURED — loop disarmed, `token` will refuse to refresh")
|
|
print(f" api_base: {API_BASE}")
|
|
return 0
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
sub.add_parser("seed", help="copy the CLI credential into the broker's own store")
|
|
p = sub.add_parser("probe-rotation", help="measure whether refresh rotates AND invalidates")
|
|
p.add_argument("--i-accept-this-may-end-the-cli-session", action="store_true")
|
|
sub.add_parser("refresh", help="refresh now (requires a safe verdict)")
|
|
t = sub.add_parser("token", help="emit a valid access token, refreshing if near expiry")
|
|
t.add_argument("--raw", action="store_true")
|
|
sub.add_parser("status", help="report state without touching anything")
|
|
a = ap.parse_args()
|
|
return {"seed": cmd_seed, "probe-rotation": cmd_probe, "refresh": cmd_refresh,
|
|
"token": cmd_token, "status": cmd_status}[a.cmd](a)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|