diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fa242e..e5b91fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,16 @@ # Changelog -## Unreleased +## 0.3.0 (unreleased) + +- SDK: OAuth login. `Discolike(auth=...)` / `AsyncDiscolike(auth=...)` accept an `ApiKeyCredential` or `OAuthCredential` (both exported from `discolike`); `api_key=`, `DISCOLIKE_API_KEY`, and the config file keep working unchanged, and `auth=` wins over all of them. OAuth credentials send `Authorization: Bearer`, refresh proactively within 60s of expiry and once more after a 401, and write rotated refresh tokens back to the config file when they were loaded from it (an injected `auth=` is never persisted). A refresh that fails raises `AuthenticationError("OAuth session expired; run `discolike auth login`")`. Config file gains the shape `{"auth_method": "oauth", "oauth": {...}}` next to the existing `api_key` shape. +- SDK: `http_client=` now has its `.auth` set by the SDK (the `X-discolike-key` header moved from a static default header into an `httpx2.Auth`); an `auth` already set on a user-supplied client is replaced. +- CLI: `discolike auth login` now logs in through the browser by default (PKCE authorization-code flow against the platform's OAuth server, loopback redirect on `127.0.0.1`). `--no-browser` prints the URL only, `--port` pins the loopback port for SSH forwarding, and `--method api_key` (or `--api-key KEY`) keeps the API-key flow, including the prompt. Login-flow failures (timeout, denied consent, state mismatch) exit 1 with `{"error": "LoginError", ...}` on stderr. +- CLI: `discolike auth login` remembers the OAuth client it registered (`oauth_client` in the config file) and reuses it on the next login for the same server, so the browser consent screen is only asked once per machine. `auth logout` drops the credential but keeps the registration (a public PKCE client, not a secret), so the next login skips consent too. If PropelAuth no longer recognises the stored client (`invalid_client` / `unauthorized_client`), login registers a fresh one and retries once. +- CLI: `discolike auth status` adds `method` (`api_key` / `oauth`); for OAuth it reports `expires_at` and `expired` instead of a masked key. - SDK: `JobStatus` gains `estimated_cost` and `cost_metadata` for DiscoGen-family jobs (`discogen`, `validate_icp`, contacts generate). `cost_metadata` has one entry per `provider/model` and a `search_provider` entry with `queries_executed` / `queries_succeeded` / `est_cost_usd` when a BYOS search provider ran. `search_calls` on the model entries counts only the model's built-in search tool and is `0` on every BYOS run; read `search_provider.queries_executed` to confirm web search happened. -## 0.3.0 (2026-08-27) +### Request models - SDK (breaking): every request-taking method now takes a single request model instead of keyword arguments, and validates it locally before any HTTP call — a bad enum value, an out-of-range number, or a missing required field raises `pydantic.ValidationError` instead of a server 422. Models live in `discolike.requests` and are generated from the platform OpenAPI spec (`scripts/gen_requests.py`); query-param routes use `Params` (`MatchCompanyParams`, `ContactsSearchParams`, `DiscoverParams`, `CountParams`, `AppendParams`, `SegmentParams`, ...) and JSON-body routes use the platform's own names (`FindEmailRequest`, `ContactFilters`, `DiscoGenProcessRequest`, `UpdateQueryRequest`, ...). Path params and file uploads stay keyword arguments next to the model. Unknown fields pass through to the wire, so the SDK never blocks a platform field it does not know about yet. diff --git a/packages/discolike-cli/README.md b/packages/discolike-cli/README.md index 2f4f463..bab4b42 100644 --- a/packages/discolike-cli/README.md +++ b/packages/discolike-cli/README.md @@ -22,7 +22,7 @@ Requires Python 3.10+. Installing this package gives you the `discolike` command discolike auth login ``` -Prompts for an API key (or pass `--api-key`) and verifies it against your account. Create a key at [app.discolike.com/account/management/keys](https://app.discolike.com/account/management/keys). You can also set `DISCOLIKE_API_KEY` in the environment instead. +Opens your browser to log in (add `--no-browser` to print the URL instead, `--port` to pin the loopback port when forwarding over SSH). To use an API key instead, pass `--api-key KEY` or `--method api_key` to be prompted; create a key at [app.discolike.com/account/management/keys](https://app.discolike.com/account/management/keys). You can also set `DISCOLIKE_API_KEY` in the environment instead. ## Quickstart diff --git a/packages/discolike-cli/src/discolike_cli/_loopback.py b/packages/discolike-cli/src/discolike_cli/_loopback.py new file mode 100644 index 0000000..dab0ecc --- /dev/null +++ b/packages/discolike-cli/src/discolike_cli/_loopback.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import threading +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler +from http.server import HTTPServer +from urllib.parse import parse_qs +from urllib.parse import urlparse + +LOOPBACK_HOST = "127.0.0.1" +CALLBACK_PATH = "/callback" +CALLBACK_HTML = "DiscoLike

Login complete. You can close this window.

" + + +class _LoopbackServer(HTTPServer): + def __init__(self, *, host: str, port: int) -> None: + super().__init__((host, port), _CallbackHandler) + self.query: dict[str, str] = {} + self.received = threading.Event() + + +class _CallbackHandler(BaseHTTPRequestHandler): + server: _LoopbackServer + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path != CALLBACK_PATH: + self.send_error(HTTPStatus.NOT_FOUND) + return + self.server.query = {key: values[0] for key, values in parse_qs(parsed.query).items()} + body = CALLBACK_HTML.encode() + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + self.server.received.set() + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 -- BaseHTTPRequestHandler signature + _ = (format, args) + + +class CallbackServer: + """Loopback redirect target for the authorization-code flow; serves on a daemon thread.""" + + def __init__(self, *, port: int) -> None: + self._server = _LoopbackServer(host=LOOPBACK_HOST, port=port) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + + @property + def redirect_uri(self) -> str: + return f"http://{LOOPBACK_HOST}:{self._server.server_port}{CALLBACK_PATH}" + + def wait(self, *, timeout: float) -> dict[str, str] | None: + return self._server.query if self._server.received.wait(timeout) else None + + def __enter__(self) -> CallbackServer: + self._thread.start() + return self + + def __exit__(self, *exc_info: object) -> None: + self._server.shutdown() + self._server.server_close() diff --git a/packages/discolike-cli/src/discolike_cli/auth.py b/packages/discolike-cli/src/discolike_cli/auth.py index e1a0f6a..597385d 100644 --- a/packages/discolike-cli/src/discolike_cli/auth.py +++ b/packages/discolike-cli/src/discolike_cli/auth.py @@ -1,81 +1,285 @@ from __future__ import annotations import json +import secrets import sys +import webbrowser +from datetime import datetime +from datetime import timezone from typing import Any +from typing import NoReturn +from urllib.parse import urlparse +import httpx2 import typer +from discolike._config import AUTH_METHOD_API_KEY +from discolike._config import AUTH_METHOD_OAUTH +from discolike._config import DEFAULT_BASE_URL from discolike._config import KEYS_URL -from discolike._config import delete_config -from discolike._config import load_config -from discolike._config import resolve_api_key +from discolike._config import NO_CREDENTIAL_MESSAGE +from discolike._config import delete_credential +from discolike._config import delete_oauth_client +from discolike._config import load_credential +from discolike._config import load_oauth_client from discolike._config import save_config +from discolike._config import save_credential +from discolike._config import save_oauth_client +from discolike._credentials import OAuthClientRegistration +from discolike._credentials import OAuthCredential +from discolike._exceptions import AuthenticationError +from discolike._oauth import AuthServerMetadata +from discolike._oauth import OAuthError +from discolike._oauth import build_authorization_url +from discolike._oauth import discover +from discolike._oauth import exchange_code +from discolike._oauth import pkce_pair +from discolike._oauth import register_client +from discolike_cli._loopback import CallbackServer from discolike_cli._output import emit from discolike_cli._output import handle_errors -app = typer.Typer(help="Manage API credentials: log in, check key status, log out.") +app = typer.Typer(help="Manage credentials: log in (browser or API key), check status, log out.") MASKED_VISIBLE_CHARS = 4 +LOGIN_TIMEOUT_SECONDS = 180.0 +OAUTH_HTTP_TIMEOUT_SECONDS = 30.0 +STATE_BYTES = 16 +RANDOM_PORT = 0 SOURCE_OPTION = "option" SOURCE_ENV = "env" SOURCE_CONFIG = "config" +LOGIN_METHODS = (AUTH_METHOD_OAUTH, AUTH_METHOD_API_KEY) +DEAD_CLIENT_ERRORS = frozenset({"invalid_client", "unauthorized_client"}) + + +class _DeadClientError(Exception): + """The authorization server no longer recognises the registered client_id.""" + def _mask(key: str) -> str: return "…" + key[-MASKED_VISIBLE_CHARS:] +def _iso(epoch_seconds: float) -> str: + return datetime.fromtimestamp(epoch_seconds, tz=timezone.utc).isoformat() + + def _global_key_source(ctx: typer.Context) -> str: # typer vendors click without re-exporting ParameterSource, so match on the enum member name. source = ctx.find_root().get_parameter_source("api_key") return SOURCE_ENV if source is not None and source.name == "ENVIRONMENT" else SOURCE_OPTION -def _verify(ctx: typer.Context, *, api_key: str) -> None: +def _verify(ctx: typer.Context, **kwargs: Any) -> None: # noqa: ANN401 -- forwarded verbatim to build_client from discolike_cli.main import build_client - kwargs: dict[str, Any] = {"api_key": api_key} base_url = ctx.obj.get("base_url") if base_url is not None: kwargs["base_url"] = base_url build_client(**kwargs).account.usage() +def _abort_login(message: str) -> NoReturn: + print(json.dumps({"error": "LoginError", "message": message}), file=sys.stderr) + raise typer.Exit(code=1) + + +def _registered_port(registration: OAuthClientRegistration) -> int: + return int(urlparse(registration.redirect_uri).port or RANDOM_PORT) + + +def _reusable_registration(*, issuer: str, port: int) -> OAuthClientRegistration | None: + stored = load_oauth_client() + if stored is None or stored.issuer != issuer: + return None + if port != RANDOM_PORT and _registered_port(stored) != port: + return None + return stored + + +def _bind_stored_port(registration: OAuthClientRegistration) -> CallbackServer | None: + try: + return CallbackServer(port=_registered_port(registration)) + except OSError: + return None + + +def _register_or_reuse( + metadata: AuthServerMetadata, *, port: int, http: httpx2.Client +) -> tuple[CallbackServer, OAuthClientRegistration, bool]: + # PropelAuth matches the redirect URI literally (port included) and remembers consent per client_id, + # so a stored registration is only worth reusing when its exact port can be bound again. + stored = _reusable_registration(issuer=metadata.issuer, port=port) + if stored is not None: + server = _bind_stored_port(stored) + if server is not None: + return server, stored, True + server = CallbackServer(port=port) + client_id = register_client(metadata, redirect_uris=[server.redirect_uri], client=http) + registration = OAuthClientRegistration( + client_id=client_id, redirect_uri=server.redirect_uri, issuer=metadata.issuer + ) + save_oauth_client(registration) + return server, registration, False + + +def _authorize( + metadata: AuthServerMetadata, + registration: OAuthClientRegistration, + server: CallbackServer, + *, + resource: str, + open_browser: bool, + http: httpx2.Client, +) -> OAuthCredential: + verifier, challenge = pkce_pair() + state = secrets.token_urlsafe(STATE_BYTES) + url = build_authorization_url( + metadata, + client_id=registration.client_id, + redirect_uri=registration.redirect_uri, + code_challenge=challenge, + state=state, + resource=resource, + ) + print(f"Open this URL in your browser to log in:\n{url}", file=sys.stderr) + if open_browser and not webbrowser.open(url): + print("Could not open a browser; open the URL above manually.", file=sys.stderr) + callback = server.wait(timeout=LOGIN_TIMEOUT_SECONDS) + if callback is None: + _abort_login(f"Timed out after {LOGIN_TIMEOUT_SECONDS:.0f}s waiting for the browser login") + if callback.get("state") != state: + _abort_login("Invalid OAuth callback (state mismatch)") + if "error" in callback: + message = f"Authorization failed: {callback.get('error_description') or callback['error']}" + if callback["error"] in DEAD_CLIENT_ERRORS: + raise _DeadClientError(message) + _abort_login(message) + if "code" not in callback: + _abort_login("Invalid OAuth callback (missing code)") + try: + return exchange_code( + metadata, + client_id=registration.client_id, + code=callback["code"], + code_verifier=verifier, + redirect_uri=registration.redirect_uri, + resource=resource, + client=http, + ) + except OAuthError as exc: + if exc.error in DEAD_CLIENT_ERRORS: + raise _DeadClientError(f"Token exchange failed: {exc}") from exc + raise + + +def _oauth_login(ctx: typer.Context, *, open_browser: bool, port: int) -> OAuthCredential: + base_url = str(ctx.obj.get("base_url") or DEFAULT_BASE_URL).rstrip("/") + with httpx2.Client(timeout=OAUTH_HTTP_TIMEOUT_SECONDS) as http: + metadata = discover(base_url, client=http) + server, registration, reused = _register_or_reuse(metadata, port=port, http=http) + try: + with server: + return _authorize( + metadata, registration, server, resource=base_url, open_browser=open_browser, http=http + ) + except _DeadClientError as exc: + if not reused: + _abort_login(str(exc)) + delete_oauth_client() + server, registration, _ = _register_or_reuse(metadata, port=port, http=http) + try: + with server: + return _authorize( + metadata, registration, server, resource=base_url, open_browser=open_browser, http=http + ) + except _DeadClientError as exc: + _abort_login(str(exc)) + + +def _api_key_login(ctx: typer.Context, *, api_key: str | None) -> None: + # An ambient DISCOLIKE_API_KEY must not silently become the saved key; only an explicit flag may. + passed_globally = ctx.obj.get("api_key") if _global_key_source(ctx) == SOURCE_OPTION else None + key = api_key or passed_globally or typer.prompt("API key", hide_input=True) + _verify(ctx, api_key=key) + save_config({"auth_method": AUTH_METHOD_API_KEY, "api_key": key}) + print(json.dumps({"logged_in": True, "source": AUTH_METHOD_API_KEY}), file=sys.stderr) + + @app.command() @handle_errors def login( ctx: typer.Context, - api_key: str | None = typer.Option(None, help=f"API key. Create one at {KEYS_URL}. Prompted for if omitted."), + api_key: str | None = typer.Option( + None, help=f"Log in with an API key instead of the browser. Create one at {KEYS_URL}." + ), + method: str = typer.Option( + AUTH_METHOD_OAUTH, + "--method", + help="oauth (browser login, default) or api_key (prompts for a key unless --api-key is given).", + ), + no_browser: bool = typer.Option(False, "--no-browser", help="Print the login URL instead of opening a browser."), + port: int = typer.Option( + RANDOM_PORT, "--port", help="Fixed loopback port for the browser redirect (default: random; use with SSH)." + ), ) -> None: - """Verify an API key and save it to the local config file.""" - # An ambient DISCOLIKE_API_KEY must not silently become the saved key; only an explicit flag may. - passed_globally = ctx.obj.get("api_key") if _global_key_source(ctx) == SOURCE_OPTION else None - key = api_key or passed_globally or typer.prompt("API key", hide_input=True) - _verify(ctx, api_key=key) - save_config({"auth_method": "api_key", "api_key": key}) - print(json.dumps({"logged_in": True, "source": "api_key"}), file=sys.stderr) + """Log in via the browser (OAuth) or with an API key, verify, and save to the local config file.""" + if method not in LOGIN_METHODS: + raise typer.BadParameter(f"must be one of {', '.join(LOGIN_METHODS)}", param_hint="--method") + global_key_passed = ctx.obj.get("api_key") is not None and _global_key_source(ctx) == SOURCE_OPTION + if api_key or global_key_passed or method == AUTH_METHOD_API_KEY: + _api_key_login(ctx, api_key=api_key) + return + credential = _oauth_login(ctx, open_browser=not no_browser, port=port) + _verify(ctx, auth=credential) + save_credential(credential) + print( + json.dumps({"logged_in": True, "method": AUTH_METHOD_OAUTH, "expires_at": _iso(credential.expires_at)}), + file=sys.stderr, + ) @app.command() @handle_errors def status(ctx: typer.Context) -> None: - """Show which API key is in use (option, env, or config) and verify it against the API.""" + """Show which credential is in use (option, env, or config) and verify it against the API.""" key = ctx.obj.get("api_key") - source = _global_key_source(ctx) - if not key: - key = load_config().get("api_key") - source = SOURCE_CONFIG - if not key: - resolve_api_key(None) - _verify(ctx, api_key=str(key)) - emit({"source": source, "api_key": _mask(str(key)), "valid": True}) + if key: + _verify(ctx, api_key=str(key)) + emit( + { + "source": _global_key_source(ctx), + "method": AUTH_METHOD_API_KEY, + "api_key": _mask(str(key)), + "valid": True, + } + ) + return + credential = load_credential() + if credential is None: + raise AuthenticationError(NO_CREDENTIAL_MESSAGE) + if isinstance(credential, OAuthCredential): + _verify(ctx) + emit( + { + "source": SOURCE_CONFIG, + "method": AUTH_METHOD_OAUTH, + "expires_at": _iso(credential.expires_at), + "expired": credential.expires_within(0), + "valid": True, + } + ) + return + _verify(ctx, api_key=credential.api_key) + emit({"source": SOURCE_CONFIG, "method": AUTH_METHOD_API_KEY, "api_key": _mask(credential.api_key), "valid": True}) @app.command() @handle_errors def logout() -> None: - """Delete saved credentials from the local config file.""" - delete_config() + """Delete saved credentials from the local config file (the registered OAuth client is kept).""" + delete_credential() emit({"logged_out": True}) diff --git a/packages/discolike-cli/tests/test_auth.py b/packages/discolike-cli/tests/test_auth.py index b2f5f72..f7b3de9 100644 --- a/packages/discolike-cli/tests/test_auth.py +++ b/packages/discolike-cli/tests/test_auth.py @@ -39,7 +39,7 @@ def test_login_with_api_key_option_verifies_and_saves(install_build_client: Call def test_login_prompts_for_key_when_not_given(install_build_client: Callable[[Handler], None]) -> None: install_build_client(_usage_ok) - result = runner.invoke(app, ["auth", "login"], input="dk-2\n") + result = runner.invoke(app, ["auth", "login", "--method", "api_key"], input="dk-2\n") assert result.exit_code == 0, result.output assert json.loads(config_path().read_text())["api_key"] == "dk-2" @@ -196,7 +196,7 @@ def test_login_ignores_ambient_env_key_and_still_prompts( ) -> None: install_build_client(_usage_ok) monkeypatch.setenv(ENV_API_KEY, "dk-from-env") - result = runner.invoke(app, ["auth", "login"], input="dk-typed\n") + result = runner.invoke(app, ["auth", "login", "--method", "api_key"], input="dk-typed\n") assert result.exit_code == 0, result.output assert json.loads(config_path().read_text())["api_key"] == "dk-typed" diff --git a/packages/discolike-cli/tests/test_auth_oauth.py b/packages/discolike-cli/tests/test_auth_oauth.py new file mode 100644 index 0000000..a241008 --- /dev/null +++ b/packages/discolike-cli/tests/test_auth_oauth.py @@ -0,0 +1,411 @@ +from __future__ import annotations + +import dataclasses +import json +import socket +import threading +import time +from collections.abc import Callable +from typing import Any +from urllib.parse import parse_qs +from urllib.parse import urlparse +from urllib.request import urlopen + +import httpx2 +import pytest +from typer.testing import CliRunner + +import discolike_cli.auth as auth_module +from discolike._config import DEFAULT_BASE_URL +from discolike._config import config_path +from discolike._config import load_credential +from discolike._config import load_oauth_client +from discolike._config import save_credential +from discolike._config import save_oauth_client +from discolike._credentials import ApiKeyCredential +from discolike._credentials import OAuthClientRegistration +from discolike._credentials import OAuthCredential +from discolike._oauth import AuthServerMetadata +from discolike._oauth import OAuthError +from discolike_cli.main import app +from discolike_testkit import Handler + +runner = CliRunner() + +METADATA = AuthServerMetadata( + authorization_endpoint="https://auth.test/oauth/2.1/authorize", + token_endpoint="https://auth.test/oauth/2.1/token", + registration_endpoint="https://auth.test/oauth/2.1/register", + issuer="https://auth.test/oauth/2.1", +) +CREDENTIAL = OAuthCredential( + access_token="at-1", + refresh_token="rt-1", + expires_at=1_800_000_000.0, + client_id="client-1", + token_endpoint=METADATA.token_endpoint, +) + + +def _usage_ok(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"requests_mtd": 1}) + + +def _usage_unauthorized(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(401, json={"detail": "Invalid API Key or Session"}) + + +class FakeProvider: + """Stands in for the authorization server and the user's browser.""" + + def __init__(self) -> None: + self.real_build_authorization_url = auth_module.build_authorization_url + self.discover_calls: list[str] = [] + self.register_calls: list[list[str]] = [] + self.exchange_calls: list[dict[str, Any]] = [] + self.opened_urls: list[str] = [] + self.exchange_failures: list[OAuthError] = [] + self.callback_query: Callable[[dict[str, str]], str] = lambda query: f"code=the-code&state={query['state']}" + + def discover(self, base_url: str, *, client: httpx2.Client) -> AuthServerMetadata: + self.discover_calls.append(base_url) + return METADATA + + def register_client(self, metadata: AuthServerMetadata, *, redirect_uris: list[str], client: httpx2.Client) -> str: + self.register_calls.append(redirect_uris) + return "client-1" + + def exchange_code(self, metadata: AuthServerMetadata, *, client: httpx2.Client, **kwargs: Any) -> OAuthCredential: + self.exchange_calls.append(kwargs) + if self.exchange_failures: + raise self.exchange_failures.pop(0) + return CREDENTIAL + + def build_authorization_url(self, metadata: AuthServerMetadata, **kwargs: Any) -> str: + url = self.real_build_authorization_url(metadata, **kwargs) + query = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()} + callback = f"{query['redirect_uri']}?{self.callback_query(query)}" + threading.Thread(target=lambda: urlopen(callback).read(), daemon=True).start() # noqa: S310 -- loopback test server + return url + + def open(self, url: str) -> bool: + self.opened_urls.append(url) + return True + + +@pytest.fixture +def provider(monkeypatch: pytest.MonkeyPatch) -> FakeProvider: + fake = FakeProvider() + monkeypatch.setattr(auth_module, "discover", fake.discover) + monkeypatch.setattr(auth_module, "register_client", fake.register_client) + monkeypatch.setattr(auth_module, "exchange_code", fake.exchange_code) + monkeypatch.setattr(auth_module, "build_authorization_url", fake.build_authorization_url) + monkeypatch.setattr(auth_module.webbrowser, "open", fake.open) + return fake + + +def test_login_default_runs_oauth_loopback_flow( + provider: FakeProvider, + install_build_client: Callable[[Handler], None], + build_client_calls: list[dict[str, Any]], +) -> None: + install_build_client(_usage_ok) + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + assert provider.discover_calls == [DEFAULT_BASE_URL] + redirect_uri = provider.register_calls[0][0] + assert redirect_uri.startswith("http://127.0.0.1:") + assert redirect_uri.endswith("/callback") + exchange = provider.exchange_calls[0] + assert exchange["code"] == "the-code" + assert exchange["client_id"] == "client-1" + assert exchange["redirect_uri"] == redirect_uri + assert exchange["resource"] == DEFAULT_BASE_URL + assert len(provider.opened_urls) == 1 + assert build_client_calls == [{"auth": CREDENTIAL}] + stored = json.loads(config_path().read_text()) + assert (stored["auth_method"], stored["oauth"]) == ("oauth", CREDENTIAL.to_config()) + assert stored["oauth_client"] == {"client_id": "client-1", "redirect_uri": redirect_uri, "issuer": METADATA.issuer} + payload = json.loads(result.stderr.splitlines()[-1]) + assert payload == {"logged_in": True, "method": "oauth", "expires_at": "2027-01-15T08:00:00+00:00"} + assert provider.opened_urls[0] in result.stderr + + +def test_login_no_browser_and_fixed_port( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + result = runner.invoke( + app, ["--base-url", "https://api.dev.test/v1/", "auth", "login", "--no-browser", "--port", "18484"] + ) + assert result.exit_code == 0, result.output + assert provider.opened_urls == [] + assert provider.register_calls == [["http://127.0.0.1:18484/callback"]] + assert provider.discover_calls == ["https://api.dev.test/v1"] + assert provider.exchange_calls[0]["resource"] == "https://api.dev.test/v1" + + +def test_login_state_mismatch_exits_1_and_saves_nothing( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + provider.callback_query = lambda query: "code=the-code&state=forged" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert load_credential() is None + assert provider.exchange_calls == [] + assert json.loads(result.stderr.splitlines()[-1])["error"] == "LoginError" + + +def test_login_user_denied_exits_1(provider: FakeProvider, install_build_client: Callable[[Handler], None]) -> None: + install_build_client(_usage_ok) + provider.callback_query = lambda query: f"error=access_denied&error_description=nope&state={query['state']}" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert "nope" in json.loads(result.stderr.splitlines()[-1])["message"] + + +def test_login_timeout_exits_1( + provider: FakeProvider, install_build_client: Callable[[Handler], None], monkeypatch: pytest.MonkeyPatch +) -> None: + install_build_client(_usage_ok) + monkeypatch.setattr(auth_module, "LOGIN_TIMEOUT_SECONDS", 0.2) + monkeypatch.setattr(auth_module, "build_authorization_url", lambda metadata, **kwargs: "https://auth.test/never") + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert "Timed out" in json.loads(result.stderr.splitlines()[-1])["message"] + assert load_credential() is None + + +def test_login_oauth_verify_failure_exits_3_and_saves_nothing( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_unauthorized) + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 3 + assert load_credential() is None + assert json.loads(result.stderr.splitlines()[-1])["error"] == "AuthenticationError" + + +def test_login_rejects_unknown_method(provider: FakeProvider) -> None: + result = runner.invoke(app, ["auth", "login", "--method", "magic"]) + assert result.exit_code == 2 + assert provider.discover_calls == [] + + +def test_status_reports_oauth_credential( + install_build_client: Callable[[Handler], None], build_client_calls: list[dict[str, Any]] +) -> None: + install_build_client(_usage_ok) + save_credential(CREDENTIAL) + result = runner.invoke(app, ["auth", "status"]) + assert result.exit_code == 0, result.output + assert build_client_calls == [{}] + assert json.loads(result.stdout) == { + "source": "config", + "method": "oauth", + "expires_at": "2027-01-15T08:00:00+00:00", + "expired": False, + "valid": True, + } + + +def test_status_flags_expired_oauth_credential(install_build_client: Callable[[Handler], None]) -> None: + install_build_client(_usage_ok) + save_credential(dataclasses.replace(CREDENTIAL, expires_at=time.time() - 1)) + result = runner.invoke(app, ["auth", "status"]) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["expired"] is True + + +def test_status_reports_api_key_method(install_build_client: Callable[[Handler], None]) -> None: + install_build_client(_usage_ok) + result = runner.invoke(app, ["--api-key", "dk-abcdefgh1234", "auth", "status"]) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["method"] == "api_key" + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _registration(port: int, *, issuer: str = METADATA.issuer) -> OAuthClientRegistration: + return OAuthClientRegistration( + client_id="stored-client", redirect_uri=f"http://127.0.0.1:{port}/callback", issuer=issuer + ) + + +def test_login_reuses_stored_client_when_its_port_is_free( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + port = _free_port() + save_oauth_client(_registration(port)) + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + assert provider.register_calls == [] + assert provider.exchange_calls[0]["client_id"] == "stored-client" + assert provider.exchange_calls[0]["redirect_uri"] == f"http://127.0.0.1:{port}/callback" + assert load_oauth_client() == _registration(port) + assert json.loads(config_path().read_text())["auth_method"] == "oauth" + + +def test_login_registers_anew_when_stored_port_is_busy( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + with socket.socket() as blocker: + blocker.bind(("127.0.0.1", 0)) + blocker.listen() + busy_port = blocker.getsockname()[1] + save_oauth_client(_registration(busy_port)) + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + new_uri = provider.register_calls[0][0] + assert new_uri != f"http://127.0.0.1:{busy_port}/callback" + stored = load_oauth_client() + assert stored is not None + assert (stored.client_id, stored.redirect_uri) == ("client-1", new_uri) + + +def test_login_registers_anew_for_a_different_issuer( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(_free_port(), issuer="https://auth.other/oauth/2.1")) + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + assert len(provider.register_calls) == 1 + stored = load_oauth_client() + assert stored is not None + assert stored.issuer == METADATA.issuer + + +def test_login_explicit_port_differing_from_stored_registers_anew( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(_free_port())) + wanted = _free_port() + result = runner.invoke(app, ["auth", "login", "--port", str(wanted)]) + assert result.exit_code == 0, result.output + assert provider.register_calls == [[f"http://127.0.0.1:{wanted}/callback"]] + + +def test_logout_keeps_stored_client_and_drops_credential(install_build_client: Callable[[Handler], None]) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(18484)) + save_credential(CREDENTIAL) + result = runner.invoke(app, ["auth", "logout"]) + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == {"logged_out": True} + assert load_credential() is None + assert load_oauth_client() == _registration(18484) + status = runner.invoke(app, ["auth", "status"]) + assert status.exit_code == 3 + assert "discolike auth login" in json.loads(status.stderr)["message"] + + +def test_logout_with_api_key_config_removes_the_key_but_keeps_stored_client() -> None: + save_oauth_client(_registration(18484)) + save_credential(ApiKeyCredential(api_key="dk-1")) + result = runner.invoke(app, ["auth", "logout"]) + assert result.exit_code == 0, result.output + stored = json.loads(config_path().read_text()) + assert "api_key" not in stored + assert "auth_method" not in stored + assert load_oauth_client() == _registration(18484) + + +def _dead_then_ok(query: dict[str, str]) -> str: + if query["client_id"] == "stored-client": + return f"error=invalid_client&error_description=unknown+client&state={query['state']}" + return f"code=the-code&state={query['state']}" + + +def test_login_reregisters_when_authorize_rejects_the_stored_client( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(_free_port())) + provider.callback_query = _dead_then_ok + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + assert len(provider.register_calls) == 1 + assert [call["client_id"] for call in provider.exchange_calls] == ["client-1"] + stored = load_oauth_client() + assert stored is not None + assert (stored.client_id, stored.redirect_uri) == ("client-1", provider.register_calls[0][0]) + assert load_credential() == CREDENTIAL + + +def test_login_reregisters_when_exchange_rejects_the_stored_client( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(_free_port())) + provider.exchange_failures = [OAuthError("invalid_client: gone", error="invalid_client", status_code=400)] + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 0, result.output + assert [call["client_id"] for call in provider.exchange_calls] == ["stored-client", "client-1"] + assert len(provider.register_calls) == 1 + stored = load_oauth_client() + assert stored is not None + assert stored.client_id == "client-1" + assert load_credential() == CREDENTIAL + + +def test_login_fresh_client_rejected_is_a_login_error_without_retry( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + provider.callback_query = lambda query: f"error=unauthorized_client&state={query['state']}" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert len(provider.register_calls) == 1 + assert provider.exchange_calls == [] + assert json.loads(result.stderr.splitlines()[-1]) == { + "error": "LoginError", + "message": "Authorization failed: unauthorized_client", + } + + +def test_login_reused_client_rejected_twice_is_a_login_error( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + save_oauth_client(_registration(_free_port())) + provider.callback_query = lambda query: f"error=invalid_client&state={query['state']}" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert len(provider.register_calls) == 1 + assert json.loads(result.stderr.splitlines()[-1])["error"] == "LoginError" + + +def test_login_access_denied_keeps_stored_client( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + registration = _registration(_free_port()) + save_oauth_client(registration) + provider.callback_query = lambda query: f"error=access_denied&state={query['state']}" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert provider.register_calls == [] + assert load_oauth_client() == registration + + +def test_login_forged_error_callback_cannot_evict_stored_client( + provider: FakeProvider, install_build_client: Callable[[Handler], None] +) -> None: + install_build_client(_usage_ok) + registration = _registration(_free_port()) + save_oauth_client(registration) + provider.callback_query = lambda query: "error=invalid_client&state=forged" + result = runner.invoke(app, ["auth", "login"]) + assert result.exit_code == 1 + assert "state mismatch" in json.loads(result.stderr.splitlines()[-1])["message"] + assert provider.register_calls == [] + assert load_oauth_client() == registration diff --git a/packages/discolike-testkit/src/discolike_testkit/__init__.py b/packages/discolike-testkit/src/discolike_testkit/__init__.py index c6a084d..8afa009 100644 --- a/packages/discolike-testkit/src/discolike_testkit/__init__.py +++ b/packages/discolike-testkit/src/discolike_testkit/__init__.py @@ -13,9 +13,16 @@ from discolike import AsyncDiscolike from discolike import Discolike +from discolike._auth import DiscolikeAuth +from discolike._credentials import ApiKeyCredential -__all__ = ["AsyncClientFactory", "ClientFactory", "Handler"] +__all__ = ["AsyncClientFactory", "ClientFactory", "Handler", "api_key_auth"] Handler = Callable[[httpx2.Request], httpx2.Response] ClientFactory = Callable[[Handler], Discolike] AsyncClientFactory = Callable[[Handler], AsyncDiscolike] + + +def api_key_auth(api_key: str) -> DiscolikeAuth: + """Auth for tests that build a ``Transport`` directly instead of going through ``Discolike``.""" + return DiscolikeAuth(ApiKeyCredential(api_key=api_key)) diff --git a/packages/discolike/README.md b/packages/discolike/README.md index e7d6aff..db8ef62 100644 --- a/packages/discolike/README.md +++ b/packages/discolike/README.md @@ -18,7 +18,7 @@ Requires Python 3.10+. export DISCOLIKE_API_KEY="dl_..." ``` -Create a key at [app.discolike.com/account/management/keys](https://app.discolike.com/account/management/keys). You can also pass `api_key=...` explicitly to `Discolike()`. +Create a key at [app.discolike.com/account/management/keys](https://app.discolike.com/account/management/keys). You can also pass `api_key=...` explicitly to `Discolike()`, or run `discolike auth login` from the CLI to log in through the browser — the SDK then picks up the saved OAuth session and refreshes it automatically. ## Quickstart diff --git a/packages/discolike/src/discolike/__init__.py b/packages/discolike/src/discolike/__init__.py index 7f4b96b..48419cd 100644 --- a/packages/discolike/src/discolike/__init__.py +++ b/packages/discolike/src/discolike/__init__.py @@ -1,5 +1,7 @@ from discolike._client import AsyncDiscolike from discolike._client import Discolike +from discolike._credentials import ApiKeyCredential +from discolike._credentials import OAuthCredential from discolike._exceptions import APIConnectionError from discolike._exceptions import AuthenticationError from discolike._exceptions import DiscolikeError @@ -26,6 +28,7 @@ __all__ = [ "APIConnectionError", + "ApiKeyCredential", "AsyncDiscolike", "AsyncJob", "AuthenticationError", @@ -44,6 +47,7 @@ "JobStatus", "JobTimeoutError", "NotFoundError", + "OAuthCredential", "PlanAccessError", "RateLimitError", "ServerError", diff --git a/packages/discolike/src/discolike/_auth.py b/packages/discolike/src/discolike/_auth.py new file mode 100644 index 0000000..e77c2ae --- /dev/null +++ b/packages/discolike/src/discolike/_auth.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import asyncio +import threading +from collections.abc import AsyncGenerator +from collections.abc import Callable +from collections.abc import Generator +from typing import cast + +import httpx2 + +from discolike._credentials import ApiKeyCredential +from discolike._credentials import Credential +from discolike._credentials import OAuthCredential +from discolike._exceptions import AuthenticationError +from discolike._oauth import REFRESH_LEEWAY_SECONDS +from discolike._oauth import SESSION_EXPIRED_MESSAGE +from discolike._oauth import parse_refresh_response +from discolike._oauth import refresh_request + +# TODO: replace this module and _oauth.py with authlib's httpx2 OAuth2Client once a release +# includes authlib/authlib@e4fb941 (httpx2 support merged 2026-08-27; 1.7.2 predates it). + +API_KEY_HEADER = "X-discolike-key" +UNAUTHORIZED = 401 + + +def _set_bearer(request: httpx2.Request, credential: OAuthCredential) -> None: + request.headers["Authorization"] = f"Bearer {credential.access_token}" + + +class DiscolikeAuth(httpx2.Auth): + """Sends the API key header, or a bearer token that is refreshed before expiry and once after a 401. + + Refreshes go through the same client as the request they precede, so tests drive them via ``MockTransport``. + """ + + requires_response_body = False + + def __init__( + self, + credential: Credential, + *, + on_update: Callable[[OAuthCredential], None] | None = None, + reload: Callable[[], Credential | None] | None = None, + ) -> None: + self._credential = credential + self.on_update = on_update + self.reload = reload + self._lock = threading.Lock() + self._async_lock = asyncio.Lock() + + @property + def credential(self) -> Credential: + return self._credential + + def _latest(self, credential: OAuthCredential) -> OAuthCredential: + return cast(OAuthCredential, self._credential) if self._credential is not credential else credential + + def _adopt_stored(self, credential: OAuthCredential) -> OAuthCredential | None: + """Another process may have rotated the tokens already; a refresh with our old refresh token would fail.""" + if self.reload is None: + return None + stored = self.reload() + if ( + not isinstance(stored, OAuthCredential) + or stored.access_token == credential.access_token + or stored.expires_within(REFRESH_LEEWAY_SECONDS) + ): + return None + self._credential = stored + return stored + + def _store(self, response: httpx2.Response, *, credential: OAuthCredential) -> OAuthCredential: + try: + rotated = parse_refresh_response(response, credential=credential) + except AuthenticationError as exc: + raise AuthenticationError( + SESSION_EXPIRED_MESSAGE, status_code=exc.status_code, payload=exc.payload + ) from exc + self._credential = rotated + if self.on_update is not None: + self.on_update(rotated) + return rotated + + def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + credential = self._credential + if isinstance(credential, ApiKeyCredential): + request.headers[API_KEY_HEADER] = credential.api_key + yield request + return + with self._lock: + credential = self._latest(credential) + if credential.expires_within(REFRESH_LEEWAY_SECONDS): + adopted = self._adopt_stored(credential) + if adopted is not None: + credential = adopted + else: + credential = yield from self._sync_refresh(credential) + _set_bearer(request, credential) + response = yield request + if response.status_code != UNAUTHORIZED: + return + with self._lock: + latest = self._latest(credential) + if latest is credential: + adopted = self._adopt_stored(credential) + if adopted is not None: + latest = adopted + else: + latest = yield from self._sync_refresh(credential) + _set_bearer(request, latest) + yield request + + def _sync_refresh(self, credential: OAuthCredential) -> Generator[httpx2.Request, httpx2.Response, OAuthCredential]: + response = yield refresh_request(credential) + response.read() + return self._store(response, credential=credential) + + async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + credential = self._credential + if isinstance(credential, ApiKeyCredential): + request.headers[API_KEY_HEADER] = credential.api_key + yield request + return + async with self._async_lock: + credential = self._latest(credential) + if credential.expires_within(REFRESH_LEEWAY_SECONDS): + adopted = self._adopt_stored(credential) + if adopted is not None: + credential = adopted + else: + response = yield refresh_request(credential) + await response.aread() + credential = self._store(response, credential=credential) + _set_bearer(request, credential) + response = yield request + if response.status_code != UNAUTHORIZED: + return + async with self._async_lock: + latest = self._latest(credential) + if latest is credential: + adopted = self._adopt_stored(credential) + if adopted is not None: + latest = adopted + else: + response = yield refresh_request(credential) + await response.aread() + latest = self._store(response, credential=credential) + _set_bearer(request, latest) + yield request diff --git a/packages/discolike/src/discolike/_client.py b/packages/discolike/src/discolike/_client.py index 623455e..fc4998d 100644 --- a/packages/discolike/src/discolike/_client.py +++ b/packages/discolike/src/discolike/_client.py @@ -2,8 +2,12 @@ import httpx2 +from discolike._auth import DiscolikeAuth from discolike._config import DEFAULT_BASE_URL -from discolike._config import resolve_api_key +from discolike._config import load_credential +from discolike._config import resolve_credential +from discolike._config import save_credential +from discolike._credentials import Credential from discolike._jobs import AsyncJob from discolike._jobs import Job from discolike._transport import AsyncTransport @@ -47,11 +51,20 @@ DEFAULT_MAX_RETRIES = 3 +def _build_auth(*, api_key: str | None, auth: Credential | None) -> DiscolikeAuth: + # The config file is read back and written only when the credential came from it. + credential = resolve_credential(api_key=api_key, auth=auth) + if auth is not None: + return DiscolikeAuth(credential) + return DiscolikeAuth(credential, on_update=save_credential, reload=load_credential) + + class Discolike: def __init__( self, *, api_key: str | None = None, + auth: Credential | None = None, base_url: str = DEFAULT_BASE_URL, timeout: float = DEFAULT_TIMEOUT_SECONDS, max_retries: int = DEFAULT_MAX_RETRIES, @@ -59,7 +72,7 @@ def __init__( ) -> None: self._attach( Transport( - resolve_api_key(api_key), + _build_auth(api_key=api_key, auth=auth), base_url=base_url, timeout=timeout, max_retries=max_retries, @@ -121,6 +134,7 @@ def __init__( self, *, api_key: str | None = None, + auth: Credential | None = None, base_url: str = DEFAULT_BASE_URL, timeout: float = DEFAULT_TIMEOUT_SECONDS, max_retries: int = DEFAULT_MAX_RETRIES, @@ -128,7 +142,7 @@ def __init__( ) -> None: self._attach( AsyncTransport( - resolve_api_key(api_key), + _build_auth(api_key=api_key, auth=auth), base_url=base_url, timeout=timeout, max_retries=max_retries, diff --git a/packages/discolike/src/discolike/_config.py b/packages/discolike/src/discolike/_config.py index 65501d4..12924fe 100644 --- a/packages/discolike/src/discolike/_config.py +++ b/packages/discolike/src/discolike/_config.py @@ -2,16 +2,24 @@ import json import os +import tempfile from pathlib import Path from typing import Any +from discolike._credentials import ApiKeyCredential +from discolike._credentials import Credential +from discolike._credentials import OAuthClientRegistration +from discolike._credentials import OAuthCredential from discolike._exceptions import AuthenticationError DEFAULT_BASE_URL = "https://api.discolike.com/v1" ENV_API_KEY = "DISCOLIKE_API_KEY" # foxguard: ignore[py/no-hardcoded-secret] KEYS_URL = "https://app.discolike.com/account/management/keys" +AUTH_METHOD_API_KEY = "api_key" +AUTH_METHOD_OAUTH = "oauth" +OAUTH_CLIENT_KEY = "oauth_client" -_NO_KEY_MESSAGE = ( +NO_CREDENTIAL_MESSAGE = ( "No API key found. Set the DISCOLIKE_API_KEY environment variable, pass api_key=..., " f"or run `discolike auth login`. Create a key at {KEYS_URL}" ) @@ -36,23 +44,74 @@ def load_config() -> dict[str, Any]: def save_config(config: dict[str, Any]) -> None: path = config_path() path.parent.mkdir(parents=True, exist_ok=True) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + fd, temp_path = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") with os.fdopen(fd, "w") as handle: handle.write(json.dumps(config, indent=2) + "\n") - path.chmod(0o600) + os.chmod(temp_path, 0o600) + os.replace(temp_path, path) def delete_config() -> None: config_path().unlink(missing_ok=True) -def resolve_api_key(explicit: str | None = None) -> str: - if explicit: - return explicit +def delete_credential() -> None: + """Forget the credential but keep the OAuth client registration; it is a public PKCE client, not a secret.""" + stored_client = load_config().get(OAUTH_CLIENT_KEY) + if stored_client is None: + delete_config() + return + save_config({OAUTH_CLIENT_KEY: stored_client}) + + +def load_credential() -> Credential | None: + config = load_config() + if config.get("auth_method") == AUTH_METHOD_OAUTH: + try: + return OAuthCredential.from_config(config["oauth"]) + except (KeyError, TypeError, ValueError): + return None + api_key = config.get("api_key") + return ApiKeyCredential(api_key=str(api_key)) if api_key else None + + +def save_credential(credential: Credential) -> None: + if isinstance(credential, OAuthCredential): + config: dict[str, Any] = {"auth_method": AUTH_METHOD_OAUTH, "oauth": credential.to_config()} + else: + config = {"auth_method": AUTH_METHOD_API_KEY, "api_key": credential.api_key} + stored_client = load_config().get(OAUTH_CLIENT_KEY) + if stored_client is not None: + config[OAUTH_CLIENT_KEY] = stored_client + save_config(config) + + +def load_oauth_client() -> OAuthClientRegistration | None: + try: + return OAuthClientRegistration.from_config(load_config()[OAUTH_CLIENT_KEY]) + except (KeyError, TypeError, ValueError): + return None + + +def save_oauth_client(registration: OAuthClientRegistration) -> None: + save_config({**load_config(), OAUTH_CLIENT_KEY: registration.to_config()}) + + +def delete_oauth_client() -> None: + config = load_config() + config.pop(OAUTH_CLIENT_KEY, None) + save_config(config) + + +def resolve_credential(*, api_key: str | None = None, auth: Credential | None = None) -> Credential: + if auth is not None: + return auth + if api_key: + return ApiKeyCredential(api_key=api_key) from_env = os.environ.get(ENV_API_KEY) if from_env: - return from_env - from_file = load_config().get("api_key") - if from_file: - return str(from_file) - raise AuthenticationError(_NO_KEY_MESSAGE) + return ApiKeyCredential(api_key=from_env) + credential = load_credential() + if credential is not None: + return credential + raise AuthenticationError(NO_CREDENTIAL_MESSAGE) diff --git a/packages/discolike/src/discolike/_credentials.py b/packages/discolike/src/discolike/_credentials.py new file mode 100644 index 0000000..c08644a --- /dev/null +++ b/packages/discolike/src/discolike/_credentials.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import time +from dataclasses import asdict +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class ApiKeyCredential: + api_key: str + + +@dataclass(frozen=True) +class OAuthCredential: + access_token: str + refresh_token: str + expires_at: float + client_id: str + token_endpoint: str + + def expires_within(self, seconds: float, *, now: float | None = None) -> bool: + current = time.time() if now is None else now + return self.expires_at - current <= seconds + + @classmethod + def from_config(cls, data: dict[str, Any]) -> OAuthCredential: + return cls( + access_token=str(data["access_token"]), + refresh_token=str(data["refresh_token"]), + expires_at=float(data["expires_at"]), + client_id=str(data["client_id"]), + token_endpoint=str(data["token_endpoint"]), + ) + + def to_config(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class OAuthClientRegistration: + client_id: str + redirect_uri: str + issuer: str + + @classmethod + def from_config(cls, data: dict[str, Any]) -> OAuthClientRegistration: + return cls(client_id=str(data["client_id"]), redirect_uri=str(data["redirect_uri"]), issuer=str(data["issuer"])) + + def to_config(self) -> dict[str, Any]: + return asdict(self) + + +Credential = ApiKeyCredential | OAuthCredential diff --git a/packages/discolike/src/discolike/_oauth.py b/packages/discolike/src/discolike/_oauth.py new file mode 100644 index 0000000..fb295b2 --- /dev/null +++ b/packages/discolike/src/discolike/_oauth.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import base64 +import hashlib +import secrets +import time +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlencode + +import httpx2 + +from discolike._credentials import OAuthCredential +from discolike._exceptions import AuthenticationError + +OAUTH_SCOPE = "offline_access" +REFRESH_LEEWAY_SECONDS = 60.0 +CLIENT_NAME = "discolike-cli" +METADATA_PATH = "/.well-known/oauth-authorization-server" +GRANT_TYPES = ["authorization_code", "refresh_token"] +RESPONSE_TYPES = ["code"] +PKCE_METHOD = "S256" +PKCE_VERIFIER_BYTES = 32 +TOKEN_HEADERS = {"Accept": "application/json"} +TOKEN_KEYS = frozenset({"access_token", "refresh_token", "id_token"}) +SESSION_EXPIRED_MESSAGE = "OAuth session expired; run `discolike auth login`" + + +class OAuthError(AuthenticationError): + """An RFC 6749 error body from the authorization server; `error` is its machine-readable code.""" + + def __init__( + self, + message: str, + *, + error: str, + status_code: int | None = None, + payload: Any = None, # noqa: ANN401 -- decoded JSON body, shape is server-defined + ) -> None: + super().__init__(message, status_code=status_code, payload=payload) + self.error = error + + +@dataclass(frozen=True) +class AuthServerMetadata: + authorization_endpoint: str + token_endpoint: str + registration_endpoint: str + issuer: str + + +def _b64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def pkce_pair() -> tuple[str, str]: + verifier = _b64url(secrets.token_bytes(PKCE_VERIFIER_BYTES)) + return verifier, _b64url(hashlib.sha256(verifier.encode("ascii")).digest()) + + +def _payload(response: httpx2.Response) -> dict[str, Any]: + try: + payload = response.json() + except ValueError as exc: + raise AuthenticationError( + f"OAuth server returned a non-JSON response (HTTP {response.status_code})", status_code=response.status_code + ) from exc + if isinstance(payload, dict) and "error" in payload: + description = payload.get("error_description") + error = str(payload["error"]) + message = f"{error}: {description}" if description else error + raise OAuthError(message, error=error, status_code=response.status_code, payload=payload) + if response.status_code >= 400 or not isinstance(payload, dict): + raise AuthenticationError( + f"OAuth server returned HTTP {response.status_code}", status_code=response.status_code, payload=payload + ) + return payload + + +def _require(payload: dict[str, Any], key: str) -> str: + if key not in payload: + raise AuthenticationError(f"OAuth server response is missing `{key}`", payload=payload) + return str(payload[key]) + + +def _redacted(payload: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in payload.items() if key not in TOKEN_KEYS} + + +def _credential_from_token_payload( + payload: dict[str, Any], *, client_id: str, token_endpoint: str, fallback_refresh_token: str | None +) -> OAuthCredential: + # Exceptions raised here may be logged by SDK consumers; never attach live tokens to them. + safe_payload = _redacted(payload) + refresh_token = payload.get("refresh_token") or fallback_refresh_token + if not refresh_token: + raise AuthenticationError("OAuth token response has no `refresh_token`", payload=safe_payload) + for key in ("access_token", "expires_in"): + if key not in payload: + raise AuthenticationError(f"OAuth server response is missing `{key}`", payload=safe_payload) + return OAuthCredential( + access_token=str(payload["access_token"]), + refresh_token=str(refresh_token), + expires_at=time.time() + float(payload["expires_in"]), + client_id=client_id, + token_endpoint=token_endpoint, + ) + + +def discover(base_url: str, *, client: httpx2.Client) -> AuthServerMetadata: + payload = _payload(client.get(base_url.rstrip("/") + METADATA_PATH)) + return AuthServerMetadata( + authorization_endpoint=_require(payload, "authorization_endpoint"), + token_endpoint=_require(payload, "token_endpoint"), + registration_endpoint=_require(payload, "registration_endpoint"), + issuer=_require(payload, "issuer"), + ) + + +def register_client(metadata: AuthServerMetadata, *, redirect_uris: list[str], client: httpx2.Client) -> str: + body = { + "client_name": CLIENT_NAME, + "redirect_uris": redirect_uris, + "grant_types": GRANT_TYPES, + "response_types": RESPONSE_TYPES, + "token_endpoint_auth_method": "none", + } + return _require(_payload(client.post(metadata.registration_endpoint, json=body)), "client_id") + + +def build_authorization_url( + metadata: AuthServerMetadata, + *, + client_id: str, + redirect_uri: str, + code_challenge: str, + state: str, + resource: str, + scope: str = OAUTH_SCOPE, +) -> str: + query = urlencode( + { + "response_type": "code", + "client_id": client_id, + "redirect_uri": redirect_uri, + "code_challenge": code_challenge, + "code_challenge_method": PKCE_METHOD, + "state": state, + "resource": resource, + "scope": scope, + } + ) + separator = "&" if "?" in metadata.authorization_endpoint else "?" + return f"{metadata.authorization_endpoint}{separator}{query}" + + +def exchange_code( + metadata: AuthServerMetadata, + *, + client_id: str, + code: str, + code_verifier: str, + redirect_uri: str, + resource: str, + client: httpx2.Client, +) -> OAuthCredential: + form = { + "grant_type": "authorization_code", + "client_id": client_id, + "code": code, + "code_verifier": code_verifier, + "redirect_uri": redirect_uri, + "resource": resource, + } + response = client.post(metadata.token_endpoint, data=form, headers=TOKEN_HEADERS) + return _credential_from_token_payload( + _payload(response), client_id=client_id, token_endpoint=metadata.token_endpoint, fallback_refresh_token=None + ) + + +def refresh_request(credential: OAuthCredential) -> httpx2.Request: + form = { + "grant_type": "refresh_token", + "refresh_token": credential.refresh_token, + "client_id": credential.client_id, + } + return httpx2.Request("POST", credential.token_endpoint, data=form, headers=TOKEN_HEADERS) + + +def parse_refresh_response(response: httpx2.Response, *, credential: OAuthCredential) -> OAuthCredential: + return _credential_from_token_payload( + _payload(response), + client_id=credential.client_id, + token_endpoint=credential.token_endpoint, + fallback_refresh_token=credential.refresh_token, + ) diff --git a/packages/discolike/src/discolike/_transport.py b/packages/discolike/src/discolike/_transport.py index 95296cd..0b54f57 100644 --- a/packages/discolike/src/discolike/_transport.py +++ b/packages/discolike/src/discolike/_transport.py @@ -22,8 +22,8 @@ def drop_none(params: Mapping[str, Any] | None) -> dict[str, Any]: return {key: value for key, value in (params or {}).items() if value is not None} -def _default_headers(api_key: str) -> dict[str, str]: - return {"X-discolike-key": api_key, "User-Agent": f"discolike-python/{__version__}"} +def _default_headers() -> dict[str, str]: + return {"User-Agent": f"discolike-python/{__version__}"} def _retryable_statuses(method: str) -> frozenset[int]: @@ -45,7 +45,7 @@ def _retry_delay(response: httpx2.Response | None, attempt: int) -> float: class Transport: def __init__( self, - api_key: str, + auth: httpx2.Auth, *, base_url: str, timeout: float, @@ -55,7 +55,8 @@ def __init__( if http_client is not None and not str(http_client.base_url): http_client.base_url = base_url self._client = http_client or httpx2.Client(base_url=base_url, timeout=timeout) - self._client.headers.update(_default_headers(api_key)) + self._client.auth = auth + self._client.headers.update(_default_headers()) self._max_retries = max_retries self._timeout_override: float | httpx2.Timeout | None = None self._is_view = False @@ -109,7 +110,7 @@ def close(self) -> None: class AsyncTransport: def __init__( self, - api_key: str, + auth: httpx2.Auth, *, base_url: str, timeout: float, @@ -119,7 +120,8 @@ def __init__( if http_client is not None and not str(http_client.base_url): http_client.base_url = base_url self._client = http_client or httpx2.AsyncClient(base_url=base_url, timeout=timeout) - self._client.headers.update(_default_headers(api_key)) + self._client.auth = auth + self._client.headers.update(_default_headers()) self._max_retries = max_retries self._timeout_override: float | httpx2.Timeout | None = None self._is_view = False diff --git a/packages/discolike/tests/test_auth_flow.py b/packages/discolike/tests/test_auth_flow.py new file mode 100644 index 0000000..a4eb24b --- /dev/null +++ b/packages/discolike/tests/test_auth_flow.py @@ -0,0 +1,243 @@ +import asyncio +import threading +import time + +import httpx2 +import pytest + +from discolike import AuthenticationError +from discolike._auth import DiscolikeAuth +from discolike._credentials import ApiKeyCredential +from discolike._credentials import OAuthCredential +from discolike._oauth import REFRESH_LEEWAY_SECONDS + +API = "https://api.test/v1" +TOKEN_ENDPOINT = "https://auth.test/oauth/2.1/token" +ONE_HOUR = 3600.0 + + +def make_oauth(*, expires_in: float = ONE_HOUR, access_token: str = "at-1") -> OAuthCredential: + return OAuthCredential( + access_token=access_token, + refresh_token="rt-1", + expires_at=time.time() + expires_in, + client_id="client-1", + token_endpoint=TOKEN_ENDPOINT, + ) + + +class Server: + """Mock API + token endpoint; counts refreshes and rejects any bearer it did not mint.""" + + def __init__(self, *, valid_tokens: set[str], refresh_ok: bool = True) -> None: + self.valid_tokens = valid_tokens + self.refresh_ok = refresh_ok + self.refreshes = 0 + self.bearers: list[str] = [] + self.token_calls: list[httpx2.Request] = [] + + def __call__(self, request: httpx2.Request) -> httpx2.Response: + if str(request.url) == TOKEN_ENDPOINT: + self.token_calls.append(request) + if not self.refresh_ok: + return httpx2.Response(400, json={"error": "invalid_grant", "error_description": "revoked"}) + self.refreshes += 1 + token = f"at-refreshed-{self.refreshes}" + self.valid_tokens.add(token) + return httpx2.Response( + 200, json={"access_token": token, "refresh_token": f"rt-{self.refreshes + 1}", "expires_in": 3600} + ) + bearer = request.headers.get("Authorization", "") + self.bearers.append(bearer) + if bearer.removeprefix("Bearer ") in self.valid_tokens: + return httpx2.Response(200, json={"ok": True}) + return httpx2.Response(401, json={"detail": "Invalid API Key or Session"}) + + +def sync_client(auth: DiscolikeAuth, handler) -> httpx2.Client: + return httpx2.Client(transport=httpx2.MockTransport(handler), base_url=API, auth=auth) + + +def async_client(auth: DiscolikeAuth, handler) -> httpx2.AsyncClient: + return httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url=API, auth=auth) + + +def test_api_key_credential_sets_header() -> None: + seen: dict[str, str] = {} + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.update(request.headers) + return httpx2.Response(200) + + sync_client(DiscolikeAuth(ApiKeyCredential(api_key="dk-1")), handler).get("/usage") + assert seen["x-discolike-key"] == "dk-1" + assert "authorization" not in seen + + +def test_oauth_credential_sets_bearer_without_refresh() -> None: + server = Server(valid_tokens={"at-1"}) + response = sync_client(DiscolikeAuth(make_oauth()), server).get("/usage") + assert response.status_code == 200 + assert server.refreshes == 0 + assert server.bearers[0] == "Bearer at-1" + + +def test_proactive_refresh_when_near_expiry() -> None: + server = Server(valid_tokens={"at-1"}) + updates: list[OAuthCredential] = [] + auth = DiscolikeAuth(make_oauth(expires_in=REFRESH_LEEWAY_SECONDS / 2), on_update=updates.append) + response = sync_client(auth, server).get("/usage") + assert response.status_code == 200 + assert server.refreshes == 1 + assert server.bearers[0] == "Bearer at-refreshed-1" + assert [update.access_token for update in updates] == ["at-refreshed-1"] + assert updates[0].refresh_token == "rt-2" + assert updates[0].client_id == "client-1" + assert updates[0].token_endpoint == TOKEN_ENDPOINT + assert auth.credential is updates[0] + + +def test_refresh_request_shape() -> None: + server = Server(valid_tokens=set()) + sync_client(DiscolikeAuth(make_oauth(expires_in=0)), server).get("/usage") + token_request = server.token_calls[0] + assert token_request.method == "POST" + assert token_request.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert token_request.content.decode() == "grant_type=refresh_token&refresh_token=rt-1&client_id=client-1" + + +def test_401_triggers_refresh_and_single_replay() -> None: + server = Server(valid_tokens=set()) + updates: list[OAuthCredential] = [] + response = sync_client(DiscolikeAuth(make_oauth(), on_update=updates.append), server).get("/usage") + assert response.status_code == 200 + assert server.refreshes == 1 + assert server.bearers == ["Bearer at-1", "Bearer at-refreshed-1"] + assert updates[0].access_token == "at-refreshed-1" + + +def test_second_401_after_replay_is_returned_not_retried() -> None: + server = Server(valid_tokens=set()) + + def reject_everything(request: httpx2.Request) -> httpx2.Response: + response = server(request) + if str(request.url) != TOKEN_ENDPOINT: + return httpx2.Response(401, json={"detail": "Invalid API Key or Session"}) + return response + + response = sync_client(DiscolikeAuth(make_oauth()), reject_everything).get("/usage") + assert response.status_code == 401 + assert server.refreshes == 1 + assert len(server.bearers) == 2 + + +def test_refresh_failure_raises_authentication_error() -> None: + server = Server(valid_tokens=set(), refresh_ok=False) + with pytest.raises(AuthenticationError, match="discolike auth login"): + sync_client(DiscolikeAuth(make_oauth(expires_in=0)), server).get("/usage") + + +def test_concurrent_requests_refresh_once() -> None: + server = Server(valid_tokens=set()) + auth = DiscolikeAuth(make_oauth(expires_in=0)) + client = sync_client(auth, server) + statuses: list[int] = [] + + def worker() -> None: + statuses.append(client.get("/usage").status_code) + + threads = [threading.Thread(target=worker) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert statuses == [200, 200, 200, 200] + assert server.refreshes == 1 + + +async def test_async_bearer_and_401_replay() -> None: + server = Server(valid_tokens=set()) + updates: list[OAuthCredential] = [] + async with async_client(DiscolikeAuth(make_oauth(), on_update=updates.append), server) as client: + response = await client.get("/usage") + assert response.status_code == 200 + assert server.refreshes == 1 + assert server.bearers == ["Bearer at-1", "Bearer at-refreshed-1"] + assert updates[0].access_token == "at-refreshed-1" + + +async def test_async_proactive_refresh_and_concurrency() -> None: + server = Server(valid_tokens=set()) + async with async_client(DiscolikeAuth(make_oauth(expires_in=0)), server) as client: + responses = await asyncio.gather(*(client.get("/usage") for _ in range(4))) + assert [response.status_code for response in responses] == [200, 200, 200, 200] + assert server.refreshes == 1 + + +async def test_async_refresh_failure_raises_authentication_error() -> None: + server = Server(valid_tokens=set(), refresh_ok=False) + async with async_client(DiscolikeAuth(make_oauth(expires_in=0)), server) as client: + with pytest.raises(AuthenticationError, match="discolike auth login"): + await client.get("/usage") + + +async def test_async_api_key_header() -> None: + seen: dict[str, str] = {} + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.update(request.headers) + return httpx2.Response(200) + + async with async_client(DiscolikeAuth(ApiKeyCredential(api_key="dk-2")), handler) as client: + await client.get("/usage") + assert seen["x-discolike-key"] == "dk-2" + + +def test_reload_adopts_credential_rotated_by_another_process() -> None: + server = Server(valid_tokens={"at-other"}) + updates: list[OAuthCredential] = [] + rotated_elsewhere = make_oauth(access_token="at-other") + auth = DiscolikeAuth(make_oauth(expires_in=0), on_update=updates.append, reload=lambda: rotated_elsewhere) + response = sync_client(auth, server).get("/usage") + assert response.status_code == 200 + assert server.token_calls == [] + assert server.bearers == ["Bearer at-other"] + assert auth.credential is rotated_elsewhere + assert updates == [] + + +def test_reload_with_stale_stored_credential_still_refreshes() -> None: + server = Server(valid_tokens=set()) + stale = make_oauth(expires_in=0) + auth = DiscolikeAuth(stale, reload=lambda: stale) + response = sync_client(auth, server).get("/usage") + assert response.status_code == 200 + assert server.refreshes == 1 + + +def test_reload_with_api_key_or_missing_config_still_refreshes() -> None: + server = Server(valid_tokens=set()) + auth = DiscolikeAuth(make_oauth(expires_in=0), reload=lambda: None) + assert sync_client(auth, server).get("/usage").status_code == 200 + assert server.refreshes == 1 + + +def test_reload_adopts_after_401() -> None: + server = Server(valid_tokens={"at-other"}) + rotated_elsewhere = make_oauth(access_token="at-other") + auth = DiscolikeAuth(make_oauth(), reload=lambda: rotated_elsewhere) + response = sync_client(auth, server).get("/usage") + assert response.status_code == 200 + assert server.token_calls == [] + assert server.bearers == ["Bearer at-1", "Bearer at-other"] + + +async def test_async_reload_adopts_credential_rotated_by_another_process() -> None: + server = Server(valid_tokens={"at-other"}) + rotated_elsewhere = make_oauth(access_token="at-other") + auth = DiscolikeAuth(make_oauth(expires_in=0), reload=lambda: rotated_elsewhere) + async with async_client(auth, server) as client: + response = await client.get("/usage") + assert response.status_code == 200 + assert server.token_calls == [] + assert server.bearers == ["Bearer at-other"] diff --git a/packages/discolike/tests/test_client.py b/packages/discolike/tests/test_client.py index 9232cbe..fecf942 100644 --- a/packages/discolike/tests/test_client.py +++ b/packages/discolike/tests/test_client.py @@ -1,8 +1,16 @@ +import time + import httpx2 import pytest +from discolike import AsyncDiscolike from discolike import AuthenticationError from discolike import Discolike +from discolike import OAuthCredential +from discolike._auth import DiscolikeAuth +from discolike._config import config_path +from discolike._config import load_credential +from discolike._config import save_credential from discolike_testkit import AsyncClientFactory from discolike_testkit import ClientFactory @@ -92,3 +100,96 @@ def handler(request: httpx2.Request) -> httpx2.Response: async with client.with_options(timeout=30.0) as view: await view.account.usage() await client.account.usage() + + +def test_client_accepts_injected_credential_and_does_not_persist_refresh(monkeypatch) -> None: + credential = OAuthCredential( + access_token="at", refresh_token="rt", expires_at=time.time() + 3600, client_id="c", token_endpoint="https://t" + ) + seen: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.append(request.headers["Authorization"]) + return httpx2.Response(200, json={"requests_mtd": 1}) + + http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") + with Discolike(auth=credential, base_url="https://api.test/v1", http_client=http) as client: + client.account.usage() + auth = client._transport._client.auth + assert isinstance(auth, DiscolikeAuth) + assert auth.on_update is None + assert seen == ["Bearer at"] + assert not config_path().exists() + + +def test_client_from_config_oauth_persists_rotated_tokens() -> None: + save_credential( + OAuthCredential( + access_token="stale", refresh_token="rt-1", expires_at=0.0, client_id="c", token_endpoint="https://t/token" + ) + ) + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == "https://t/token": + return httpx2.Response(200, json={"access_token": "fresh", "refresh_token": "rt-2", "expires_in": 3600}) + assert request.headers["Authorization"] == "Bearer fresh" + return httpx2.Response(200, json={"requests_mtd": 1}) + + http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") + with Discolike(base_url="https://api.test/v1", http_client=http) as client: + assert client.account.usage().requests_mtd == 1 + stored = load_credential() + assert isinstance(stored, OAuthCredential) + assert (stored.access_token, stored.refresh_token) == ("fresh", "rt-2") + + +def test_with_options_view_shares_auth(make_client: ClientFactory) -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"balance": 1}) + + with make_client(handler) as client: + assert client.with_options(timeout=1.0)._transport._client.auth is client._transport._client.auth + + +async def test_async_client_accepts_injected_credential() -> None: + seen: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.append(request.headers["Authorization"]) + return httpx2.Response(200, json={"requests_mtd": 1}) + + credential = OAuthCredential( + access_token="at", refresh_token="rt", expires_at=time.time() + 3600, client_id="c", token_endpoint="https://t" + ) + http = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") + async with AsyncDiscolike(auth=credential, base_url="https://api.test/v1", http_client=http) as client: + await client.account.usage() + assert seen == ["Bearer at"] + + +def test_client_from_config_reloads_before_refreshing() -> None: + save_credential( + OAuthCredential( + access_token="stale", refresh_token="rt-1", expires_at=0.0, client_id="c", token_endpoint="https://t/token" + ) + ) + seen: list[str] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + assert str(request.url) != "https://t/token" + seen.append(request.headers["Authorization"]) + return httpx2.Response(200, json={"requests_mtd": 1}) + + http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") + with Discolike(base_url="https://api.test/v1", http_client=http) as client: + save_credential( + OAuthCredential( + access_token="fresh-elsewhere", + refresh_token="rt-2", + expires_at=time.time() + 3600, + client_id="c", + token_endpoint="https://t/token", + ) + ) + client.account.usage() + assert seen == ["Bearer fresh-elsewhere"] diff --git a/packages/discolike/tests/test_config.py b/packages/discolike/tests/test_config.py index 9aaee96..5daa832 100644 --- a/packages/discolike/tests/test_config.py +++ b/packages/discolike/tests/test_config.py @@ -4,9 +4,18 @@ from discolike import AuthenticationError from discolike._config import config_path +from discolike._config import delete_credential +from discolike._config import delete_oauth_client from discolike._config import load_config -from discolike._config import resolve_api_key +from discolike._config import load_credential +from discolike._config import load_oauth_client +from discolike._config import resolve_credential from discolike._config import save_config +from discolike._config import save_credential +from discolike._config import save_oauth_client +from discolike._credentials import ApiKeyCredential +from discolike._credentials import OAuthClientRegistration +from discolike._credentials import OAuthCredential @pytest.fixture(autouse=True) @@ -37,22 +46,6 @@ def test_load_missing_returns_empty(isolated_config) -> None: assert load_config() == {} -def test_resolve_precedence_explicit_wins(isolated_config, monkeypatch) -> None: - save_config({"auth_method": "api_key", "api_key": "from-file"}) - monkeypatch.setenv("DISCOLIKE_API_KEY", "from-env") - assert resolve_api_key("explicit") == "explicit" - assert resolve_api_key(None) == "from-env" - monkeypatch.delenv("DISCOLIKE_API_KEY") - assert resolve_api_key(None) == "from-file" - - -def test_resolve_nothing_raises_with_guidance(isolated_config) -> None: - with pytest.raises(AuthenticationError) as exc_info: - resolve_api_key(None) - assert "DISCOLIKE_API_KEY" in str(exc_info.value) - assert "discolike auth login" in str(exc_info.value) - - def test_corrupt_config_returns_empty(isolated_config) -> None: config_path().parent.mkdir(parents=True, exist_ok=True) config_path().write_text("{not json") @@ -63,3 +56,137 @@ def test_binary_garbage_config_returns_empty(isolated_config) -> None: config_path().parent.mkdir(parents=True, exist_ok=True) config_path().write_bytes(b"\xff\xfe\x00garbage") assert load_config() == {} + + +def _oauth_credential() -> OAuthCredential: + return OAuthCredential( + access_token="at", refresh_token="rt", expires_at=1.0, client_id="c", token_endpoint="https://t/token" + ) + + +def test_save_and_load_oauth_credential(isolated_config) -> None: + save_credential(_oauth_credential()) + stored = load_config() + assert stored["auth_method"] == "oauth" + assert stored["oauth"] == { + "access_token": "at", + "refresh_token": "rt", + "expires_at": 1.0, + "client_id": "c", + "token_endpoint": "https://t/token", + } + assert load_credential() == _oauth_credential() + + +def test_save_api_key_credential_keeps_legacy_shape(isolated_config) -> None: + save_credential(ApiKeyCredential(api_key="dk-1")) + assert load_config() == {"auth_method": "api_key", "api_key": "dk-1"} + + +def test_load_credential_without_auth_method_is_api_key(isolated_config) -> None: + save_config({"api_key": "legacy"}) + assert load_credential() == ApiKeyCredential(api_key="legacy") + assert resolve_credential() == ApiKeyCredential(api_key="legacy") + + +def test_load_credential_missing_returns_none(isolated_config) -> None: + assert load_credential() is None + + +def test_resolve_credential_precedence(isolated_config, monkeypatch) -> None: + save_credential(_oauth_credential()) + monkeypatch.setenv("DISCOLIKE_API_KEY", "from-env") + injected = ApiKeyCredential(api_key="injected") + assert resolve_credential(api_key="explicit", auth=injected) is injected + assert resolve_credential(api_key="explicit") == ApiKeyCredential(api_key="explicit") + assert resolve_credential() == ApiKeyCredential(api_key="from-env") + monkeypatch.delenv("DISCOLIKE_API_KEY") + assert resolve_credential() == _oauth_credential() + + +def test_resolve_credential_nothing_raises_with_guidance(isolated_config) -> None: + with pytest.raises(AuthenticationError, match="discolike auth login"): + resolve_credential() + + +def test_save_config_is_atomic_and_leaves_no_temp_files(isolated_config) -> None: + save_config({"auth_method": "api_key", "api_key": "first"}) + save_config({"auth_method": "api_key", "api_key": "second"}) + assert [entry.name for entry in config_path().parent.iterdir()] == [config_path().name] + assert load_config()["api_key"] == "second" + assert stat.S_IMODE(config_path().stat().st_mode) == 0o600 + + +REGISTRATION = OAuthClientRegistration( + client_id="client-1", redirect_uri="http://127.0.0.1:18484/callback", issuer="https://auth.test/oauth/2.1" +) + + +def test_oauth_client_registration_roundtrip_and_missing(isolated_config) -> None: + assert load_oauth_client() is None + save_oauth_client(REGISTRATION) + assert load_oauth_client() == REGISTRATION + assert load_config()["oauth_client"] == REGISTRATION.to_config() + + +def test_save_credential_preserves_oauth_client(isolated_config) -> None: + save_oauth_client(REGISTRATION) + save_credential(_oauth_credential()) + assert load_oauth_client() == REGISTRATION + assert load_credential() == _oauth_credential() + save_credential(ApiKeyCredential(api_key="dk-1")) + assert load_oauth_client() == REGISTRATION + assert load_credential() == ApiKeyCredential(api_key="dk-1") + + +def test_save_oauth_client_preserves_credential(isolated_config) -> None: + save_credential(_oauth_credential()) + save_oauth_client(REGISTRATION) + assert load_credential() == _oauth_credential() + + +def test_delete_credential_keeps_oauth_client(isolated_config) -> None: + save_oauth_client(REGISTRATION) + save_credential(_oauth_credential()) + delete_credential() + assert load_credential() is None + assert load_oauth_client() == REGISTRATION + assert load_config() == {"oauth_client": REGISTRATION.to_config()} + + +def test_delete_credential_without_oauth_client_removes_file(isolated_config) -> None: + save_credential(ApiKeyCredential(api_key="dk-1")) + delete_credential() + assert not config_path().exists() + delete_credential() + + +@pytest.mark.parametrize( + "config", + [ + {"auth_method": "oauth"}, + {"auth_method": "oauth", "oauth": "not-a-dict"}, + {"auth_method": "oauth", "oauth": {"access_token": "a", "refresh_token": "r"}}, + {"auth_method": "oauth", "oauth": {**_oauth_credential().to_config(), "expires_at": "soon"}}, + ], +) +def test_malformed_oauth_section_is_no_credential(isolated_config, config) -> None: + save_config(config) + assert load_credential() is None + with pytest.raises(AuthenticationError, match="discolike auth login"): + resolve_credential() + + +@pytest.mark.parametrize("stored", ["not-a-dict", {"client_id": "c"}, 7]) +def test_malformed_oauth_client_is_none(isolated_config, stored) -> None: + save_config({"oauth_client": stored}) + assert load_oauth_client() is None + + +def test_delete_oauth_client_keeps_credential(isolated_config) -> None: + save_credential(_oauth_credential()) + save_oauth_client(REGISTRATION) + delete_oauth_client() + assert load_oauth_client() is None + assert load_credential() == _oauth_credential() + delete_oauth_client() diff --git a/packages/discolike/tests/test_jobs.py b/packages/discolike/tests/test_jobs.py index c9d81ab..d71a463 100644 --- a/packages/discolike/tests/test_jobs.py +++ b/packages/discolike/tests/test_jobs.py @@ -9,6 +9,7 @@ from discolike._jobs import Job from discolike._transport import AsyncTransport from discolike._transport import Transport +from discolike_testkit import api_key_auth BASE = "https://api.test/v1" @@ -25,7 +26,7 @@ async def fake_sleep(seconds: float) -> None: def make_job(handler) -> Job: http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url=BASE) - transport = Transport("k", base_url=BASE, timeout=5.0, max_retries=0, http_client=http) + transport = Transport(api_key_auth("k"), base_url=BASE, timeout=5.0, max_retries=0, http_client=http) return Job(transport, task_family=FAMILY_DISCOGEN, task_id="t-1") @@ -123,6 +124,6 @@ async def test_async_job_wait() -> None: [{"status": "in_progress", "progress": 5}, {"status": "completed", "progress": 100, "results": []}] ) http = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url=BASE) - transport = AsyncTransport("k", base_url=BASE, timeout=5.0, max_retries=0, http_client=http) + transport = AsyncTransport(api_key_auth("k"), base_url=BASE, timeout=5.0, max_retries=0, http_client=http) final = await AsyncJob(transport, task_family=FAMILY_DISCOGEN, task_id="t-1").wait(timeout=60.0) assert final.status == "completed" diff --git a/packages/discolike/tests/test_oauth.py b/packages/discolike/tests/test_oauth.py new file mode 100644 index 0000000..daa4fb0 --- /dev/null +++ b/packages/discolike/tests/test_oauth.py @@ -0,0 +1,249 @@ +import base64 +import hashlib +import time +from urllib.parse import parse_qs +from urllib.parse import urlparse + +import httpx2 +import pytest + +from discolike import AuthenticationError +from discolike._credentials import OAuthCredential +from discolike._oauth import CLIENT_NAME +from discolike._oauth import AuthServerMetadata +from discolike._oauth import OAuthError +from discolike._oauth import build_authorization_url +from discolike._oauth import discover +from discolike._oauth import exchange_code +from discolike._oauth import parse_refresh_response +from discolike._oauth import pkce_pair +from discolike._oauth import refresh_request +from discolike._oauth import register_client + +BASE_URL = "https://api.test/v1" +METADATA = AuthServerMetadata( + authorization_endpoint="https://auth.test/oauth/2.1/authorize", + token_endpoint="https://auth.test/oauth/2.1/token", + registration_endpoint="https://auth.test/oauth/2.1/register", + issuer="https://auth.test/oauth/2.1", +) + + +def client_for(handler) -> httpx2.Client: + return httpx2.Client(transport=httpx2.MockTransport(handler)) + + +def form(request: httpx2.Request) -> dict[str, str]: + return {key: values[0] for key, values in parse_qs(request.content.decode()).items()} + + +def test_pkce_pair_is_s256() -> None: + verifier, challenge = pkce_pair() + expected = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() + assert challenge == expected + assert "=" not in verifier + assert 43 <= len(verifier) <= 128 + + +def test_discover_reads_well_known_under_base_url() -> None: + seen: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.append(request) + return httpx2.Response( + 200, + json={ + "issuer": METADATA.issuer, + "authorization_endpoint": METADATA.authorization_endpoint, + "token_endpoint": METADATA.token_endpoint, + "registration_endpoint": METADATA.registration_endpoint, + }, + ) + + assert discover(BASE_URL + "/", client=client_for(handler)) == METADATA + assert str(seen[0].url) == "https://api.test/v1/.well-known/oauth-authorization-server" + + +def test_discover_missing_field_raises() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"authorization_endpoint": "x", "token_endpoint": "y"}) + + with pytest.raises(AuthenticationError, match="registration_endpoint"): + discover(BASE_URL, client=client_for(handler)) + + +def test_register_client_sends_public_client_metadata() -> None: + seen: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.append(request) + return httpx2.Response(201, json={"client_id": "client-abc", "client_secret": ""}) + + client_id = register_client(METADATA, redirect_uris=["http://127.0.0.1:9999/callback"], client=client_for(handler)) + assert client_id == "client-abc" + request = seen[0] + assert request.method == "POST" + assert str(request.url) == METADATA.registration_endpoint + body = httpx2.Response(200, content=request.content).json() + assert body["client_name"] == CLIENT_NAME + assert body["redirect_uris"] == ["http://127.0.0.1:9999/callback"] + assert body["grant_types"] == ["authorization_code", "refresh_token"] + assert body["response_types"] == ["code"] + assert body["token_endpoint_auth_method"] == "none" + + +def test_build_authorization_url_carries_pkce_state_and_resource() -> None: + url = build_authorization_url( + METADATA, + client_id="client-abc", + redirect_uri="http://127.0.0.1:9999/callback", + code_challenge="chal", + state="st", + resource=BASE_URL, + ) + parsed = urlparse(url) + assert f"{parsed.scheme}://{parsed.netloc}{parsed.path}" == METADATA.authorization_endpoint + query = {key: values[0] for key, values in parse_qs(parsed.query).items()} + assert query == { + "response_type": "code", + "client_id": "client-abc", + "redirect_uri": "http://127.0.0.1:9999/callback", + "code_challenge": "chal", + "code_challenge_method": "S256", + "state": "st", + "resource": BASE_URL, + "scope": "offline_access", + } + + +def test_exchange_code_posts_form_and_builds_credential() -> None: + seen: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + seen.append(request) + return httpx2.Response( + 200, json={"access_token": "at", "refresh_token": "rt", "expires_in": 3600, "token_type": "Bearer"} + ) + + before = time.time() + credential = exchange_code( + METADATA, + client_id="client-abc", + code="the-code", + code_verifier="ver", + redirect_uri="http://127.0.0.1:9999/callback", + resource=BASE_URL, + client=client_for(handler), + ) + request = seen[0] + assert request.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert form(request) == { + "grant_type": "authorization_code", + "client_id": "client-abc", + "code": "the-code", + "code_verifier": "ver", + "redirect_uri": "http://127.0.0.1:9999/callback", + "resource": BASE_URL, + } + assert credential.access_token == "at" + assert credential.refresh_token == "rt" + assert credential.client_id == "client-abc" + assert credential.token_endpoint == METADATA.token_endpoint + assert before + 3600 <= credential.expires_at <= time.time() + 3600 + + +def test_exchange_code_without_refresh_token_raises() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"access_token": "at", "expires_in": 3600}) + + with pytest.raises(AuthenticationError, match="refresh_token"): + exchange_code( + METADATA, + client_id="c", + code="x", + code_verifier="v", + redirect_uri="http://127.0.0.1:1/callback", + resource=BASE_URL, + client=client_for(handler), + ) + + +def test_malformed_token_response_error_payload_carries_no_tokens() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"access_token": "at", "refresh_token": "rt", "token_type": "Bearer"}) + + with pytest.raises(AuthenticationError, match="expires_in") as info: + exchange_code( + METADATA, + client_id="c", + code="x", + code_verifier="v", + redirect_uri="http://127.0.0.1:1/callback", + resource=BASE_URL, + client=client_for(handler), + ) + assert info.value.payload == {"token_type": "Bearer"} + + +def test_oauth_error_body_maps_to_authentication_error() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(400, json={"error": "invalid_grant", "error_description": "code expired"}) + + with pytest.raises(OAuthError, match="invalid_grant: code expired") as exc_info: + exchange_code( + METADATA, + client_id="c", + code="x", + code_verifier="v", + redirect_uri="http://127.0.0.1:1/callback", + resource=BASE_URL, + client=client_for(handler), + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.error == "invalid_grant" + + +def test_non_json_error_maps_to_authentication_error() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(502, text="bad gateway") + + with pytest.raises(AuthenticationError, match="502"): + discover(BASE_URL, client=client_for(handler)) + + +def test_refresh_rotates_tokens_and_keeps_old_refresh_token_when_absent() -> None: + credential = OAuthCredential( + access_token="old", + refresh_token="rt-old", + expires_at=0.0, + client_id="c", + token_endpoint=METADATA.token_endpoint, + ) + seen: list[httpx2.Request] = [] + + def rotating(request: httpx2.Request) -> httpx2.Response: + seen.append(request) + return httpx2.Response(200, json={"access_token": "new", "refresh_token": "rt-new", "expires_in": 60}) + + with client_for(rotating) as client: + rotated = parse_refresh_response(client.send(refresh_request(credential)), credential=credential) + assert form(seen[0]) == {"grant_type": "refresh_token", "refresh_token": "rt-old", "client_id": "c"} + assert (rotated.access_token, rotated.refresh_token) == ("new", "rt-new") + assert rotated.client_id == "c" + assert rotated.token_endpoint == METADATA.token_endpoint + + def not_rotating(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, json={"access_token": "newer", "expires_in": 60}) + + with client_for(not_rotating) as client: + kept = parse_refresh_response(client.send(refresh_request(rotated)), credential=rotated) + assert (kept.access_token, kept.refresh_token) == ("newer", "rt-new") + + +def test_credential_config_roundtrip_and_expiry() -> None: + credential = OAuthCredential( + access_token="a", refresh_token="r", expires_at=1000.0, client_id="c", token_endpoint="https://t" + ) + assert OAuthCredential.from_config(credential.to_config()) == credential + assert credential.expires_within(60, now=950.0) + assert not credential.expires_within(60, now=900.0) diff --git a/packages/discolike/tests/test_transport.py b/packages/discolike/tests/test_transport.py index 69149cc..1d83e01 100644 --- a/packages/discolike/tests/test_transport.py +++ b/packages/discolike/tests/test_transport.py @@ -7,6 +7,7 @@ from discolike._transport import AsyncTransport from discolike._transport import Transport from discolike._transport import drop_none +from discolike_testkit import api_key_auth @pytest.fixture(autouse=True) @@ -23,7 +24,9 @@ async def fake_async_sleep(seconds: float) -> None: def make_transport(handler) -> Transport: http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") - return Transport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http) + return Transport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http + ) def test_drop_none() -> None: @@ -102,7 +105,9 @@ def handler(request: httpx2.Request) -> httpx2.Response: from discolike import ServerError http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") - transport = Transport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http) + transport = Transport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http + ) with pytest.raises(ServerError): transport.request("GET", "/usage") assert len(calls) == 1 @@ -117,7 +122,9 @@ def handler(request: httpx2.Request) -> httpx2.Response: return httpx2.Response(502) if len(calls) < 2 else httpx2.Response(200, json={"ok": True}) http = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") - transport = AsyncTransport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http) + transport = AsyncTransport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http + ) response = await transport.request("GET", "/usage") assert response.json() == {"ok": True} await transport.aclose() @@ -201,7 +208,9 @@ def handler(request: httpx2.Request) -> httpx2.Response: from discolike import ServerError http = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") - transport = AsyncTransport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http) + transport = AsyncTransport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http + ) with pytest.raises(ServerError): await transport.request("POST", "/discogen/process") assert len(calls) == 1 @@ -218,7 +227,9 @@ def handler(request: httpx2.Request) -> httpx2.Response: return httpx2.Response(200, json={"ok": True}) http = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1") - transport = AsyncTransport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http) + transport = AsyncTransport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=2, http_client=http + ) response = await transport.request("POST", "/discogen/process") assert response.json() == {"ok": True} assert len(calls) == 2 @@ -233,14 +244,16 @@ def handler(request: httpx2.Request) -> httpx2.Response: return httpx2.Response(200, json={"ok": True}) http = httpx2.Client(transport=httpx2.MockTransport(handler)) - transport = Transport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http) + transport = Transport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http + ) transport.request("GET", "/usage") assert seen["url"] == "https://api.test/v1/usage" def test_byo_http_client_with_base_url_is_left_alone() -> None: http = httpx2.Client(base_url="https://custom.example/v2") - Transport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http) + Transport(api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http) assert str(http.base_url) == "https://custom.example/v2/" @@ -252,7 +265,9 @@ def handler(request: httpx2.Request) -> httpx2.Response: return httpx2.Response(200, json={"ok": True}) http = httpx2.Client(transport=httpx2.MockTransport(handler), base_url="https://api.test/v1", timeout=5.0) - transport = Transport("test-key", base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http) + transport = Transport( + api_key_auth("test-key"), base_url="https://api.test/v1", timeout=5.0, max_retries=0, http_client=http + ) transport.with_timeout(120.0).request("GET", "/usage") transport.request("GET", "/usage") diff --git a/pyproject.toml b/pyproject.toml index a534ea7..f5aa784 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,4 +70,4 @@ ban-relative-imports = "all" force-single-line = true [tool.ruff.lint.per-file-ignores] -"packages/*/tests/*" = ["ANN"] +"packages/*/tests/*" = ["ANN", "S105", "S106", "S107"] # fake tokens in fixtures