diff --git a/docs/security.md b/docs/security.md index bfaf401..e48f5aa 100644 --- a/docs/security.md +++ b/docs/security.md @@ -120,8 +120,9 @@ Things worth knowing before you enrol: - Registrations, revocations and passkey sign-ins are recorded in the audit log (`dashboard.passkey.*`, `auth.passkey.*`, and `login`/`authorize` events tagged `via=passkey`). -Passkey ceremonies are rate-limited per IP by the same limiter that guards `/app/login`, and a TOTP -lockout also blocks the passkey path for that client. +Passkey ceremonies are rate-limited per IP by the same limiter that guards `/app/login`. A TOTP +lockout does not block them: a passkey cannot be guessed, so it stays the owner's way in while +someone burns the client's TOTP budget. Manage enrolled credentials from `/app/tokens`, or drop them all for a client with: diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 837d596..0309043 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -16,7 +16,7 @@ | MCP returns 401/403 only via the public Cloudflare URL (works on `localhost`) | Cloudflare Bot Fight Mode / WAF blocks the headless MCP client, or Cloudflare Access strips the `Authorization` header | Add a WAF skip rule + cache bypass for the MCP/OAuth paths and exclude `/mcp` from Access. See **[docs/cloudflare.md](cloudflare.md)**. A 401 body whose `hint` mentions `cf-ray` confirms the header was stripped at the edge. | | `{"error":"unauthorized"}` on `/authorize` | Client id not registered server-side | Create the client with `beaconmcp auth create`, then paste the returned id into the connector. | | `invalid_grant` with `missing or invalid totp` | 2FA code wrong, expired (>30 s), or already used | Generate a fresh code. Compare the server clock against the phone (`timedatectl` vs. authenticator). | -| `Too many attempts` on the 2FA page | 5 consecutive failed TOTP codes triggered a 5-minute lockout | Wait. The counter resets on the next valid code. | +| `Too many attempts` on the 2FA page | 5 failed TOTP codes from your address (or 10 for the client from any mix of addresses) triggered a 5-minute lockout | Wait, or sign in with a passkey, which the lockout does not block. Behind a reverse proxy, make sure `server.trusted_proxies` lists it, otherwise every visitor shares one budget. | | Clients silently rejected after upgrade | 2FA migration: clients without a TOTP secret are refused at startup | Inspect `journalctl -u beaconmcp` for the list, recreate them with `beaconmcp auth create`. | | Dashboard loops between `/app/refresh` and `/app/chat` | Bearer wiped by a service restart while the session cookie is still valid | Enter the TOTP code on the refresh page to mint a new bearer. | | External MCP clients disconnected after a restart | `TokenStore` is in-memory and wiped on restart | Recreate tokens from `/app/tokens` (max 3, 24 h expiry). | diff --git a/src/beaconmcp/__main__.py b/src/beaconmcp/__main__.py index 64f277e..ff7c9be 100644 --- a/src/beaconmcp/__main__.py +++ b/src/beaconmcp/__main__.py @@ -1106,7 +1106,7 @@ def _run_http(mcp, host: str, port: int): TotpResult, current_bearer_token, ) - from .ratelimit import RateLimiter, client_ip, forwarded_host + from .ratelimit import RateLimiter, TotpLockout, client_ip, forwarded_host from .server import config from .metrics import REGISTRY, auth_events, http_requests @@ -1230,31 +1230,18 @@ async def dispatch(self, request: Request, call_next): # the caller's bearer without an import cycle. auth.register_token_store(token_store) - # In-memory TOTP bruteforce guard: per client_id, (failures, cooldown_until). - # 5 failed attempts → 5-minute lockout. Reset on first success. - totp_fail_max = 5 - totp_lockout_seconds = 300 - totp_failures: dict[str, tuple[int, float]] = {} + # In-memory TOTP brute-force guard, keyed on (client_id, source IP): 5 + # wrong codes from one IP lock that IP out for 5 minutes, and 10 from any + # mix of IPs lock the client's TOTP (passkeys stay usable). Keyed on + # client_id alone, anyone who knew a client_id could lock its owner out + # of every 2FA path. + totp_lockout = TotpLockout(per_source=5, per_client=10, lockout_seconds=300) + totp_locked = totp_lockout.locked + totp_record_failure = totp_lockout.record_failure + totp_record_success = totp_lockout.record_success - def totp_locked(client_id: str) -> bool: - entry = totp_failures.get(client_id) - if not entry: - return False - count, until = entry - if count < totp_fail_max: - return False - if time.time() >= until: - totp_failures.pop(client_id, None) - return False - return True - - def totp_record_failure(client_id: str) -> None: - count, _ = totp_failures.get(client_id, (0, 0.0)) - count += 1 - totp_failures[client_id] = (count, time.time() + totp_lockout_seconds) - - def totp_record_success(client_id: str) -> None: - totp_failures.pop(client_id, None) + def _ip(request: Request) -> str: + return client_ip(request, tuple(config.server.trusted_proxies)) def _issuer(request: Request) -> str: scheme = request.headers.get("x-forwarded-proto", request.url.scheme) @@ -1619,7 +1606,7 @@ async def oauth_authorize_get(request: Request) -> Response: return _authorize_page( normalized, request=request, - locked=totp_locked(normalized["client_id"]), + locked=totp_locked(normalized["client_id"], _ip(request)), ) async def oauth_authorize_post(request: Request) -> Response: @@ -1629,8 +1616,21 @@ async def oauth_authorize_post(request: Request) -> Response: if err is not None: return err + # This endpoint takes a TOTP guess with nothing but a client_id, so + # it gets the same per-IP limit as /app/login. + ip = _ip(request) + if not _login_limiter.check(ip): + retry = _login_limiter.retry_after(ip) + response = _authorize_failed( + request, normalized, + f"Too many attempts. Retry in {retry}s.", + status=429, + ) + response.headers["Retry-After"] = str(retry) + return response + client_id = normalized["client_id"] - if totp_locked(client_id): + if totp_locked(client_id, ip): return _authorize_failed( request, normalized, "Too many attempts. Try again in 5 minutes.", @@ -1643,11 +1643,10 @@ async def oauth_authorize_post(request: Request) -> Response: # A replayed code is the operator re-submitting one they already # spent, not an attack -- it must not count towards the lockout. if totp_result is TotpResult.INVALID: - totp_record_failure(client_id) + totp_record_failure(client_id, ip) audit.emit( "auth.authorize.fail", client_id=client_id, - reason=totp_result.value, - ip=client_ip(request, tuple(config.server.trusted_proxies)), + reason=totp_result.value, ip=ip, ) return _authorize_failed( request, normalized, @@ -1656,13 +1655,10 @@ async def oauth_authorize_post(request: Request) -> Response: if totp_result is TotpResult.REPLAY else "Incorrect code. Check that your device clock is in sync." ), - locked=totp_locked(client_id), + locked=totp_locked(client_id, ip), ) - totp_record_success(client_id) - audit.emit( - "auth.authorize.2fa", client_id=client_id, - ip=client_ip(request, tuple(config.server.trusted_proxies)), - ) + totp_record_success(client_id, ip) + audit.emit("auth.authorize.2fa", client_id=client_id, ip=ip) return _authorize_approved(request, normalized) async def oauth_authorize_finalize(request: Request) -> Response: @@ -1726,11 +1722,9 @@ async def oauth_passkey_options(request: Request) -> Response: return JSONResponse( {"ok": False, "error": "Unknown client."}, status_code=400, ) - if totp_locked(client_id): - return JSONResponse( - {"ok": False, "error": "Too many attempts. Try again in 5 minutes."}, - status_code=429, - ) + # No TOTP-lockout check on the passkey path: a passkey cannot be + # guessed, so it stays the owner's way in while someone burns the + # client's TOTP budget. The per-IP guard above still applies. try: options, state = passkey_service.authentication_options( request, client_id=_passkey_owner(client_id), @@ -1759,11 +1753,6 @@ async def oauth_passkey_verify(request: Request) -> Response: status_code=400, ) client_id = normalized["client_id"] - if totp_locked(client_id): - return JSONResponse( - {"ok": False, "error": "Too many attempts. Try again in 5 minutes."}, - status_code=429, - ) credential = body.get("credential") state = str(body.get("state") or "") if not isinstance(credential, dict) or not state: @@ -1789,7 +1778,7 @@ async def oauth_passkey_verify(request: Request) -> Response: {"ok": False, "error": "This passkey belongs to another client."}, status_code=401, ) - totp_record_success(client_id) + totp_record_success(client_id, _ip(request)) audit.emit( "auth.authorize.2fa", client_id=client_id, via="passkey", ip=client_ip(request, tuple(config.server.trusted_proxies)), @@ -1882,7 +1871,7 @@ async def oauth_token(request: Request) -> Response: # 2FA mandatory on every client_credentials exchange (design choice: # no non-interactive escape hatch, the operator must re-type a TOTP # code at every 24 h token refresh). - if totp_locked(client_id): + if totp_locked(client_id, ip): return JSONResponse( { "error": "invalid_grant", @@ -1895,7 +1884,7 @@ async def oauth_token(request: Request) -> Response: if totp_result is not TotpResult.OK: # Replays don't count towards the lockout -- see /authorize. if totp_result is TotpResult.INVALID: - totp_record_failure(client_id) + totp_record_failure(client_id, ip) return JSONResponse( { "error": "invalid_grant", @@ -1907,7 +1896,7 @@ async def oauth_token(request: Request) -> Response: }, status_code=400, ) - totp_record_success(client_id) + totp_record_success(client_id, ip) # No refresh token here: every client_credentials exchange must # carry a fresh TOTP code, a refresh token would bypass that. grant = token_store.issue_oauth(client_id) diff --git a/src/beaconmcp/dashboard/app.py b/src/beaconmcp/dashboard/app.py index 46eab7b..f34ac6a 100644 --- a/src/beaconmcp/dashboard/app.py +++ b/src/beaconmcp/dashboard/app.py @@ -75,9 +75,10 @@ class DashboardDeps: session_store: SessionStore client_store: object # beaconmcp.auth.ClientStore -- avoid import cycle token_store: object # beaconmcp.auth.TokenStore - totp_locked: Callable[[str], bool] - totp_record_failure: Callable[[str], None] - totp_record_success: Callable[[str], None] + # TOTP lockout, keyed on (client_id, source IP): see ratelimit.TotpLockout. + totp_locked: Callable[[str, str], bool] + totp_record_failure: Callable[[str, str], None] + totp_record_success: Callable[[str, str], None] conversations: ConversationStore | None = None engine: ChatEngine | None = None confirmations: ConfirmationStore | None = None @@ -252,7 +253,15 @@ def _bearer_live(deps: DashboardDeps, session: Session) -> bool: return True -def _check_totp(deps: DashboardDeps, client_id: str, code: str) -> TotpResult: +def _source_ip(request: Request, deps: DashboardDeps) -> str: + """Client IP used to key the TOTP lockout (honours trusted proxies).""" + from ..ratelimit import client_ip # local import: avoid cycle at module load + return client_ip(request, deps.trusted_proxies) + + +def _check_totp( + deps: DashboardDeps, client_id: str, code: str, request: Request, +) -> TotpResult: """Check a 2FA code, recording a lockout strike only for a wrong one. The dashboard prompts for 2FA on several actions (login, minting a token, @@ -263,7 +272,7 @@ def _check_totp(deps: DashboardDeps, client_id: str, code: str) -> TotpResult: """ result = deps.client_store.check_totp(client_id, code) # type: ignore[attr-defined] if result is TotpResult.INVALID: - deps.totp_record_failure(client_id) + deps.totp_record_failure(client_id, _source_ip(request, deps)) return result @@ -508,14 +517,14 @@ def _fail(message: str, *, status: int = 400, locked: bool = False) -> Response: if not deps.client_store.verify(client_id, client_secret): # type: ignore[attr-defined] return _fail("Invalid credentials.", status=401) - if deps.totp_locked(client_id): + if deps.totp_locked(client_id, _source_ip(request, deps)): return _fail( "Too many attempts. Try again in 5 minutes.", status=429, locked=True, ) - totp_result = _check_totp(deps, client_id, totp) + totp_result = _check_totp(deps, client_id, totp, request) if totp_result is not TotpResult.OK: audit.emit( "dashboard.login.fail", client_id=client_id, @@ -527,10 +536,10 @@ def _fail(message: str, *, status: int = 400, locked: bool = False) -> Response: "Invalid 2FA code. Check that your device clock is in sync.", ), status=401, - locked=deps.totp_locked(client_id), + locked=deps.totp_locked(client_id, _source_ip(request, deps)), ) - deps.totp_record_success(client_id) + deps.totp_record_success(client_id, _source_ip(request, deps)) audit.emit("dashboard.login.ok", client_id=client_id) _, _, response = _start_session( request, @@ -558,7 +567,7 @@ async def refresh_get(request: Request) -> Response: client_name=client_name, next=request.query_params.get("next", ""), banner=None, - locked=deps.totp_locked(session.client_id), + locked=deps.totp_locked(session.client_id, _source_ip(request, deps)), ) async def refresh_post(request: Request) -> Response: @@ -597,21 +606,21 @@ def _fail(message: str, *, status: int = 400, locked: bool = False) -> Response: if not totp: return _fail("2FA code is required.") - if deps.totp_locked(session.client_id): + if deps.totp_locked(session.client_id, _source_ip(request, deps)): return _fail( "Too many attempts. Try again in 5 minutes.", status=429, locked=True, ) - totp_result = _check_totp(deps, session.client_id, totp) + totp_result = _check_totp(deps, session.client_id, totp, request) if totp_result is not TotpResult.OK: return _fail( _totp_error(totp_result, "Invalid 2FA code."), status=401, - locked=deps.totp_locked(session.client_id), + locked=deps.totp_locked(session.client_id, _source_ip(request, deps)), ) - deps.totp_record_success(session.client_id) + deps.totp_record_success(session.client_id, _source_ip(request, deps)) # Re-issue MCP bearer using the stored client_secret. secret = deps.session_store.get_client_secret(session.session_id) @@ -716,20 +725,20 @@ async def tokens_create(request: Request) -> Response: form_error="2FA code is required.", form_name=name, ) - if deps.totp_locked(session.client_id): + if deps.totp_locked(session.client_id, _source_ip(request, deps)): return _render_tokens_page( request, session, deps, form_error="Too many 2FA attempts; try again in 5 minutes.", form_name=name, ) - totp_result = _check_totp(deps, session.client_id, totp) + totp_result = _check_totp(deps, session.client_id, totp, request) if totp_result is not TotpResult.OK: return _render_tokens_page( request, session, deps, form_error=_totp_error(totp_result, "Invalid 2FA code."), form_name=name, ) - deps.totp_record_success(session.client_id) + deps.totp_record_success(session.client_id, _source_ip(request, deps)) try: token, ttl = deps.token_store.issue( # type: ignore[attr-defined] @@ -803,7 +812,7 @@ def _render_connectors_page( form_error=form_error, form_label=form_label, just_created=just_created, - locked=deps.totp_locked(session.client_id), + locked=deps.totp_locked(session.client_id, _source_ip(request, deps)), ) async def connectors_get(request: Request) -> Response: @@ -847,20 +856,20 @@ async def connectors_mint(request: Request) -> Response: form_error="Label cannot exceed 60 characters.", form_label=label, ) - if deps.totp_locked(session.client_id): + if deps.totp_locked(session.client_id, _source_ip(request, deps)): return _render_connectors_page( request, session, form_error="Too many 2FA attempts; try again in 5 minutes.", form_label=label, ) - totp_result = _check_totp(deps, session.client_id, totp_code) + totp_result = _check_totp(deps, session.client_id, totp_code, request) if totp_result is not TotpResult.OK: return _render_connectors_page( request, session, form_error=_totp_error(totp_result, "Invalid 2FA code."), form_label=label, ) - deps.totp_record_success(session.client_id) + deps.totp_record_success(session.client_id, _source_ip(request, deps)) store = deps.dyn_reg assert store is not None @@ -1153,7 +1162,7 @@ async def api_passkeys_auth_verify(request: Request) -> Response: next_url = _default_landing() # A passkey sign-in is a clean second factor: clear any TOTP # lockout the same way a correct code would. - deps.totp_record_success(client_id) + deps.totp_record_success(client_id, _source_ip(request, deps)) audit.emit("dashboard.login.ok", client_id=client_id, via="passkey") _, _, response = _start_session( request, @@ -1220,18 +1229,18 @@ async def api_update_apply(request: Request) -> Response: totp = str(body.get("totp") or "").strip() if not totp: return _json({"ok": False, "error": "2FA code is required."}, status=400) - if deps.totp_locked(session.client_id): + if deps.totp_locked(session.client_id, _source_ip(request, deps)): return _json( {"ok": False, "error": "Too many 2FA attempts; try again in 5 minutes."}, status=429, ) - totp_result = _check_totp(deps, session.client_id, totp) + totp_result = _check_totp(deps, session.client_id, totp, request) if totp_result is not TotpResult.OK: return _json( {"ok": False, "error": _totp_error(totp_result, "Invalid 2FA code.")}, status=401, ) - deps.totp_record_success(session.client_id) + deps.totp_record_success(session.client_id, _source_ip(request, deps)) from .. import updates as updates_mod @@ -1889,7 +1898,7 @@ def _expires_label(expires_at: float) -> str: form_name=form_name, just_created=just_created, mcp_url=mcp_url, - locked=deps.totp_locked(session.client_id), + locked=deps.totp_locked(session.client_id, _source_ip(request, deps)), chat_enabled=deps.engine is not None, dcr_enabled=deps.dyn_reg is not None, passkeys_enabled=_passkeys_ready(deps), diff --git a/src/beaconmcp/ratelimit.py b/src/beaconmcp/ratelimit.py index 9e48640..27845de 100644 --- a/src/beaconmcp/ratelimit.py +++ b/src/beaconmcp/ratelimit.py @@ -15,6 +15,7 @@ from __future__ import annotations import ipaddress +import logging import threading import time from collections import deque @@ -90,6 +91,77 @@ def retry_after(self, key: str) -> int: return max(0, int(self._window - (time.monotonic() - oldest)) + 1) +class TotpLockout: + """TOTP brute-force lockout keyed on ``(client_id, source IP)``. + + ``client_id`` is not a secret (it travels in authorize URLs, browser + history, proxy logs), so a lockout keyed on it alone let anyone who knew + it lock the owner out of every 2FA path with a handful of wrong codes. + Strikes are counted per source instead: ``per_source`` wrong codes from + one IP lock that IP out for that client. A second, higher per-client cap + (``per_client`` wrong codes from any mix of IPs) still stops a guesser + who spreads attempts over many addresses. + + Every failure pushes the cooldown of its counters forward; a counter + whose cooldown has passed starts again from zero. A success clears the + caller's own per-source counter only, so a distributed guesser can't + reset the per-client cap by having the owner log in. + """ + + def __init__( + self, + *, + per_source: int = 5, + per_client: int = 10, + lockout_seconds: float = 300.0, + ) -> None: + self._per_source = per_source + self._per_client = per_client + self._lockout = lockout_seconds + # key -> (failures, cooldown_until). Keys are ("src", client_id, ip) + # and ("client", client_id). + self._strikes: dict[tuple[str, ...], tuple[int, float]] = {} + self._lock = threading.Lock() + + @staticmethod + def _src_key(client_id: str, ip: str | None) -> tuple[str, ...]: + return ("src", client_id, ip or "-") + + def _count_locked(self, key: tuple[str, ...], now: float) -> int: + entry = self._strikes.get(key) + if entry is None: + return 0 + count, until = entry + if now >= until: + del self._strikes[key] + return 0 + return count + + def locked(self, client_id: str, ip: str | None = None) -> bool: + now = time.time() + with self._lock: + return ( + self._count_locked(("client", client_id), now) >= self._per_client + or self._count_locked(self._src_key(client_id, ip), now) + >= self._per_source + ) + + def record_failure(self, client_id: str, ip: str | None = None) -> None: + now = time.time() + until = now + self._lockout + with self._lock: + for key in (("client", client_id), self._src_key(client_id, ip)): + self._strikes[key] = (self._count_locked(key, now) + 1, until) + if len(self._strikes) > 4096: + for key, (_, key_until) in list(self._strikes.items()): + if now >= key_until: + del self._strikes[key] + + def record_success(self, client_id: str, ip: str | None = None) -> None: + with self._lock: + self._strikes.pop(self._src_key(client_id, ip), None) + + def _coerce_ip(value: object) -> str | None: raw = str(value).strip() if not raw: @@ -138,6 +210,29 @@ def _is_trusted_proxy(ip_value: str, trusted_proxies: tuple[str, ...]) -> bool: return False +_forwarding_warned = False + + +def _warn_untrusted_forwarding(forwarded_for: str | None, peer: str | None) -> None: + """Say once that a proxy is forwarding but ``trusted_proxies`` is empty. + + Behind a reverse proxy every request then comes from the proxy's own + address, so the per-IP limits (login, authorize, TOTP lockout) become + one budget shared by every visitor, and one attacker can spend it for + everyone. + """ + global _forwarding_warned + if _forwarding_warned or not forwarded_for: + return + _forwarding_warned = True + logging.getLogger("beaconmcp").warning( + "X-Forwarded-For received from %s but server.trusted_proxies is empty: " + "every client behind that proxy shares one rate-limit and TOTP-lockout " + "budget. Add the proxy's address to server.trusted_proxies.", + peer or "an unknown peer", + ) + + def client_ip(request: object, trusted_proxies: tuple[str, ...] = ()) -> str: """Best-effort client IP for a Starlette ``Request``. @@ -151,6 +246,8 @@ def client_ip(request: object, trusted_proxies: tuple[str, ...] = ()) -> str: direct_ip = _coerce_ip(direct_peer_raw) headers = getattr(request, "headers", None) + if not trusted_proxies and headers is not None and hasattr(headers, "get"): + _warn_untrusted_forwarding(headers.get("x-forwarded-for"), direct_ip) if headers is not None and direct_ip and _is_trusted_proxy(direct_ip, trusted_proxies): fwd = headers.get("x-forwarded-for") if hasattr(headers, "get") else None if fwd: diff --git a/tests/test_dashboard_chat.py b/tests/test_dashboard_chat.py index 351374d..562cdbf 100644 --- a/tests/test_dashboard_chat.py +++ b/tests/test_dashboard_chat.py @@ -124,9 +124,9 @@ def deps(tmp_path, engine): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, confirmations=ConfirmationStore(), @@ -374,9 +374,9 @@ def validate(self, token): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=WipedTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, ) @@ -404,9 +404,9 @@ def validate(self, token): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=WipedTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, ) @@ -442,9 +442,9 @@ def validate(self, token): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=token_store, - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, ) @@ -792,9 +792,9 @@ def test_chat_stream_remote_mode_still_routes_public_url(tmp_path, engine): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, mcp_public_url="https://mcp.example/", diff --git a/tests/test_dashboard_integration.py b/tests/test_dashboard_integration.py index ea0ab99..4f13ec7 100644 --- a/tests/test_dashboard_integration.py +++ b/tests/test_dashboard_integration.py @@ -91,18 +91,18 @@ def deps(tmp_path, monkeypatch): session_store = SessionStore(db, key=os.urandom(32)) failures: dict[str, tuple[int, float]] = {} - def totp_locked(cid): + def totp_locked(cid, ip=None): e = failures.get(cid) if not e: return False count, until = e return count >= 5 and time.time() < until - def totp_record_failure(cid): + def totp_record_failure(cid, ip=None): c, _ = failures.get(cid, (0, 0.0)) failures[cid] = (c + 1, time.time() + 300) - def totp_record_success(cid): + def totp_record_success(cid, ip=None): failures.pop(cid, None) # A sentinel non-None engine so the post-login landing stays /app/chat. @@ -143,15 +143,15 @@ def tokens_only_client(tmp_path): session_store=session_store, client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: ( + totp_locked=lambda cid, ip=None: ( failures.get(cid, (0, 0.0))[0] >= 5 and time.time() < failures.get(cid, (0, 0.0))[1] ), - totp_record_failure=lambda cid: failures.__setitem__( + totp_record_failure=lambda cid, ip=None: failures.__setitem__( cid, (failures.get(cid, (0, 0.0))[0] + 1, time.time() + 300), ), - totp_record_success=lambda cid: (failures.pop(cid, None), None)[1], + totp_record_success=lambda cid, ip=None: (failures.pop(cid, None), None)[1], engine=None, ) app = Starlette(routes=build_dashboard_routes(deps_local)) diff --git a/tests/test_dashboard_mcp_apps.py b/tests/test_dashboard_mcp_apps.py index 86722d9..9dba410 100644 --- a/tests/test_dashboard_mcp_apps.py +++ b/tests/test_dashboard_mcp_apps.py @@ -291,9 +291,9 @@ def deps(tmp_path, engine): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, confirmations=ConfirmationStore(), diff --git a/tests/test_dashboard_usage.py b/tests/test_dashboard_usage.py index acbf1b0..4be2c60 100644 --- a/tests/test_dashboard_usage.py +++ b/tests/test_dashboard_usage.py @@ -308,9 +308,9 @@ def _build_deps(tmp_path, *, engine, budget): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=engine, confirmations=ConfirmationStore(), @@ -415,9 +415,9 @@ def test_usage_endpoint_returns_disabled_snapshot_when_store_missing(tmp_path): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, conversations=ConversationStore(db), engine=FakeChatEngine(FakeScript()), confirmations=ConfirmationStore(), diff --git a/tests/test_passkeys.py b/tests/test_passkeys.py index 06279ac..579b412 100644 --- a/tests/test_passkeys.py +++ b/tests/test_passkeys.py @@ -516,9 +516,9 @@ def web(tmp_path, monkeypatch): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, passkeys=PasskeyService(passkey_store), ) app = Starlette(routes=build_dashboard_routes(deps)) @@ -768,9 +768,9 @@ def test_routes_report_unavailable_without_a_service(tmp_path, monkeypatch): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, passkeys=None, ) client = TestClient( @@ -795,9 +795,9 @@ def test_login_page_hides_passkeys_when_unavailable(tmp_path, monkeypatch): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, passkeys=None, ) client = TestClient( diff --git a/tests/test_totp_lockout.py b/tests/test_totp_lockout.py new file mode 100644 index 0000000..d43ad38 --- /dev/null +++ b/tests/test_totp_lockout.py @@ -0,0 +1,157 @@ +"""TOTP lockout keyed on (client_id, source IP) with a per-client cap. + +Keyed on client_id alone, anyone who knew a client_id (not a secret: it is in +authorize URLs and proxy logs) could lock the owner out of the dashboard, +token issuance and /oauth/authorize with five wrong codes. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from beaconmcp import ratelimit # noqa: E402 +from beaconmcp.ratelimit import TotpLockout # noqa: E402 + +ATTACKER = "203.0.113.7" +OWNER = "198.51.100.1" + + +class Clock: + def __init__(self) -> None: + self.now = 1_000_000.0 + + def time(self) -> float: + return self.now + + +@pytest.fixture() +def clock(monkeypatch): + c = Clock() + monkeypatch.setattr(ratelimit.time, "time", c.time) + return c + + +def test_failures_from_one_ip_do_not_lock_another(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20, lockout_seconds=300) + for _ in range(5): + guard.record_failure("c1", ATTACKER) + assert guard.locked("c1", ATTACKER) + assert not guard.locked("c1", OWNER) + + +def test_source_lock_is_per_client(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + for _ in range(5): + guard.record_failure("c1", ATTACKER) + assert not guard.locked("c2", ATTACKER) + + +def test_per_client_cap_stops_distributed_guessing(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + # 4 wrong codes from each of 5 addresses: no single source is locked... + for i in range(5): + for _ in range(4): + guard.record_failure("c1", f"203.0.113.{i}") + # ...but the client as a whole is, from everywhere. + assert guard.locked("c1", "203.0.113.0") + assert guard.locked("c1", OWNER) + + +def test_below_thresholds_nothing_is_locked(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + for _ in range(4): + guard.record_failure("c1", ATTACKER) + assert not guard.locked("c1", ATTACKER) + + +def test_lock_expires_after_cooldown(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20, lockout_seconds=300) + for _ in range(5): + guard.record_failure("c1", ATTACKER) + clock.now += 299 + assert guard.locked("c1", ATTACKER) + clock.now += 2 + assert not guard.locked("c1", ATTACKER) + + +def test_old_strikes_do_not_accumulate_across_cooldowns(clock) -> None: + # Four wrong codes today and one tomorrow must not lock. + guard = TotpLockout(per_source=5, per_client=20, lockout_seconds=300) + for _ in range(4): + guard.record_failure("c1", ATTACKER) + clock.now += 86_400 + guard.record_failure("c1", ATTACKER) + assert not guard.locked("c1", ATTACKER) + + +def test_success_clears_only_the_callers_source(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + for _ in range(4): + guard.record_failure("c1", ATTACKER) + for _ in range(4): + guard.record_failure("c1", OWNER) + guard.record_success("c1", OWNER) + guard.record_failure("c1", OWNER) + assert not guard.locked("c1", OWNER) + guard.record_failure("c1", ATTACKER) + assert guard.locked("c1", ATTACKER) + + +def test_success_does_not_reset_the_per_client_cap(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + for i in range(5): + for _ in range(4): + guard.record_failure("c1", f"203.0.113.{i}") + guard.record_success("c1", OWNER) + assert guard.locked("c1", OWNER) + + +def test_unknown_ip_is_its_own_source(clock) -> None: + guard = TotpLockout(per_source=5, per_client=20) + for _ in range(5): + guard.record_failure("c1") + assert guard.locked("c1") + assert not guard.locked("c1", OWNER) + + +def test_default_per_client_cap_is_ten(clock) -> None: + """Two addresses at the per-source limit are enough to hit the client cap.""" + guard = TotpLockout() + for ip in ("203.0.113.1", "203.0.113.2"): + for _ in range(5): + guard.record_failure("c1", ip) + assert guard.locked("c1", OWNER) + + +def test_forwarded_header_without_trusted_proxies_warns_once(caplog, monkeypatch) -> None: + from types import SimpleNamespace + + monkeypatch.setattr(ratelimit, "_forwarding_warned", False) + request = SimpleNamespace( + client=SimpleNamespace(host="10.0.0.2"), + headers={"x-forwarded-for": ATTACKER}, + ) + with caplog.at_level("WARNING", logger="beaconmcp"): + assert ratelimit.client_ip(request) == "10.0.0.2" + ratelimit.client_ip(request) + warnings = [r for r in caplog.records if "trusted_proxies is empty" in r.getMessage()] + assert len(warnings) == 1 + assert "10.0.0.2" in warnings[0].getMessage() + + +def test_no_warning_when_the_proxy_is_trusted(caplog, monkeypatch) -> None: + from types import SimpleNamespace + + monkeypatch.setattr(ratelimit, "_forwarding_warned", False) + request = SimpleNamespace( + client=SimpleNamespace(host="10.0.0.2"), + headers={"x-forwarded-for": ATTACKER}, + ) + with caplog.at_level("WARNING", logger="beaconmcp"): + assert ratelimit.client_ip(request, ("10.0.0.2",)) == ATTACKER + assert not [r for r in caplog.records if "trusted_proxies" in r.getMessage()] diff --git a/tests/test_updates.py b/tests/test_updates.py index 655217a..223816a 100644 --- a/tests/test_updates.py +++ b/tests/test_updates.py @@ -613,9 +613,9 @@ def _make_client(tmp_path, monkeypatch, **overrides): session_store=SessionStore(db, key=os.urandom(32)), client_store=FakeClientStore(), token_store=FakeTokenStore(), - totp_locked=lambda cid: False, - totp_record_failure=lambda cid: None, - totp_record_success=lambda cid: None, + totp_locked=lambda cid, ip=None: False, + totp_record_failure=lambda cid, ip=None: None, + totp_record_success=lambda cid, ip=None: None, **overrides, ) client = TestClient(