pilot.py
Rama main · réplica de lectura del repositorio original.
"""
Pilot programme: share usage COUNTS with the Squidbrake team: with a code you were given, or, for every other
install, as the anonymous usage stats that start with the first `squidbrake` command run in a terminal
(telemetry.py; they join the community code below).
squidbrake pilot join CODE --server URL start sharing (shows exactly what is sent, and asks first)
squidbrake pilot status what is shared, where, and when it was last sent
squidbrake pilot leave stop sharing
Nothing is sent unless you join. What is sent, at start, every 2 minutes while something changed (at least every
6 hours otherwise), and once more when the gateway stops:
- the Squidbrake version, operating system, enforce/shadow mode and number of rules
- the agents connected, by their labels (e.g. "claude-code", "antigravity"), and which agents on this computer
have Squidbrake's hook in their settings (so a pilot that installed but never connected an agent can be helped)
- per day, for the last 7 days: how many actions were allowed, held, approved, rejected, blocked, timed out or failed,
and how many were paused by an emergency stop (counted apart from what rules blocked), and how many a person
marked "this was legit" (squidbrake feedback)
- how long the gateway took to decide, as p50 / p95 / p99 milliseconds
- which rules blocked or held things (rule ids such as "command:catastrophic_command")
- for each action held or blocked in the last 7 days: when, which agent, the program only (e.g. "rm", "git"),
the rule and its reason, what happened (approved, rejected, blocked, timed out), how long a person took,
and the size of what it would have changed, as numbers only (e.g. 3 commits, 1,204 files, 4,312 rows)
When you join through the first-run question (telemetry.py), also: where you heard about Squidbrake, and the work
email and name you gave, if you did (so your usage shows under your name, not as an anonymous install).
Never sent: commands or their arguments, file or folder names, file contents, prompts, tool inputs or outputs,
keys, names of people.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
import os
import platform
import re
import sys
import time
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path
from urllib.parse import urlparse
import httpx
from sqlalchemy import case, func, select
import weekly
log = logging.getLogger("gateway")
# How often the counts are looked at; they are only sent when something changed since the last send, or HEARTBEAT
# went by (so the dashboard still sees a running gateway that's idle). Every send is the last 7 days in full and the
# server overwrites what it had, so sending often never counts anything twice.
INTERVAL = int(os.getenv("SQUIDBRAKE_PILOT_INTERVAL", "120"))
HEARTBEAT = int(os.getenv("SQUIDBRAKE_PILOT_HEARTBEAT", str(6 * 3600)))
_last: dict[str, tuple[str, float]] = {} # data folder -> (fingerprint of the last payload sent, when)
WHAT_IS_SENT = __doc__.split("Nothing is sent unless you join.")[1].strip()
def _path(home: Path) -> Path:
return home / "pilot.json" # `home` is the gateway's data folder (next to keys.json): writable everywhere
def load(home: Path) -> dict | None:
try:
return json.loads(_path(home).read_text(encoding="utf-8"))
except (OSError, ValueError):
return None
def _save(home: Path, cfg: dict) -> None:
home.mkdir(parents=True, exist_ok=True)
_path(home).write_text(json.dumps(cfg, indent=2), encoding="utf-8")
SIZE = re.compile(r"(\d[\d,]*)\+?\s+(Kubernetes objects?|files?|rows?|commits?|resources?|objects?)\b", re.I)
WHY = {"catastrophic": "could wipe a home folder, a drive or the system", "irreversible": "can't be undone",
"hidden": "runs code that can't be read first"}
HISTORY_WHY = {"repeat_of_rejected": "asked again after a person said no", "impersonation": "money after a look-alike sender",
"payment_request_in_message": "money a message asked for", "duplicate_change": "the same change again",
"untrusted_destination": "sends to an address found only in outside content"}
def _program(name: str, raw_input: str | None) -> tuple[str | None, str | None]:
"""(program, category) of a shell command, e.g. ("rm", "irreversible"): never its arguments."""
try:
import commands
cmd = commands.command_of(json.loads(raw_input)) if raw_input else None
if not cmd:
return None, None
reading = commands.read(cmd)
worst = reading.worst()
prog = re.sub(r"[^a-z0-9._+-]", "", (worst.program if worst else "").lower())[:30] or None
return prog, reading.kind
except Exception:
return None, None
def catches(conn, events, since: str, reasons: dict) -> list[dict]:
"""What was held or blocked, without what it was about: program, rule, outcome, and sizes as numbers."""
rows = conn.execute(select(
events.c.created_at, events.c.decided_at, events.c.source, events.c.client, events.c.name, events.c.input,
events.c.rule_id, events.c.decision, events.c.decided_by, events.c.approval_deadline, events.c.metadata,
).where(events.c.created_at >= since, events.c.approval_deadline.isnot(None) |
((events.c.status == "denied") & events.c.decided_by.is_(None))
).order_by(events.c.created_at.desc()).limit(60)).all()
out = []
for r in rows:
prog, kind = _program(r.name or "", r.input)
rule = (r.rule_id or "")[:80]
why = reasons.get(rule) or (HISTORY_WHY.get(rule.split(":", 1)[1]) if rule.startswith(("history:", "taint:")) else None) \
or WHY.get(kind or "") or ""
if r.approval_deadline is None:
outcome = "blocked"
else:
outcome = {None: "waiting", "timeout": "timed out"}.get(r.decided_by) or ("approved" if r.decision == "allow" else "rejected")
saved: dict[str, int] = {}
try:
for line in (json.loads(r.metadata or "{}").get("effects") or []):
for n, unit in SIZE.findall(str(line)):
u = unit.lower().replace("kubernetes ", "")
u = u if u.endswith("s") else u + "s"
saved[u] = saved.get(u, 0) + int(n.replace(",", ""))
except (ValueError, AttributeError):
pass
decide_s = None
if r.decided_at and r.decided_by not in (None, "timeout"):
try:
decide_s = int((datetime.fromisoformat(r.decided_at.replace("Z", "+00:00"))
- datetime.fromisoformat(r.created_at.replace("Z", "+00:00"))).total_seconds())
except ValueError:
pass
tool = re.sub(r"[^A-Za-z0-9_.:-]", "", r.name or "")[:60]
out.append({"t": (r.created_at or "")[:16], "agent": str(r.source or r.client or "")[:40], "tool": tool,
"program": prog, "category": kind if prog else None, "rule": rule, "why": str(why)[:100],
"outcome": outcome, "decide_s": decide_s, "saved": saved})
return out
PAUSE_RULES = ("emergency-stop", "session-stop")
def usage(engine, events, mode: str, rules: int, version: str, reasons: dict | None = None,
connected: list[str] | None = None) -> dict:
"""Counts, and what was held or blocked as program + rule + outcome + sizes (see the top of this file).
No inputs, outputs, arguments, file names, names of people."""
since = (datetime.now(timezone.utc) - timedelta(days=7)).strftime("%Y-%m-%d")
day = func.substr(events.c.created_at, 1, 10)
human = events.c.decided_by.isnot(None) & (events.c.decided_by != "timeout")
n = lambda cond: func.sum(case((cond, 1), else_=0))
in_week = events.c.created_at >= since
paused = func.coalesce(events.c.rule_id.in_(PAUSE_RULES), False)
with engine.connect() as conn:
days = {d: {"events": int(t or 0), "allowed": int(a or 0), "held": int(h or 0), "approved": int(ap or 0),
"rejected": int(rj or 0), "blocked": int(b or 0), "timed_out": int(to or 0), "failed": int(f or 0),
"would_block": int(wb or 0), "would_hold": int(wh or 0), "paused": int(ps or 0),
"legit": int(lg or 0)}
for d, t, a, h, ap, rj, b, to, f, wb, wh, ps, lg in conn.execute(select(
day, func.count(),
n((events.c.decision == "allow") & events.c.decided_by.is_(None)),
n(events.c.approval_deadline.isnot(None)),
n(human & (events.c.decision == "allow")),
n(human & (events.c.decision == "deny")),
n((events.c.status == "denied") & events.c.decided_by.is_(None) & paused.is_(False)),
n(events.c.decided_by == "timeout"),
n(events.c.status == "failed"),
n(events.c.would == "deny"),
n(events.c.would == "review"),
n((events.c.status == "denied") & events.c.decided_by.is_(None) & paused.is_(True)),
n(events.c.feedback == "legit"), # a person said a hold or block should have run
).where(in_week).group_by(day)).all()}
agent = func.coalesce(events.c.source, events.c.client)
agents = dict(conn.execute(select(agent, func.count()).where(in_week).group_by(agent)).all())
rules_hit = dict(conn.execute(select(events.c.rule_id, func.count()).where(
in_week, events.c.rule_id.isnot(None),
(events.c.status == "denied") | events.c.approval_deadline.isnot(None),
).group_by(events.c.rule_id).order_by(func.count().desc()).limit(15)).all())
total, first = conn.execute(select(func.count(), func.min(events.c.created_at)).select_from(events)).one()
held = catches(conn, events, since, reasons or {})
took = [ms for (ms,) in conn.execute(select(events.c.decision_ms).where(
in_week, events.c.decision_ms.isnot(None)).limit(50000))]
return {"version": version, "os": f"{platform.system()} {platform.release()}", "python": platform.python_version(),
"mode": mode, "rules": rules, "agents": {str(k): v for k, v in agents.items() if k},
"days": days, "rules_hit": rules_hit, "total_events": total, "first_event": (first or "")[:10],
"catches": held, "connected": connected, "latency_ms": weekly.percentiles(took)}
def _post(server: str, path: str, body: dict, timeout: float = 15) -> dict:
r = httpx.post(server.rstrip("/") + path, json=body, timeout=timeout)
if r.status_code >= 400:
try:
detail = r.json().get("detail")
except ValueError:
detail = r.text[:200]
raise RuntimeError(f"{r.status_code}: {detail}")
return r.json()
def send(home: Path, payload: dict, timeout: float = 15) -> bool:
cfg = load(home)
if not cfg:
return False
import telemetry
if cfg.get("code") == telemetry.COMMUNITY_CODE and telemetry._switched_off():
return False # the anonymous counts stop with SQUIDBRAKE_TELEMETRY=0 / DO_NOT_TRACK=1, even after joining
body = {"code": cfg["code"], "install_id": cfg["install_id"], "usage": payload}
try:
try:
_post(cfg["server"], "/v1/ping", body, timeout)
except RuntimeError as e:
# 403: the server doesn't know this install any more (e.g. its record was removed): join again, same id
if not str(e).startswith("403") or not str(cfg["code"]).startswith("community-"):
raise
_post(cfg["server"], "/v1/pilot/join", {"code": cfg["code"], "install_id": cfg["install_id"],
"version": str(payload.get("version", ""))[:40],
"os": f"{platform.system()} {platform.release()}"[:80]}, timeout)
_post(cfg["server"], "/v1/ping", body, timeout)
except (httpx.HTTPError, RuntimeError) as e:
log.info("pilot: usage not sent (%s); will try again later", e)
return False
cfg["last_sent"] = datetime.now(timezone.utc).isoformat(timespec="seconds")
_save(home, cfg)
return True
def _fingerprint(payload: dict) -> str:
return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode()).hexdigest()
def send_if_changed(home: Path, make_payload, timeout: float = 15) -> bool | None:
"""Send the counts if they changed since the last send (or HEARTBEAT went by). None: nothing needed sending."""
if not load(home):
return None
payload = make_payload()
fp, key = _fingerprint(payload), str(home)
sent_fp, sent_at = _last.get(key, (None, 0.0))
if fp == sent_fp and time.monotonic() - sent_at < HEARTBEAT:
return None
ok = send(home, payload, timeout)
if ok:
_last[key] = (fp, time.monotonic())
return ok
async def loop(home: Path, make_payload) -> None:
"""Runs inside the gateway. Does nothing at all unless this install joined a pilot."""
await asyncio.sleep(20)
while True:
try:
await asyncio.to_thread(send_if_changed, home, make_payload)
except Exception as e: # never let reporting affect the gateway
log.info("pilot: %s", e)
await asyncio.sleep(INTERVAL)
def flush(home: Path, make_payload) -> None:
"""When the gateway stops (a laptop closing, a restart): send what changed since the last send, waiting a few
seconds at most, so the last minutes of work aren't held back until it starts again."""
try:
send_if_changed(home, make_payload, timeout=3)
except Exception as e:
log.info("pilot: %s", e)
# --------------------------------------------------------------------------- command line
def join(home: Path, code: str, server: str | None, yes: bool, version: str, ref: str | None = None,
contact: tuple[str, str | None] | None = None) -> int:
server = (server or os.getenv("SQUIDBRAKE_PILOT_SERVER") or "").rstrip("/")
host = urlparse(server).hostname or ""
# https, or plain http only to this machine or an internal name without dots (a hosted gateway's Docker network)
if not (server.startswith("https://") or (server.startswith("http://") and (host in ("localhost", "127.0.0.1") or "." not in host))):
print("Give the pilot server you were sent, e.g. --server https://... (https only)", file=sys.stderr)
return 2
print(f"\nJoining the Squidbrake pilot with code {code}.\nThis sends usage counts to {server}:\n")
print(" " + WHAT_IS_SENT.replace("\n", "\n ") + "\n")
if not yes:
try:
answer = input("Share these counts? [y/N] ")
except EOFError: # no keyboard here (piped, CI): don't guess
print(f"\nNot joined: nothing to answer with here. To join, run: squidbrake pilot join {code} --server {server}")
return 1
if answer.strip().lower() not in ("y", "yes"):
print("Not joined. Nothing is shared.")
return 1
cfg = load(home) or {}
install_id = cfg.get("install_id") or uuid.uuid4().hex
try:
r = _post(server, "/v1/pilot/join", {"code": code, "install_id": install_id, "version": version,
"os": f"{platform.system()} {platform.release()}",
**({"ref": ref} if ref else {}),
**({"email": contact[0], "name": contact[1] or ""} if contact else {})})
except (httpx.HTTPError, RuntimeError) as e:
print(f"Couldn't join: {e}", file=sys.stderr)
return 1
code = r.get("code") or code # an email on the community code gets its own pilot (insights/app.py)
_save(home, {"code": code, "server": server, "install_id": install_id, "company": r.get("company"),
"joined": datetime.now(timezone.utc).isoformat(timespec="seconds")})
print(f"Joined{(' as ' + r['company']) if r.get('company') else ''}. Thank you!"
f"\nCounts are sent while the gateway runs (restart it if it's running). Stop anytime: squidbrake pilot leave")
return 0
def leave(home: Path) -> int:
cfg = load(home)
if not cfg:
print("Not in a pilot; nothing is shared.")
return 0
try:
_post(cfg["server"], "/v1/pilot/leave", {"code": cfg["code"], "install_id": cfg["install_id"]})
except (httpx.HTTPError, RuntimeError):
pass
_path(home).unlink(missing_ok=True)
print("Left the pilot. Nothing more is shared.")
return 0
def status(home: Path) -> int:
cfg = load(home)
if not cfg:
print("Not in a pilot; nothing is shared.")
return 0
print(f"In the pilot {cfg.get('company') or cfg['code']} (code {cfg['code']}), sharing usage counts with {cfg['server']}.")
print(f"Last sent: {cfg.get('last_sent') or 'not yet (sent while the gateway runs)'}. Stop: squidbrake pilot leave")
return 0