'
-def _render_callback_html() -> str:
- return _CALLBACK_HTML.replace("", _logo_img_tag())
+def _render_callback_html(account_label: str = "your ChatGPT subscription") -> str:
+ return _CALLBACK_HTML.replace("", _logo_img_tag()).replace(
+ "", account_label
+ )
_CALLBACK_HTML = """
@@ -399,7 +658,7 @@ _CALLBACK_HTML = """
Strix is connected to your ChatGPT subscription. Head back to your +
Strix is connected to . Head back to your terminal — your security test runs there.
Autonomous AI hackers that find and fix your app's diff --git a/strix/interface/environment.py b/strix/interface/environment.py index 522067df..57562b85 100644 --- a/strix/interface/environment.py +++ b/strix/interface/environment.py @@ -8,7 +8,7 @@ from rich.console import Console from rich.panel import Panel from rich.text import Text -from strix.config import IntegrationSettings, codex, load_settings +from strix.config import IntegrationSettings, anthropic_wif, codex, load_settings from strix.interface.utils import ( check_docker_connection, image_exists, @@ -53,7 +53,30 @@ def validate_environment() -> None: if not settings.llm.model: missing_required_vars.append("STRIX_LLM") - if not settings.llm.api_key: + uses_wif = anthropic_wif.applies(settings.llm.model, settings.llm.api_key) + if uses_wif: + try: + anthropic_wif.get_valid_token() + except anthropic_wif.WifAuthError as exc: + console.print( + f"[red]Anthropic workload identity federation failed:[/] {exc}", + highlight=False, + ) + report_error("workload_identity_exchange_failed") + sys.exit(1) + logger.info("Anthropic workload identity token exchange OK") + elif ( + settings.llm.api_key + and anthropic_wif.is_anthropic_model(settings.llm.model or "") + and anthropic_wif.is_configured() + ): + console.print( + "[yellow]LLM_API_KEY is set, so it is used instead of the configured " + "Anthropic workload identity federation.[/] Unset it to federate.", + highlight=False, + ) + + if not settings.llm.api_key and not uses_wif: missing_optional_vars.append("LLM_API_KEY") if not settings.llm.api_base: diff --git a/tests/test_anthropic_wif.py b/tests/test_anthropic_wif.py new file mode 100644 index 00000000..03b82ace --- /dev/null +++ b/tests/test_anthropic_wif.py @@ -0,0 +1,476 @@ +"""Tests for Anthropic Workload Identity Federation: exchange, refresh, routing.""" + +from __future__ import annotations + +import asyncio +import json +import shlex +import subprocess +import sys +import time +from typing import TYPE_CHECKING, Any +from unittest import mock + +import pytest +import requests +from agents.extensions.models.litellm_model import LitellmModel + +from strix.config import anthropic_wif, codex, entra_user, loader +from strix.config.models import StrixProvider + + +if TYPE_CHECKING: + from pathlib import Path + + +_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_rule", + "ANTHROPIC_ORGANIZATION_ID": "00000000-0000-0000-0000-000000000000", + "ANTHROPIC_SERVICE_ACCOUNT_ID": "svac_sa", +} + + +@pytest.fixture +def token_file(tmp_path: Path) -> Path: + path = tmp_path / "jwt" + path.write_text("jwt-1\n", encoding="utf-8") + return path + + +@pytest.fixture(autouse=True) +def _isolated(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, token_file: Path) -> None: + for key in ( + "STRIX_LLM", + "LLM_API_KEY", + "OPENAI_API_KEY", + "LLM_API_BASE", + "OPENAI_API_BASE", + "OPENAI_BASE_URL", + "LITELLM_BASE_URL", + "OLLAMA_API_BASE", + "ANTHROPIC_WORKSPACE_ID", + "ANTHROPIC_IDENTITY_TOKEN", + "ANTHROPIC_IDENTITY_TOKEN_COMMAND", + "ACTIONS_ID_TOKEN_REQUEST_URL", + "ACTIONS_ID_TOKEN_REQUEST_TOKEN", + "STRIX_ENTRA_TENANT_ID", + "STRIX_ENTRA_APP_ID", + ): + monkeypatch.delenv(key, raising=False) + for key, value in _WIF_ENV.items(): + monkeypatch.setenv(key, value) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", str(token_file)) + monkeypatch.setattr(entra_user, "AUTH_PATH", tmp_path / "entra-auth.json") + monkeypatch.setattr(loader, "_override", tmp_path / "missing-config.json") + monkeypatch.setattr(loader, "_cached", None) + anthropic_wif.reset_cache() + + +def _response(status: int, body: dict[str, Any], headers: dict[str, str] | None = None) -> Any: + resp = mock.MagicMock(spec=requests.Response) + resp.status_code = status + resp.json.return_value = body + resp.headers = headers or {} + return resp + + +def _ok(token: str = "sk-ant-oat01-a", expires_in: int = 600) -> Any: # noqa: S107 + return _response(200, {"access_token": token, "token_type": "Bearer", "expires_in": expires_in}) + + +def test_configured_only_with_ids_and_a_token_source(monkeypatch: pytest.MonkeyPatch) -> None: + assert anthropic_wif.is_configured() + monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE") + monkeypatch.setattr(loader, "_cached", None) + assert not anthropic_wif.is_configured() + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "jwt-literal") + assert anthropic_wif.is_configured() + monkeypatch.delenv("ANTHROPIC_SERVICE_ACCOUNT_ID") + monkeypatch.setattr(loader, "_cached", None) + assert not anthropic_wif.is_configured() + + +def test_exchange_sends_jwt_bearer_grant(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ANTHROPIC_WORKSPACE_ID", "wrkspc_ws") + with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post: + token, expires_at = anthropic_wif.exchange() + assert token == "sk-ant-oat01-a" # noqa: S105 + assert expires_at > 0 + url = post.call_args.args[0] + assert url == "https://api.anthropic.com/v1/oauth/token" + assert post.call_args.kwargs["json"] == { + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "assertion": "jwt-1", + "federation_rule_id": "fdrl_rule", + "organization_id": "00000000-0000-0000-0000-000000000000", + "service_account_id": "svac_sa", + "workspace_id": "wrkspc_ws", + } + + +def test_exchange_omits_unset_workspace() -> None: + with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post: + anthropic_wif.exchange() + assert "workspace_id" not in post.call_args.kwargs["json"] + + +def test_token_url_follows_api_base() -> None: + assert anthropic_wif.token_url("https://gw.example/v1/") == "https://gw.example/v1/oauth/token" + assert anthropic_wif.token_url(None) == "https://api.anthropic.com/v1/oauth/token" + + +def test_exchange_rereads_the_token_file(token_file: Path) -> None: + with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post: + anthropic_wif.exchange() + token_file.write_text("jwt-2", encoding="utf-8") + anthropic_wif.exchange() + assertions = [call.kwargs["json"]["assertion"] for call in post.call_args_list] + assert assertions == ["jwt-1", "jwt-2"] + + +def test_exchange_error_names_status_but_not_the_jwt() -> None: + body = { + "type": "error", + "error": {"type": "authentication_error", "message": "Authentication failed"}, + } + resp = _response(401, body, {"request-id": "req_1"}) + with ( + mock.patch.object(anthropic_wif.requests, "post", return_value=resp), + pytest.raises(anthropic_wif.WifAuthError) as exc_info, + ): + anthropic_wif.exchange() + message = str(exc_info.value) + assert "HTTP 401" in message + assert "authentication_error" in message + assert "req_1" in message + assert "jwt-1" not in message + + +def test_github_actions_oidc_token_used_without_a_file(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE") + monkeypatch.setattr(loader, "_cached", None) + monkeypatch.setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "https://gh.example/token?x=1") + monkeypatch.setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "gh-req") + gh = _response(200, {"value": "gh-jwt"}) + with mock.patch.object(anthropic_wif.requests, "get", return_value=gh) as get: + assert anthropic_wif.read_identity_token() == "gh-jwt" + assert get.call_args.kwargs["params"] == {"audience": "https://api.anthropic.com"} + assert get.call_args.kwargs["headers"] == {"Authorization": "Bearer gh-req"} + + +def _python_command(code: str) -> str: + return shlex.join([sys.executable, "-c", code]) + + +def _use_token_command(monkeypatch: pytest.MonkeyPatch, command: str) -> None: + monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE") + monkeypatch.setattr(loader, "_cached", None) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_COMMAND", command) + + +def _completed(returncode: int = 0, stdout: str = "", stderr: str = "") -> Any: + return subprocess.CompletedProcess(["cmd"], returncode, stdout, stderr) + + +def test_token_command_used_without_a_file(monkeypatch: pytest.MonkeyPatch) -> None: + _use_token_command(monkeypatch, _python_command("print('cmd-jwt')")) + assert anthropic_wif.is_configured() + assert anthropic_wif.read_identity_token() == "cmd-jwt" + + +def test_token_file_wins_over_command(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_COMMAND", "should-not-run") + monkeypatch.setattr(loader, "_cached", None) + with mock.patch.object(anthropic_wif.subprocess, "run") as run: + assert anthropic_wif.read_identity_token() == "jwt-1" + run.assert_not_called() + + +def test_token_command_runs_on_every_exchange(monkeypatch: pytest.MonkeyPatch) -> None: + _use_token_command(monkeypatch, "issue-token --audience x") + outputs = [_completed(stdout="jwt-a\n"), _completed(stdout="jwt-b\n")] + with ( + mock.patch.object(anthropic_wif.subprocess, "run", side_effect=outputs) as run, + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post, + ): + anthropic_wif.exchange() + anthropic_wif.exchange() + assert run.call_args.args[0] == ["issue-token", "--audience", "x"] + assert run.call_args.kwargs["stdin"] == subprocess.DEVNULL + assertions = [call.kwargs["json"]["assertion"] for call in post.call_args_list] + assert assertions == ["jwt-a", "jwt-b"] + + +def test_token_command_failure_names_status_and_stderr_but_not_stdout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + code = "import sys; print('secret-jwt'); print('boom', file=sys.stderr); sys.exit(3)" + _use_token_command(monkeypatch, _python_command(code)) + with pytest.raises(anthropic_wif.WifAuthError) as exc: + anthropic_wif.read_identity_token() + assert "status 3" in str(exc.value) + assert "boom" in str(exc.value) + assert "secret-jwt" not in str(exc.value) + + +@pytest.mark.parametrize( + ("stdout", "message"), + [("\n", "printed no token"), ("jwt-a jwt-b\n", "single JWT")], +) +def test_token_command_output_must_be_one_jwt( + monkeypatch: pytest.MonkeyPatch, stdout: str, message: str +) -> None: + _use_token_command(monkeypatch, "issue-token") + with ( + mock.patch.object(anthropic_wif.subprocess, "run", return_value=_completed(stdout=stdout)), + pytest.raises(anthropic_wif.WifAuthError, match=message), + ): + anthropic_wif.read_identity_token() + + +def test_token_command_not_found(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _use_token_command(monkeypatch, str(tmp_path / "no-such-command")) + with pytest.raises(anthropic_wif.WifAuthError, match="could not run"): + anthropic_wif.read_identity_token() + + +def test_token_command_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + _use_token_command(monkeypatch, "issue-token") + timeout = subprocess.TimeoutExpired(["issue-token"], 30) + with ( + mock.patch.object(anthropic_wif.subprocess, "run", side_effect=timeout), + pytest.raises(anthropic_wif.WifAuthError, match="timed out"), + ): + anthropic_wif.read_identity_token() + + +def test_token_command_expands_home_in_program(monkeypatch: pytest.MonkeyPatch) -> None: + _use_token_command(monkeypatch, "~/bin/issue-token token") + monkeypatch.setenv("HOME", "/home/someone") + with mock.patch.object( + anthropic_wif.subprocess, "run", return_value=_completed(stdout="jwt\n") + ) as run: + anthropic_wif.read_identity_token() + assert run.call_args.args[0] == ["/home/someone/bin/issue-token", "token"] + + +def test_token_command_persists_and_is_read_from_config( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + config = tmp_path / "config.json" + config.write_text('{"env": {}}', encoding="utf-8") + monkeypatch.setattr(loader, "_override", config) + command = "issue-token --from-config" + _use_token_command(monkeypatch, command) + loader.persist_current() + saved = json.loads(config.read_text(encoding="utf-8"))["env"] + assert saved.get("ANTHROPIC_IDENTITY_TOKEN_COMMAND") == command + + monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_COMMAND") + monkeypatch.setattr(loader, "_cached", None) + with mock.patch.object( + anthropic_wif.subprocess, "run", return_value=_completed(stdout="jwt\n") + ) as run: + assert anthropic_wif.read_identity_token() == "jwt" + assert run.call_args.args[0] == ["issue-token", "--from-config"] + + +def test_empty_env_clears_saved_token_command( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + config = tmp_path / "config.json" + config.write_text( + json.dumps({"env": {"ANTHROPIC_IDENTITY_TOKEN_COMMAND": "issue-token"}}), encoding="utf-8" + ) + monkeypatch.setattr(loader, "_override", config) + _use_token_command(monkeypatch, "") + assert not anthropic_wif.is_configured() + loader.persist_current() + saved = json.loads(config.read_text(encoding="utf-8"))["env"] + assert "ANTHROPIC_IDENTITY_TOKEN_COMMAND" not in saved + + +def _sign_in_to_entra(**overrides: Any) -> None: + record = { + "type": "oauth", + "provider": "entra", + "tenant": "tenant-id", + "app": "app-id", + "access": "entra-jwt", + "refresh": "entra-refresh", + "expires_at": time.time() + 3600, + } + record.update(overrides) + entra_user.save_record(record) + + +def _without_token_file(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE") + monkeypatch.setattr(loader, "_cached", None) + + +def test_entra_sign_in_used_without_other_sources(monkeypatch: pytest.MonkeyPatch) -> None: + _without_token_file(monkeypatch) + assert not anthropic_wif.is_configured() + _sign_in_to_entra() + assert anthropic_wif.is_configured() + assert anthropic_wif.read_identity_token() == "entra-jwt" + + +def test_token_command_wins_over_entra_sign_in(monkeypatch: pytest.MonkeyPatch) -> None: + _sign_in_to_entra() + _use_token_command(monkeypatch, "issue-token") + with mock.patch.object( + anthropic_wif.subprocess, "run", return_value=_completed(stdout="cmd-jwt\n") + ): + assert anthropic_wif.read_identity_token() == "cmd-jwt" + + +def test_entra_refresh_failure_names_the_fix(monkeypatch: pytest.MonkeyPatch) -> None: + _without_token_file(monkeypatch) + _sign_in_to_entra(expires_at=time.time() - 1) + refused = entra_user.EntraAuthError("invalid_grant", "invalid_grant: AADSTS70000") + with ( + mock.patch.object(entra_user, "refresh_tokens", side_effect=refused), + pytest.raises(anthropic_wif.WifAuthError, match="strix auth login entra"), + ): + anthropic_wif.read_identity_token() + + +def test_cached_token_reused_until_advisory_window() -> None: + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1000.0), + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok(expires_in=600)) as post, + ): + assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a" + assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a" + assert post.call_count == 1 + + +def test_refreshes_inside_advisory_window() -> None: + with ( + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok("sk-ant-oat01-a")), + mock.patch.object(anthropic_wif.time, "time", return_value=1000.0), + ): + anthropic_wif.get_valid_token() + # expires at 1600; 1500 is inside the 120s advisory window. + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1500.0), + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok("sk-ant-oat01-b")), + ): + assert anthropic_wif.get_valid_token() == "sk-ant-oat01-b" + + +def test_advisory_refresh_failure_keeps_cached_token() -> None: + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1000.0), + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()), + ): + anthropic_wif.get_valid_token() + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1500.0), + mock.patch.object( + anthropic_wif.requests, "post", side_effect=requests.ConnectionError("down") + ), + ): + assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a" + + +def test_mandatory_refresh_failure_raises() -> None: + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1000.0), + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()), + ): + anthropic_wif.get_valid_token() + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1580.0), + mock.patch.object( + anthropic_wif.requests, "post", side_effect=requests.ConnectionError("down") + ), + pytest.raises(anthropic_wif.WifAuthError), + ): + anthropic_wif.get_valid_token() + + +def test_slow_failed_refresh_does_not_serve_an_expired_token() -> None: + with ( + mock.patch.object(anthropic_wif.time, "time", return_value=1000.0), + mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()), + ): + anthropic_wif.get_valid_token() + # Expires at 1600. The refresh starts at 1555, inside the advisory window but + # outside the mandatory one, and fails only after the token has expired. + clock = iter([1555.0, 1555.0, 1610.0, 1610.0]) + + def _slow_failure(*_args: Any, **_kwargs: Any) -> Any: + raise requests.ConnectionError("down") + + with ( + mock.patch.object(anthropic_wif.time, "time", side_effect=lambda: next(clock)), + mock.patch.object(anthropic_wif.requests, "post", side_effect=_slow_failure), + pytest.raises(anthropic_wif.WifAuthError), + ): + anthropic_wif.get_valid_token() + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("anthropic/claude-sonnet-5-5", True), + ("litellm/anthropic/claude-sonnet-5-5", True), + ("claude-sonnet-5-5", False), + ("bedrock/anthropic.claude-sonnet-5-5", False), + ("vertex_ai/claude-sonnet-5-5", False), + ("openrouter/anthropic/claude-sonnet-5-5", False), + ("openai/gpt-5", False), + ], +) +def test_applies_only_to_anthropic_direct(model: str, *, expected: bool) -> None: + assert anthropic_wif.applies(model) is expected + + +def test_explicit_api_key_wins() -> None: + assert not anthropic_wif.applies("anthropic/claude-sonnet-5-5", "sk-ant-api03-x") + + +def test_auth_mode_reports_workload_identity(monkeypatch: pytest.MonkeyPatch) -> None: + assert codex.auth_mode("anthropic/claude-sonnet-5-5") == "workload_identity" + assert codex.auth_mode("openai/gpt-5") == "api_key" + monkeypatch.setenv("LLM_API_KEY", "sk-ant-api03-x") + monkeypatch.setattr(loader, "_cached", None) + assert codex.auth_mode("anthropic/claude-sonnet-5-5") == "api_key" + + +def _unwrap(model: object) -> object: + while hasattr(model, "_inner"): + model = model._inner + return model + + +def test_provider_federates_anthropic_requests() -> None: + model = _unwrap(StrixProvider().get_model("anthropic/claude-sonnet-5-5")) + sent: dict[str, Any] = {} + + async def fake_fetch(self: Any, *_args: Any, **_kwargs: Any) -> str: + sent["api_key"] = self.api_key + return "response" + + with ( + mock.patch.object(LitellmModel, "_fetch_response", fake_fetch), + mock.patch.object(anthropic_wif, "get_valid_token", return_value="sk-ant-oat01-z"), + ): + result = asyncio.run(model._fetch_response()) # type: ignore[attr-defined] + assert result == "response" + assert sent["api_key"] == "sk-ant-oat01-z" + + +def test_provider_leaves_other_litellm_models_alone() -> None: + model = _unwrap(StrixProvider().get_model("deepseek/deepseek-chat")) + assert type(model).__name__ == "LitellmModel" + + +def test_provider_uses_api_key_when_set(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LLM_API_KEY", "sk-ant-api03-x") + monkeypatch.setattr(loader, "_cached", None) + model = _unwrap(StrixProvider().get_model("anthropic/claude-sonnet-5-5")) + assert type(model).__name__ == "LitellmModel" diff --git a/tests/test_auth_cli.py b/tests/test_auth_cli.py index 9b66a2db..8bf0d278 100644 --- a/tests/test_auth_cli.py +++ b/tests/test_auth_cli.py @@ -2,11 +2,14 @@ from __future__ import annotations +import base64 +import json +import time from typing import TYPE_CHECKING, Any import pytest -from strix.config import codex +from strix.config import anthropic_wif, codex, entra_user, loader from strix.interface import auth_cli @@ -17,6 +20,11 @@ if TYPE_CHECKING: @pytest.fixture(autouse=True) def _tmp_store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(codex, "AUTH_PATH", tmp_path / "home" / ".strix" / "subscription-auth.json") + monkeypatch.setattr(entra_user, "AUTH_PATH", tmp_path / "home" / ".strix" / "entra-auth.json") + for key in ("STRIX_ENTRA_TENANT_ID", "STRIX_ENTRA_APP_ID"): + monkeypatch.delenv(key, raising=False) + monkeypatch.setattr(loader, "_override", tmp_path / "missing-config.json") + monkeypatch.setattr(loader, "_cached", None) def test_login_provider_is_chatgpt() -> None: @@ -104,3 +112,193 @@ def test_login_accepts_provider_aliases(provider: str, monkeypatch: pytest.Monke assert auth_cli.run_auth(["login", provider]) == 0 assert reached["flow"] is True + + +def _entra_record(**overrides: Any) -> dict[str, Any]: + record = { + "type": "oauth", + "provider": "entra", + "tenant": "tenant-id", + "app": "app-id", + "access": "entra-access", + "refresh": "entra-refresh", + "username": "user@example.com", + "expires_at": time.time() + 3600, + } + record.update(overrides) + return record + + +def _codex_record() -> dict[str, Any]: + return { + "type": "oauth", + "provider": "codex", + "access": "a", + "refresh": "r", + "account_id": "acct", + "expires_at": time.time() + 3600, + } + + +def test_login_entra_manual_paste_saves_the_sign_in(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[tuple[str, ...]] = [] + + def _exchange(tenant: str, app: str, code: str, verifier: str) -> dict[str, Any]: + calls.append((tenant, app, code)) + assert verifier + return _entra_record(tenant=tenant, app=app) + + monkeypatch.setattr(entra_user, "exchange_code", _exchange) + monkeypatch.setattr( + auth_cli.Console, "input", lambda *_args, **_kwargs: "http://localhost:8400/?code=abc" + ) + + argv = ["login", "entra", "--manual", "--tenant", "tenant-id", "--app", "app-id"] + assert auth_cli.run_auth(argv) == 0 + assert calls == [("tenant-id", "app-id", "abc")] + saved = entra_user.read_record() + assert saved is not None + assert saved["tenant"] == "tenant-id" + + +def test_login_entra_uses_settings_for_ids(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("STRIX_ENTRA_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("STRIX_ENTRA_APP_ID", "app-from-env") + seen: dict[str, str] = {} + + def _fake_flow(_console: Any, url: str, *_args: Any, **_kwargs: Any) -> dict[str, Any]: + seen["url"] = url + return _entra_record() + + monkeypatch.setattr(auth_cli, "_run_oauth_flow", _fake_flow) + assert auth_cli.run_auth(["login", "entra"]) == 0 + assert "/tenant-from-env/oauth2/v2.0/authorize" in seen["url"] + assert "client_id=app-from-env" in seen["url"] + + +def test_login_entra_without_ids_is_a_usage_error(monkeypatch: pytest.MonkeyPatch) -> None: + def _should_not_run(*_args: Any, **_kwargs: Any) -> dict[str, Any]: + msg = "OAuth flow must not start without Entra IDs" + raise AssertionError(msg) + + monkeypatch.setattr(auth_cli, "_run_oauth_flow", _should_not_run) + assert auth_cli.run_auth(["login", "entra"]) == 2 + + +def test_status_reports_both_providers(capsys: pytest.CaptureFixture[str]) -> None: + entra_user.save_record(_entra_record()) + assert auth_cli.run_auth(["status"]) == 0 + out = capsys.readouterr().out + assert "ChatGPT subscription" in out + assert "Not signed in" in out + assert "Microsoft Entra ID" in out + assert "user@example.com" in out + + +def test_logout_signs_out_of_every_provider(capsys: pytest.CaptureFixture[str]) -> None: + codex.save_record(_codex_record()) + entra_user.save_record(_entra_record()) + assert auth_cli.run_auth(["logout"]) == 0 + out = capsys.readouterr().out + assert "ChatGPT subscription and Microsoft Entra ID" in out + assert codex.read_record() is None + assert entra_user.read_record() is None + + +def test_logout_when_signed_out(capsys: pytest.CaptureFixture[str]) -> None: + assert auth_cli.run_auth(["logout"]) == 0 + assert "Not signed in" in capsys.readouterr().out + + +_WIF_ENV = { + "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_test", + "ANTHROPIC_ORGANIZATION_ID": "org-test", + "ANTHROPIC_SERVICE_ACCOUNT_ID": "svac_test", + "STRIX_LLM": "anthropic/claude-sonnet-4-6", +} + + +def _fake_jwt(claims: dict[str, Any]) -> str: + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + return f"e30.{payload}.sig" + + +@pytest.fixture +def wif_env(monkeypatch: pytest.MonkeyPatch) -> str: + for key in ( + "ANTHROPIC_WORKSPACE_ID", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_TOKEN_COMMAND", + "ACTIONS_ID_TOKEN_REQUEST_URL", + ): + monkeypatch.delenv(key, raising=False) + for key, value in _WIF_ENV.items(): + monkeypatch.setenv(key, value) + monkeypatch.setenv("LLM_API_KEY", "") + jwt = _fake_jwt({"iss": "https://idp.example", "sub": "subject-123"}) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", jwt) + return jwt + + +def test_test_reports_source_claims_and_prefix_only( + wif_env: str, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + seen: list[str | None] = [] + + def _exchange(assertion: str | None = None) -> tuple[str, float]: + seen.append(assertion) + return "sk-ant-oat01-SECRETPART", time.time() + 600 + + monkeypatch.setattr(anthropic_wif, "exchange", _exchange) + assert auth_cli.run_auth(["test"]) == 0 + out = capsys.readouterr().out + assert seen == [wif_env] + assert "ANTHROPIC_IDENTITY_TOKEN" in out + assert "subject-123" in out + assert "sk-ant-oat01-…" in out + assert "SECRETPART" not in out + assert wif_env not in out + assert "Note:" not in out + + +def test_test_without_federation_ids(monkeypatch: pytest.MonkeyPatch) -> None: + for key in _WIF_ENV: + monkeypatch.delenv(key, raising=False) + assert auth_cli.run_auth(["test"]) == 1 + + +def test_test_reports_exchange_failure( + wif_env: str, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + del wif_env + + def _fail(_assertion: str | None = None) -> tuple[str, float]: + raise anthropic_wif.WifAuthError("token exchange failed: HTTP 401") + + monkeypatch.setattr(anthropic_wif, "exchange", _fail) + assert auth_cli.run_auth(["test"]) == 1 + assert "HTTP 401" in capsys.readouterr().out + + +def test_test_warns_when_an_api_key_wins( + wif_env: str, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + del wif_env + monkeypatch.setenv("LLM_API_KEY", "sk-ant-api-key") + monkeypatch.setattr(anthropic_wif, "exchange", lambda _a=None: ("sk-ant-oat01-x", time.time())) + assert auth_cli.run_auth(["test"]) == 0 + assert "LLM_API_KEY is set" in capsys.readouterr().out + + +def test_test_reads_settings_from_config( + wif_env: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + del wif_env + config = tmp_path / "strix-config.json" + config.write_text(json.dumps({"env": _WIF_ENV}), encoding="utf-8") + for key in _WIF_ENV: + monkeypatch.delenv(key, raising=False) + monkeypatch.setattr(anthropic_wif, "exchange", lambda _a=None: ("sk-ant-oat01-x", time.time())) + assert auth_cli.run_auth(["test"]) == 1 + monkeypatch.setattr(loader, "_cached", None) + assert auth_cli.run_auth(["test", "--config", str(config)]) == 0 diff --git a/tests/test_entra_user.py b/tests/test_entra_user.py new file mode 100644 index 00000000..bd6d9fff --- /dev/null +++ b/tests/test_entra_user.py @@ -0,0 +1,208 @@ +"""Tests for Microsoft Entra ID user sign-in (the `strix auth login entra` token store).""" + +from __future__ import annotations + +import base64 +import json +import threading +import time +import urllib.parse +from typing import TYPE_CHECKING, Any +from unittest import mock + +import pytest + +from strix.config import entra_user, loader + + +if TYPE_CHECKING: + from pathlib import Path + + +TENANT = "11111111-1111-1111-1111-111111111111" +APP = "22222222-2222-2222-2222-222222222222" + + +@pytest.fixture(autouse=True) +def _isolated(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(entra_user, "AUTH_PATH", tmp_path / ".strix" / "entra-auth.json") + for key in ("STRIX_ENTRA_TENANT_ID", "STRIX_ENTRA_APP_ID"): + monkeypatch.delenv(key, raising=False) + monkeypatch.setattr(loader, "_override", tmp_path / "missing-config.json") + monkeypatch.setattr(loader, "_cached", None) + + +def _jwt(claims: dict[str, Any]) -> str: + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + return f"e30.{payload}.sig" + + +def _response(status: int, body: dict[str, Any]) -> Any: + resp = mock.MagicMock() + resp.status_code = status + resp.json.return_value = body + return resp + + +def _record(**overrides: Any) -> dict[str, Any]: + record = { + "type": "oauth", + "provider": "entra", + "tenant": TENANT, + "app": APP, + "access": "access-1", + "refresh": "refresh-1", + "username": "user@example.com", + "expires_at": time.time() + 3600, + } + record.update(overrides) + return record + + +def test_authorize_url_requests_the_app_scope_with_pkce() -> None: + url = urllib.parse.urlparse(entra_user.build_authorize_url(TENANT, APP, "chal", "st")) + params = dict(urllib.parse.parse_qsl(url.query)) + assert url.netloc == "login.microsoftonline.com" + assert url.path == f"/{TENANT}/oauth2/v2.0/authorize" + assert params["client_id"] == APP + assert params["scope"] == f"api://{APP}/user_impersonation offline_access openid profile" + assert params["redirect_uri"] == "http://localhost:8400" + assert params["response_mode"] == "query" + assert params["code_challenge"] == "chal" + assert params["code_challenge_method"] == "S256" + assert params["state"] == "st" + + +def test_exchange_code_builds_a_record() -> None: + body = { + "access_token": "access-1", + "refresh_token": "refresh-1", + "expires_in": 3600, + "id_token": _jwt({"preferred_username": "user@example.com"}), + } + with mock.patch.object(entra_user.requests, "post", return_value=_response(200, body)) as post: + record = entra_user.exchange_code(TENANT, APP, "the-code", "the-verifier") + sent = post.call_args.kwargs["data"] + assert post.call_args.args[0] == f"https://login.microsoftonline.com/{TENANT}/oauth2/v2.0/token" + assert sent["grant_type"] == "authorization_code" + assert sent["client_id"] == APP + assert sent["code"] == "the-code" + assert sent["code_verifier"] == "the-verifier" + assert record["tenant"] == TENANT + assert record["app"] == APP + assert record["username"] == "user@example.com" + assert record["expires_at"] > time.time() + 3000 + + +@pytest.mark.parametrize( + ("returned", "expected"), [("refresh-2", "refresh-2"), (None, "refresh-1")] +) +def test_refresh_keeps_or_rotates_the_refresh_token(returned: str | None, expected: str) -> None: + body: dict[str, Any] = {"access_token": "access-2", "expires_in": 3600} + if returned: + body["refresh_token"] = returned + with mock.patch.object(entra_user.requests, "post", return_value=_response(200, body)) as post: + record = entra_user.refresh_tokens(_record()) + assert post.call_args.kwargs["data"]["grant_type"] == "refresh_token" + assert record["access"] == "access-2" + assert record["refresh"] == expected + assert record["username"] == "user@example.com" + + +def test_token_errors_name_the_reason_but_not_tokens_or_trace() -> None: + body = { + "error": "invalid_grant", + "error_description": "AADSTS70000: Provided grant is invalid.\r\nTrace ID: abc\r\n", + } + with ( + mock.patch.object(entra_user.requests, "post", return_value=_response(400, body)), + pytest.raises(entra_user.EntraAuthError) as exc, + ): + entra_user.refresh_tokens(_record()) + assert exc.value.code == "invalid_grant" + assert "AADSTS70000" in str(exc.value) + assert "Trace ID" not in str(exc.value) + assert "refresh-1" not in str(exc.value) + + +def test_valid_token_is_served_from_the_store() -> None: + entra_user.save_record(_record()) + with mock.patch.object(entra_user.requests, "post") as post: + assert entra_user.get_valid_token() == "access-1" + post.assert_not_called() + + +def test_token_near_expiry_is_refreshed_and_saved() -> None: + entra_user.save_record(_record(expires_at=time.time() + 60)) + body = {"access_token": "access-2", "refresh_token": "refresh-2", "expires_in": 3600} + with mock.patch.object(entra_user.requests, "post", return_value=_response(200, body)): + assert entra_user.get_valid_token() == "access-2" + saved = entra_user.read_record() + assert saved is not None + assert saved["refresh"] == "refresh-2" + + +@pytest.mark.parametrize("action", ["logout", "login"]) +def test_logout_and_login_wait_for_a_refresh_in_flight(action: str) -> None: + entra_user.save_record(_record(expires_at=time.time() + 60)) + started, release = threading.Event(), threading.Event() + body = {"access_token": "access-2", "refresh_token": "refresh-2", "expires_in": 3600} + + def _slow_post(*_args: Any, **_kwargs: Any) -> Any: + started.set() + assert release.wait(5) + return _response(200, body) + + with mock.patch.object(entra_user.requests, "post", side_effect=_slow_post): + refresher = threading.Thread(target=entra_user.get_valid_token) + refresher.start() + assert started.wait(5) + if action == "logout": + other = threading.Thread(target=entra_user.logout) + else: + other = threading.Thread( + target=entra_user.save_record, args=(_record(access="new-login"),) + ) + other.start() + other.join(0.2) + assert other.is_alive() # blocked until the refresh has saved + release.set() + refresher.join(5) + other.join(5) + + saved = entra_user.read_record() + if action == "logout": + assert saved is None + else: + assert saved is not None + assert saved["access"] == "new-login" + + +def test_not_signed_in_raises() -> None: + with pytest.raises(entra_user.EntraAuthError, match="strix auth login entra"): + entra_user.get_valid_token() + + +def test_store_is_private() -> None: + entra_user.save_record(_record()) + assert entra_user.AUTH_PATH.stat().st_mode & 0o777 == 0o600 + + +def test_configured_only_for_the_signed_in_tenant_and_app(monkeypatch: pytest.MonkeyPatch) -> None: + assert not entra_user.is_configured() + entra_user.save_record(_record()) + assert entra_user.is_configured() + monkeypatch.setenv("STRIX_ENTRA_TENANT_ID", TENANT) + monkeypatch.setenv("STRIX_ENTRA_APP_ID", APP) + monkeypatch.setattr(loader, "_cached", None) + assert entra_user.is_configured() + monkeypatch.setenv("STRIX_ENTRA_APP_ID", "33333333-3333-3333-3333-333333333333") + monkeypatch.setattr(loader, "_cached", None) + assert not entra_user.is_configured() + + +def test_logout_reports_whether_there_was_a_sign_in() -> None: + assert entra_user.logout() is False + entra_user.save_record(_record()) + assert entra_user.logout() is True + assert entra_user.read_record() is None