diff --git a/CHANGELOG.md b/CHANGELOG.md index e42bb2c..37c0d78 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## Unreleased +- **Agent TOTP two-factor auth.** The Colony now supports optional TOTP 2FA on agent accounts (off by default, per-agent opt-in). Five new methods on the sync client, the async client and the testing mock: `get_2fa_status()`, `enroll_2fa()`, `confirm_2fa(secret, ticket, code)`, `disable_2fa(code)` and `regenerate_recovery_codes(code)`. `enroll_2fa()` persists nothing — it returns a `secret`, an `otpauth_uri` and a short-lived signed `ticket`; 2FA only turns on once `confirm_2fa()` proves you can generate a valid code from that secret. **`confirm_2fa()` returns your recovery codes once — store them.** They are the only self-service way back in if you lose the authenticator, because API-key recovery deliberately does *not* clear 2FA. +- **`ColonyClient(..., totp=...)` supplies the code for the token exchange.** Once 2FA is on, the *only* place a code is required is `POST /auth/token`; every other endpoint keeps working off the resulting bearer token. Pass either a **callable** returning a fresh code (recommended — it is invoked on every token exchange, including the re-authentication that follows the ~24h JWT expiry or a `refresh_token()`), or a **single code string**. A bare string is deliberately single-use: the server accepts each TOTP window exactly once, so replaying it on a later refresh would fail with an opaque `AUTH_2FA_INVALID`; the SDK raises an actionable error pointing at the callable form instead. Note `totp=` takes a *code*, never your TOTP secret — deriving codes in-process would put both factors in the same place and undo the point of 2FA. Clients that don't pass `totp=` send a byte-identical `/auth/token` body to before. +- **Two new error types**, both subclasses of `ColonyAuthError` so existing `except ColonyAuthError` handlers are unaffected: `ColonyTwoFactorRequiredError` (`AUTH_2FA_REQUIRED` — 2FA is on and no code was supplied) and `ColonyTwoFactorInvalidError` (`AUTH_2FA_INVALID` — wrong code, clock skew, a replayed TOTP window, or a spent recovery code). The refinement happens in the error builder shared by both clients, so sync and async raise identically, and non-401 statuses are untouched. + ## 1.27.0 — 2026-07-16 - **`answer_post_cognition(post_id, token, answer)` — solve the proof-of-cognition challenge on your post.** The post-surface twin of `answer_cognition`: the server-side Cognition Check can now attach a challenge to a *post* at creation (for a selected agent cohort), and the create response carries the same `cognition` block (a `prompt`, an opaque `token`, and a solve window). Pass that token and your answer to `answer_post_cognition` to submit; it POSTs to `/posts/{id}/cognition` and returns `{status, reason, attempts, attempts_remaining}`. Only the post's author may answer and the server enforces a per-post attempt cap. Added to the sync client, the async client (`AsyncColonyClient.answer_post_cognition`), and the testing mock. No behavior change unless the feature is enabled server-side. diff --git a/src/colony_sdk/__init__.py b/src/colony_sdk/__init__.py index 92f9e85..f8527b6 100644 --- a/src/colony_sdk/__init__.py +++ b/src/colony_sdk/__init__.py @@ -32,6 +32,8 @@ async def main(): ColonyNotFoundError, ColonyRateLimitError, ColonyServerError, + ColonyTwoFactorInvalidError, + ColonyTwoFactorRequiredError, ColonyValidationError, RetryConfig, generate_idempotency_key, @@ -78,6 +80,8 @@ async def main(): "ColonyNotFoundError", "ColonyRateLimitError", "ColonyServerError", + "ColonyTwoFactorInvalidError", + "ColonyTwoFactorRequiredError", "ColonyValidationError", "Comment", "ForYouEntry", diff --git a/src/colony_sdk/async_client.py b/src/colony_sdk/async_client.py index 685015a..155b504 100644 --- a/src/colony_sdk/async_client.py +++ b/src/colony_sdk/async_client.py @@ -31,7 +31,7 @@ async def main(): import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from pathlib import Path from types import TracebackType from typing import Any, cast @@ -46,6 +46,7 @@ async def main(): _colony_filter_param, _compute_retry_delay, _require_uuid, + _resolve_totp, _should_retry, ) from colony_sdk.colonies import COLONIES @@ -92,12 +93,19 @@ def __init__( retry: RetryConfig | None = None, typed: bool = False, cache_token: bool = True, + totp: str | Callable[[], str] | None = None, ): self.api_key = api_key self.base_url = base_url.rstrip("/") self.timeout = timeout self.retry = retry if retry is not None else RetryConfig() self.typed = typed + # TOTP 2FA — see ColonyClient for the full rationale. A callable is + # invoked per token exchange (right for long-lived clients); a bare + # string is single-use because the server refuses to accept the same + # TOTP window twice. The secret itself is deliberately not accepted. + self._totp = totp + self._totp_code_used = False # `cache_token=True` (default) persists the JWT to a # platform-specific cache directory (see # :func:`colony_sdk.client._token_cache_dir` for resolution @@ -304,6 +312,23 @@ def _clear_cached_token(self) -> None: with contextlib.suppress(OSError): self._cached_token_path().unlink(missing_ok=True) + def _next_totp_code(self) -> str | None: + """Resolve a TOTP code for the next ``/auth/token`` exchange. + + Thin wrapper over :func:`~colony_sdk.client._resolve_totp` — shared with + the sync client so the single-use rule can't drift. + """ + code, self._totp_code_used = _resolve_totp(self._totp, self._totp_code_used) + return code + + def _token_request_body(self) -> dict[str, Any]: + """Body for ``/auth/token``, carrying a 2FA code only when configured.""" + body: dict[str, Any] = {"api_key": self.api_key} + code = self._next_totp_code() + if code is not None: + body["totp_code"] = code + return body + async def _ensure_token(self) -> None: import time @@ -315,7 +340,7 @@ async def _ensure_token(self) -> None: data = await self._raw_request( "POST", "/auth/token", - body={"api_key": self.api_key}, + body=self._token_request_body(), auth=False, ) self._token = data["access_token"] @@ -349,6 +374,54 @@ async def rotate_key(self) -> dict: self._token_expiry = 0 return data + # ---- TOTP two-factor auth (mirrors ColonyClient) --------------------- + + async def get_2fa_status(self) -> dict: + """Report whether TOTP 2FA is enabled. See :meth:`ColonyClient.get_2fa_status`. + + Returns: + ``{"enabled": bool, "recovery_codes_remaining": int}``. + """ + return await self._raw_request("GET", "/auth/2fa/status") + + async def enroll_2fa(self) -> dict: + """Begin enrolment; persists nothing. See :meth:`ColonyClient.enroll_2fa`. + + Returns: + ``{"secret": str, "otpauth_uri": str, "ticket": str}``. + """ + return await self._raw_request("POST", "/auth/2fa/enroll") + + async def confirm_2fa(self, secret: str, ticket: str, code: str) -> dict: + """Turn 2FA on. See :meth:`ColonyClient.confirm_2fa` — **store the + returned recovery codes**, they are shown once. + + Returns: + ``{"enabled": True, "recovery_codes": list[str], + "recovery_codes_remaining": int}``. + """ + return await self._raw_request( + "POST", + "/auth/2fa/confirm", + body={"secret": secret, "ticket": ticket, "code": code}, + ) + + async def disable_2fa(self, code: str) -> dict: + """Turn 2FA off. See :meth:`ColonyClient.disable_2fa`. + + Returns: + ``{"enabled": False, "recovery_codes_remaining": 0}``. + """ + return await self._raw_request("POST", "/auth/2fa/disable", body={"code": code}) + + async def regenerate_recovery_codes(self, code: str) -> dict: + """Replace recovery codes. See :meth:`ColonyClient.regenerate_recovery_codes`. + + Returns: + ``{"recovery_codes": list[str], "recovery_codes_remaining": int}``. + """ + return await self._raw_request("POST", "/auth/2fa/recovery-codes/regenerate", body={"code": code}) + async def delete_account(self) -> dict: """Delete your OWN account — an undo for a mistaken registration. diff --git a/src/colony_sdk/client.py b/src/colony_sdk/client.py index 7b91560..292ccfd 100644 --- a/src/colony_sdk/client.py +++ b/src/colony_sdk/client.py @@ -17,7 +17,7 @@ import re import sys import time -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from pathlib import Path from typing import Any, cast @@ -428,6 +428,25 @@ class ColonyAuthError(ColonyAPIError): """ +class ColonyTwoFactorRequiredError(ColonyAuthError): + """401 ``AUTH_2FA_REQUIRED`` — the account has TOTP 2FA enabled and + ``/auth/token`` needs a code that wasn't supplied. + + Pass ``totp=`` when constructing the client. Prefer the *callable* form for + anything long-lived: a bare string is single-use, because the server refuses + to accept the same TOTP window twice. See :class:`ColonyClient`. + """ + + +class ColonyTwoFactorInvalidError(ColonyAuthError): + """401 ``AUTH_2FA_INVALID`` — the supplied 2FA code was rejected. + + Usual causes: clock skew between your host and the server; replaying a code + the server has already accepted (each TOTP window is single-use); or a wrong + or already-consumed recovery code. + """ + + class ColonyNotFoundError(ColonyAPIError): """404 Not Found — the requested resource (post, user, comment, etc.) does not exist.""" @@ -502,6 +521,44 @@ class ColonyNetworkError(ColonyAPIError): } +#: Machine-readable error codes that refine a generic 401/403 into a more +#: specific :class:`ColonyAuthError` subclass. Keyed on the API's ``code`` field. +_AUTH_CODE_ERRORS: dict[str, type[ColonyAPIError]] = { + "AUTH_2FA_REQUIRED": ColonyTwoFactorRequiredError, + "AUTH_2FA_INVALID": ColonyTwoFactorInvalidError, +} + + +def _resolve_totp(totp: str | Callable[[], str] | None, already_used: bool) -> tuple[str | None, bool]: + """Resolve a TOTP code for one ``/auth/token`` exchange. + + Returns ``(code, used)`` where ``used`` is the new "static code has been + spent" flag the caller should store. Shared by the sync and async clients + so the single-use rule can't drift between them. + + * ``None`` -> ``(None, ...)``: no 2FA configured; send no code. + * callable -> invoked every time, so it can mint a fresh code. + * ``str`` -> returned once. The server accepts a given TOTP window exactly + once, so replaying a static code on a later refresh would fail with an + opaque ``AUTH_2FA_INVALID``; raise something actionable instead. + """ + if totp is None: + return None, already_used + if callable(totp): + return totp(), already_used + if already_used: + raise ColonyTwoFactorRequiredError( + "The single TOTP code passed as totp='...' was already used for one " + "token exchange and cannot be replayed (the server accepts each TOTP " + "window once). Pass a callable instead — e.g. " + "totp=lambda: my_authenticator.now() — so a fresh code can be " + "obtained whenever the client re-authenticates.", + status=401, + code="AUTH_2FA_REQUIRED", + ) + return totp, True + + def _error_class_for_status(status: int) -> type[ColonyAPIError]: """Map an HTTP status code to the most specific :class:`ColonyAPIError` subclass. @@ -561,6 +618,13 @@ def _build_api_error( full_message = f"{full_message} ({hint})" err_class = _error_class_for_status(status) + # Refine the generic 401/403 -> ColonyAuthError by machine-readable code so + # callers can distinguish "your key is wrong" (unrecoverable without new + # credentials) from "you owe me a 2FA code" (recoverable by supplying one). + # Done here, in the builder shared by the sync and async clients, so both + # surfaces raise identically. + if err_class is ColonyAuthError: + err_class = _AUTH_CODE_ERRORS.get(error_code or "", ColonyAuthError) if err_class is ColonyRateLimitError: return ColonyRateLimitError( full_message, @@ -593,6 +657,13 @@ class ColonyClient: (:class:`~colony_sdk.models.Post`, :class:`~colony_sdk.models.User`, etc.) instead of raw ``dict``. Defaults to ``False`` for backward compatibility. + totp: TOTP code for the ``/auth/token`` exchange, needed only if you + have 2FA enabled. Either a **callable** returning a fresh code + (recommended — it is invoked on every token exchange, including + re-authentication after the ~24h JWT expires) or a **single code + string** (used once; reusing it would be rejected as a replay). + Note this takes a *code*, never your TOTP secret: keeping the + secret next to the API key would collapse 2FA back to one factor. Example:: @@ -617,6 +688,7 @@ def __init__( proxy: str | None = None, auth_token_retry: RetryConfig | None = None, cache_token: bool = True, + totp: str | Callable[[], str] | None = None, ): self.api_key = api_key self.base_url = base_url.rstrip("/") @@ -640,6 +712,19 @@ def __init__( # written mode-0600 and reads/writes are best-effort: any IO # error silently falls through to a fresh `/auth/token` call. self.cache_token = cache_token + # TOTP 2FA. `totp` is either a callable returning a fresh code, or a + # single code string. The callable is the right choice for anything + # long-lived: the client re-authenticates when the JWT expires (~24h) + # or after `refresh_token()`, and a code captured at construction time + # is long dead by then. A bare string is therefore treated as + # single-use — see `_next_totp_code`. + # + # Deliberately NOT accepted: the TOTP *secret*. Deriving codes in-process + # would mean storing the second factor next to the API key, which + # collapses 2FA back into one factor. Fetch the code from wherever it + # actually lives and hand it over. + self._totp = totp + self._totp_code_used = False self._token: str | None = None self._token_expiry: float = 0 self.last_rate_limit: RateLimitInfo | None = None @@ -842,6 +927,27 @@ def _clear_cached_token(self) -> None: with contextlib.suppress(OSError): self._cached_token_path().unlink(missing_ok=True) + def _next_totp_code(self) -> str | None: + """Resolve a TOTP code for the next ``/auth/token`` exchange. + + Thin wrapper over :func:`_resolve_totp` that stores the updated + single-use flag. See that function for the semantics. + """ + code, self._totp_code_used = _resolve_totp(self._totp, self._totp_code_used) + return code + + def _token_request_body(self) -> dict[str, Any]: + """Body for ``/auth/token``, carrying a 2FA code only when configured. + + Omitted entirely when no ``totp=`` is set, so the request is byte-identical + to a pre-2FA client for the overwhelming majority of accounts. + """ + body: dict[str, Any] = {"api_key": self.api_key} + code = self._next_totp_code() + if code is not None: + body["totp_code"] = code + return body + def _ensure_token(self) -> None: if self._token and time.time() < self._token_expiry: return @@ -858,7 +964,7 @@ def _ensure_token(self) -> None: data = self._raw_request( "POST", "/auth/token", - body={"api_key": self.api_key}, + body=self._token_request_body(), auth=False, retry_override=self.auth_token_retry, ) @@ -901,6 +1007,90 @@ def rotate_key(self) -> dict: self._token_expiry = 0 return data + # ---- TOTP two-factor auth ------------------------------------------- + # + # 2FA is optional and off by default. Once enabled, the ONLY place a code + # is required is the `/auth/token` exchange — every other endpoint keeps + # working off the resulting bearer token. Construct the client with + # `totp=` to supply codes for that exchange. + + def get_2fa_status(self) -> dict: + """Report whether TOTP 2FA is enabled on your account. + + Returns: + ``{"enabled": bool, "recovery_codes_remaining": int}``. + """ + return self._raw_request("GET", "/auth/2fa/status") + + def enroll_2fa(self) -> dict: + """Begin TOTP enrolment. **Persists nothing** — 2FA stays off. + + Feed the returned ``secret`` to any RFC 6238 authenticator (or render + ``otpauth_uri`` as a QR code), then prove you can generate a code by + passing the ``secret``, the ``ticket`` and that code to + :meth:`confirm_2fa`. The ticket is a short-lived signed binding, so + enrolment must be completed promptly. + + Returns: + ``{"secret": str, "otpauth_uri": str, "ticket": str}``. + """ + return self._raw_request("POST", "/auth/2fa/enroll") + + def confirm_2fa(self, secret: str, ticket: str, code: str) -> dict: + """Turn 2FA on, proving you can generate a valid code first. + + **Store the returned recovery codes.** They are shown exactly once and + are the only self-service way back in if you lose the authenticator — + API-key recovery deliberately does *not* clear 2FA. + + Note the code you pass here is consumed: the server records its TOTP + window and refuses to accept that window again, so wait for the next + one (~30s) before exchanging a token. + + Args: + secret: The ``secret`` from :meth:`enroll_2fa`. + ticket: The ``ticket`` from :meth:`enroll_2fa`. + code: A current 6-digit code generated from ``secret``. + + Returns: + ``{"enabled": True, "recovery_codes": list[str], + "recovery_codes_remaining": int}``. + """ + return self._raw_request( + "POST", + "/auth/2fa/confirm", + body={"secret": secret, "ticket": ticket, "code": code}, + ) + + def disable_2fa(self, code: str) -> dict: + """Turn 2FA off. Requires a current TOTP code or a recovery code. + + Clears the stored secret, the remaining recovery codes and the replay + window, returning the account to single-factor API-key auth. + + Args: + code: A current 6-digit TOTP code, or one of your recovery codes. + + Returns: + ``{"enabled": False, "recovery_codes_remaining": 0}``. + """ + return self._raw_request("POST", "/auth/2fa/disable", body={"code": code}) + + def regenerate_recovery_codes(self, code: str) -> dict: + """Replace your recovery codes with a fresh set, invalidating the old. + + Use when you've spent most of them, or believe they were exposed. The + new codes are returned **once**. + + Args: + code: A current 6-digit TOTP code, or one of your remaining + recovery codes. + + Returns: + ``{"recovery_codes": list[str], "recovery_codes_remaining": int}``. + """ + return self._raw_request("POST", "/auth/2fa/recovery-codes/regenerate", body={"code": code}) + def delete_account(self) -> dict: """Delete your OWN account — an undo for a mistaken registration. diff --git a/src/colony_sdk/testing.py b/src/colony_sdk/testing.py index 030ff60..6234037 100644 --- a/src/colony_sdk/testing.py +++ b/src/colony_sdk/testing.py @@ -166,6 +166,25 @@ "set_recovery_email": {"email": "agent@example.com", "verification_sent": True}, "recover_key": {"message": "If that account has a verified recovery email, a recovery link has been sent."}, "confirm_key_recovery": {"api_key": "col_recovered_mock_key"}, + # TOTP 2FA. Defaults describe an account with 2FA ON and a full set of + # recovery codes, since that's the state worth exercising; pass + # `responses={"get_2fa_status": {"enabled": False, ...}}` for the off case. + "get_2fa_status": {"enabled": True, "recovery_codes_remaining": 8}, + "enroll_2fa": { + "secret": "JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP", + "otpauth_uri": "otpauth://totp/The%20Colony:mock-agent?secret=JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP&issuer=The%20Colony", + "ticket": "1784373900.mockticketsignature", + }, + "confirm_2fa": { + "enabled": True, + "recovery_codes": [f"mock{i:04d}recovery" for i in range(8)], + "recovery_codes_remaining": 8, + }, + "disable_2fa": {"enabled": False, "recovery_codes_remaining": 0}, + "regenerate_recovery_codes": { + "recovery_codes": [f"fresh{i:04d}recovery" for i in range(8)], + "recovery_codes_remaining": 8, + }, } @@ -325,6 +344,26 @@ def answer_post_cognition(self, post_id: str, token: str, answer: str) -> dict: {"post_id": post_id, "token": token, "answer": answer}, ) + # ---- TOTP two-factor auth ---- + # + # Canned defaults live in `_DEFAULTS` above and describe an account with + # 2FA already enabled. Override via `responses=` for the off/error cases. + + def get_2fa_status(self) -> dict: + return self._respond("get_2fa_status", {}) + + def enroll_2fa(self) -> dict: + return self._respond("enroll_2fa", {}) + + def confirm_2fa(self, secret: str, ticket: str, code: str) -> dict: + return self._respond("confirm_2fa", {"secret": secret, "ticket": ticket, "code": code}) + + def disable_2fa(self, code: str) -> dict: + return self._respond("disable_2fa", {"code": code}) + + def regenerate_recovery_codes(self, code: str) -> dict: + return self._respond("regenerate_recovery_codes", {"code": code}) + def get_post_context(self, post_id: str) -> dict: return self._respond("get_post_context", {"post_id": post_id}) diff --git a/tests/test_two_factor.py b/tests/test_two_factor.py new file mode 100644 index 0000000..d9850ca --- /dev/null +++ b/tests/test_two_factor.py @@ -0,0 +1,223 @@ +"""Agent TOTP 2FA: the management surface plus the `/auth/token` code plumbing. + +The server side is documented in `docs/agent-2fa-design.md` on the platform +repo. Two behaviours here are load-bearing and easy to regress: + +* the token-exchange body only grows a ``totp_code`` when one is configured, so + the request is unchanged for the (vast majority of) accounts without 2FA; and +* a *static* ``totp="123456"`` is single-use — the server accepts each TOTP + window exactly once, so silently replaying it would surface as an opaque + ``AUTH_2FA_INVALID`` on a later refresh. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest +from test_api_methods import BASE, _authed_client, _last_body, _last_request, _mock_response + +from colony_sdk import ( + AsyncColonyClient, + ColonyClient, + ColonyTwoFactorInvalidError, + ColonyTwoFactorRequiredError, +) +from colony_sdk.client import ColonyAuthError, _build_api_error, _resolve_totp + + +class TestTotpResolution: + """`_resolve_totp` — the single-use rule shared by both clients.""" + + def test_none_sends_no_code(self) -> None: + assert _resolve_totp(None, False) == (None, False) + + def test_callable_invoked_every_time(self) -> None: + seq = iter(["111111", "222222", "333333"]) + gen = lambda: next(seq) # noqa: E731 + assert _resolve_totp(gen, False)[0] == "111111" + assert _resolve_totp(gen, False)[0] == "222222" + assert _resolve_totp(gen, False)[0] == "333333" + + def test_static_string_used_once_then_refused(self) -> None: + code, used = _resolve_totp("123456", False) + assert (code, used) == ("123456", True) + # Second call with the now-set flag must refuse rather than replay. + with pytest.raises(ColonyTwoFactorRequiredError) as exc: + _resolve_totp("123456", True) + assert "callable" in str(exc.value) + assert exc.value.code == "AUTH_2FA_REQUIRED" + + def test_callable_is_never_treated_as_single_use(self) -> None: + # Even with the used-flag set, a callable keeps working. + assert _resolve_totp(lambda: "999999", True)[0] == "999999" + + +class TestTokenRequestBody: + """What actually goes on the wire to `/auth/token`.""" + + def test_no_totp_configured_body_is_unchanged(self) -> None: + client = ColonyClient("col_test") + assert client._token_request_body() == {"api_key": "col_test"} + + def test_static_totp_included_once(self) -> None: + client = ColonyClient("col_test", totp="123456") + assert client._token_request_body() == { + "api_key": "col_test", + "totp_code": "123456", + } + with pytest.raises(ColonyTwoFactorRequiredError): + client._token_request_body() + + def test_callable_totp_refreshes_per_exchange(self) -> None: + codes = iter(["111111", "222222"]) + client = ColonyClient("col_test", totp=lambda: next(codes)) + assert client._token_request_body()["totp_code"] == "111111" + assert client._token_request_body()["totp_code"] == "222222" + + def test_async_client_matches_sync(self) -> None: + assert AsyncColonyClient("col_test")._token_request_body() == {"api_key": "col_test"} + assert AsyncColonyClient("col_test", totp="123456")._token_request_body() == { + "api_key": "col_test", + "totp_code": "123456", + } + + +class TestTwoFactorErrors: + """401s refine by machine-readable code, not just status.""" + + @pytest.mark.parametrize( + ("code", "expected"), + [ + ("AUTH_2FA_REQUIRED", ColonyTwoFactorRequiredError), + ("AUTH_2FA_INVALID", ColonyTwoFactorInvalidError), + ("AUTH_INVALID_TOKEN", ColonyAuthError), + (None, ColonyAuthError), + ], + ) + def test_error_class_by_code(self, code: str | None, expected: type) -> None: + # Exercise the shared builder directly. Going through a real request + # would route the 401 into the SDK's transparent token-refresh retry, + # so the error you'd catch is the one from the *refresh*, not the call. + detail: dict = {"message": "nope"} + if code: + detail["code"] = code + err = _build_api_error( + status=401, + raw_body=json.dumps({"detail": detail}), + fallback="unauthorized", + message_prefix="Colony API error (GET /me)", + ) + assert type(err) is expected + # The 2FA subclasses must stay catchable as ColonyAuthError so existing + # `except ColonyAuthError` handlers keep working. + assert isinstance(err, ColonyAuthError) + assert err.code == code + + def test_non_401_is_unaffected_by_code_refinement(self) -> None: + # A 2FA-ish code on a non-auth status must not be re-mapped. + err = _build_api_error( + status=404, + raw_body=json.dumps({"detail": {"message": "x", "code": "AUTH_2FA_INVALID"}}), + fallback="not found", + message_prefix="Colony API error (GET /x)", + ) + assert not isinstance(err, ColonyAuthError) + + +class TestTwoFactorMethods: + """The five management endpoints: path, verb, and body.""" + + @patch("colony_sdk.client.urlopen") + def test_get_2fa_status(self, mock_urlopen: MagicMock) -> None: + mock_urlopen.return_value = _mock_response({"enabled": True, "recovery_codes_remaining": 8}) + result = _authed_client().get_2fa_status() + + req = _last_request(mock_urlopen) + assert req.get_method() == "GET" + assert req.full_url == f"{BASE}/auth/2fa/status" + assert result == {"enabled": True, "recovery_codes_remaining": 8} + + @patch("colony_sdk.client.urlopen") + def test_enroll_2fa(self, mock_urlopen: MagicMock) -> None: + mock_urlopen.return_value = _mock_response( + {"secret": "S" * 32, "otpauth_uri": "otpauth://totp/x", "ticket": "t.sig"} + ) + result = _authed_client().enroll_2fa() + + req = _last_request(mock_urlopen) + assert req.get_method() == "POST" + assert req.full_url == f"{BASE}/auth/2fa/enroll" + assert result["otpauth_uri"].startswith("otpauth://") + + @patch("colony_sdk.client.urlopen") + def test_confirm_2fa(self, mock_urlopen: MagicMock) -> None: + mock_urlopen.return_value = _mock_response( + {"enabled": True, "recovery_codes": ["a", "b"], "recovery_codes_remaining": 2} + ) + result = _authed_client().confirm_2fa("SECRET", "ticket.sig", "123456") + + req = _last_request(mock_urlopen) + assert req.get_method() == "POST" + assert req.full_url == f"{BASE}/auth/2fa/confirm" + assert _last_body(mock_urlopen) == { + "secret": "SECRET", + "ticket": "ticket.sig", + "code": "123456", + } + assert result["recovery_codes"] == ["a", "b"] + + @patch("colony_sdk.client.urlopen") + def test_disable_2fa(self, mock_urlopen: MagicMock) -> None: + mock_urlopen.return_value = _mock_response({"enabled": False, "recovery_codes_remaining": 0}) + result = _authed_client().disable_2fa("123456") + + req = _last_request(mock_urlopen) + assert req.get_method() == "POST" + assert req.full_url == f"{BASE}/auth/2fa/disable" + assert _last_body(mock_urlopen) == {"code": "123456"} + assert result["enabled"] is False + + @patch("colony_sdk.client.urlopen") + def test_regenerate_recovery_codes(self, mock_urlopen: MagicMock) -> None: + mock_urlopen.return_value = _mock_response({"recovery_codes": ["x"], "recovery_codes_remaining": 1}) + result = _authed_client().regenerate_recovery_codes("123456") + + req = _last_request(mock_urlopen) + assert req.get_method() == "POST" + assert req.full_url == f"{BASE}/auth/2fa/recovery-codes/regenerate" + assert _last_body(mock_urlopen) == {"code": "123456"} + assert result["recovery_codes"] == ["x"] + + +class TestParity: + """Sync, async and the testing mock must expose the same surface.""" + + METHODS = ( + "get_2fa_status", + "enroll_2fa", + "confirm_2fa", + "disable_2fa", + "regenerate_recovery_codes", + ) + + def test_all_three_surfaces_have_every_method(self) -> None: + from colony_sdk.testing import MockColonyClient + + for name in self.METHODS: + assert hasattr(ColonyClient, name), f"sync missing {name}" + assert hasattr(AsyncColonyClient, name), f"async missing {name}" + assert hasattr(MockColonyClient, name), f"mock missing {name}" + + def test_mock_records_calls(self) -> None: + from colony_sdk.testing import MockColonyClient + + mock = MockColonyClient(responses={"get_2fa_status": {"enabled": True, "recovery_codes_remaining": 8}}) + assert mock.get_2fa_status()["enabled"] is True + mock.confirm_2fa("s", "t", "123456") + mock.disable_2fa("123456") + + recorded = [name for name, _ in mock.calls] + assert recorded == ["get_2fa_status", "confirm_2fa", "disable_2fa"] + assert mock.calls[1][1] == {"secret": "s", "ticket": "t", "code": "123456"}