diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index a9c7c82ae..ff663c404 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -13,7 +13,7 @@ import asyncio import contextlib import json import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -29,6 +29,32 @@ _LOGIN_AS_GUEST_BODY = ( ) +def _extract_access_token(payload: Any) -> str | None: + if not isinstance(payload, dict): + return None + data = payload.get("data") + if not isinstance(data, dict): + return None + login = data.get("loginAsGuest") + if not isinstance(login, dict): + return None + token_obj = login.get("token") + if not isinstance(token_obj, dict): + return None + token = token_obj.get("accessToken") + if isinstance(token, str) and token.strip(): + return token.strip() + return None + + +def _format_stderr(raw_stderr: Any, max_len: int = 200) -> str: + if isinstance(raw_stderr, (bytes, bytearray)): + return raw_stderr.decode("utf-8", errors="replace")[:max_len] + if raw_stderr is not None: + return str(raw_stderr)[:max_len] + return "" + + async def _login_as_guest( session: BaseSandboxSession, *, @@ -43,34 +69,36 @@ async def _login_as_guest( """ last_err: str | None = None for i in range(1, attempts + 1): - result = await session.exec( - "curl", - "-fsS", - "-X", - "POST", - "-H", - "Content-Type: application/json", - "-d", - _LOGIN_AS_GUEST_BODY, - f"{container_url}/graphql", - timeout=15, - ) + try: + result = await session.exec( + "curl", + "-fsS", + "-X", + "POST", + "-H", + "Content-Type: application/json", + "-d", + _LOGIN_AS_GUEST_BODY, + f"{container_url}/graphql", + timeout=15, + ) + except Exception as exc: # noqa: BLE001 + last_err = f"session.exec failed: {exc}" + logger.debug("loginAsGuest attempt %d/%d failed: %s", i, attempts, last_err) + await asyncio.sleep(min(2.0 * i, 8.0)) + continue + if result.ok(): try: payload = json.loads(result.stdout) - token = ( - payload.get("data", {}) - .get("loginAsGuest", {}) - .get("token", {}) - .get("accessToken") - ) + token = _extract_access_token(payload) if token: - return str(token) + return token last_err = f"loginAsGuest returned no token: {payload}" - except json.JSONDecodeError as exc: + except (json.JSONDecodeError, AttributeError, TypeError) as exc: last_err = f"unparseable response: {exc}: {result.stdout!r}" else: - stderr = result.stderr.decode("utf-8", errors="replace")[:200] + stderr = _format_stderr(result.stderr) last_err = f"curl exit {result.exit_code}: {stderr}" logger.debug("loginAsGuest attempt %d/%d failed: %s", i, attempts, last_err) await asyncio.sleep(min(2.0 * i, 8.0)) diff --git a/strix/tools/proxy/caido_api.py b/strix/tools/proxy/caido_api.py index af81b5008..2f3145020 100644 --- a/strix/tools/proxy/caido_api.py +++ b/strix/tools/proxy/caido_api.py @@ -6,6 +6,7 @@ import asyncio import json import os import time +import urllib.error import urllib.request from typing import TYPE_CHECKING, Any, Literal from urllib.parse import parse_qs, urlencode, urlparse, urlunparse @@ -75,9 +76,31 @@ def _login_as_guest() -> str: headers={"Content-Type": "application/json"}, method="POST", ) - with urllib.request.urlopen(req, timeout=10) as resp: # noqa: S310 # nosec B310 - payload = json.loads(resp.read()) - return str(payload["data"]["loginAsGuest"]["token"]["accessToken"]) + try: + with urllib.request.urlopen(req, timeout=10) as resp: # noqa: S310 # nosec B310 + raw_resp = resp.read() + except urllib.error.HTTPError as err: + raw_resp = err.read() + except urllib.error.URLError as err: + raise RuntimeError(f"Failed to connect to Caido at {_graphql_url()}: {err}") from err + + try: + payload = json.loads(raw_resp) + except (json.JSONDecodeError, TypeError) as err: + raise RuntimeError(f"Unparseable response from Caido: {raw_resp!r}") from err + + token = None + if isinstance(payload, dict): + data = payload.get("data") + if isinstance(data, dict): + login = data.get("loginAsGuest") + if isinstance(login, dict): + token_obj = login.get("token") + if isinstance(token_obj, dict): + token = token_obj.get("accessToken") + if not token or not isinstance(token, str) or not token.strip(): + raise RuntimeError(f"loginAsGuest returned no token: {payload}") + return token.strip() async def _new_client() -> Client: diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index 98e42a94d..10da9a1bd 100644 --- a/tests/test_caido_bootstrap.py +++ b/tests/test_caido_bootstrap.py @@ -8,24 +8,33 @@ client exists but no caller will ever see it to close it. from __future__ import annotations import asyncio +import io import sys import types +import urllib.error +import urllib.request from typing import Any +from unittest.mock import MagicMock import pytest -from strix.runtime.caido_bootstrap import bootstrap_caido +from strix.runtime.caido_bootstrap import _login_as_guest, bootstrap_caido +from strix.tools.proxy import caido_api class _FakeExecResult: - stderr = b"" - exit_code = 0 - - def __init__(self, stdout: str) -> None: + def __init__( + self, + stdout: str, + stderr: Any = b"", + exit_code: int = 0, + ) -> None: self.stdout = stdout + self.stderr = stderr + self.exit_code = exit_code def ok(self) -> bool: - return True + return self.exit_code == 0 class _FakeSession: @@ -79,3 +88,241 @@ async def test_cancellation_during_connect_closes_the_client( async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None: client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener")) assert client.closed + + +class _MultiResultSession: + def __init__(self, responses: list[str]) -> None: + self._responses = list(responses) + self.call_count = 0 + + async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult: + self.call_count += 1 + resp = self._responses.pop(0) if self._responses else "{}" + return _FakeExecResult(resp) + + +async def _instant_sleep(_s: float) -> None: + pass + + +async def test_login_as_guest_retries_when_data_is_null( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data": null, "errors": [{"message": "Could not acquire a connection to the database"}]}', + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_ok_token"}}}}', + ]) + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "retry_ok_token" # noqa: S105 + assert session.call_count == 2 + + +async def test_login_as_guest_retries_when_login_as_guest_is_null( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data": {"loginAsGuest": null}}', + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_ok_token"}}}}', + ]) + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "retry_ok_token" # noqa: S105 + assert session.call_count == 2 + + +async def test_login_as_guest_retries_when_token_is_null( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data": {"loginAsGuest": {"token": null}}}', + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_ok_token"}}}}', + ]) + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "retry_ok_token" # noqa: S105 + assert session.call_count == 2 + + +async def test_login_as_guest_retries_when_data_is_non_dict( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data": "unexpected_string"}', + '{"data": [1, 2, 3]}', + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_ok_token"}}}}', + ]) + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=4, + ) + assert token == "retry_ok_token" # noqa: S105 + assert session.call_count == 3 + + +async def test_login_as_guest_fails_when_all_attempts_return_null_data( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data": null, "errors": [{"message": "busy"}]}', + '{"data": null, "errors": [{"message": "busy"}]}', + ]) + with pytest.raises(RuntimeError, match="loginAsGuest failed after 2 attempts"): + await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=2, + ) + assert session.call_count == 2 + + +def test_caido_api_login_as_guest_null_data(monkeypatch: pytest.MonkeyPatch) -> None: + mock_resp = MagicMock() + mock_resp.read.return_value = b'{"data": null, "errors": [{"message": "busy"}]}' + mock_resp.__enter__.return_value = mock_resp + mock_resp.__exit__.return_value = None + + monkeypatch.setattr(urllib.request, "urlopen", lambda *_a, **_k: mock_resp) + + with pytest.raises(RuntimeError, match="loginAsGuest returned no token"): + caido_api._login_as_guest() + + +def test_caido_api_login_as_guest_malformed_data(monkeypatch: pytest.MonkeyPatch) -> None: + mock_resp = MagicMock() + mock_resp.read.return_value = b'{"data": "invalid_string_instead_of_dict"}' + mock_resp.__enter__.return_value = mock_resp + mock_resp.__exit__.return_value = None + + monkeypatch.setattr(urllib.request, "urlopen", lambda *_a, **_k: mock_resp) + + with pytest.raises(RuntimeError, match="loginAsGuest returned no token"): + caido_api._login_as_guest() + + +def test_caido_api_login_as_guest_success(monkeypatch: pytest.MonkeyPatch) -> None: + mock_resp = MagicMock() + mock_resp.read.return_value = ( + b'{"data": {"loginAsGuest": {"token": {"accessToken": "host_token_123"}}}}' + ) + mock_resp.__enter__.return_value = mock_resp + mock_resp.__exit__.return_value = None + + monkeypatch.setattr(urllib.request, "urlopen", lambda *_a, **_k: mock_resp) + + assert caido_api._login_as_guest() == "host_token_123" + + +async def test_login_as_guest_retries_when_session_exec_raises_exception( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + + class _ExceptionThenSuccessSession: + def __init__(self) -> None: + self.call_count = 0 + + async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult: + self.call_count += 1 + if self.call_count == 1: + raise TimeoutError("container exec timed out") + return _FakeExecResult( + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_after_exec_err"}}}}' + ) + + session = _ExceptionThenSuccessSession() + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "retry_after_exec_err" # noqa: S105 + assert session.call_count == 2 + + +async def test_login_as_guest_handles_none_stderr_on_curl_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + + class _NoneStderrThenSuccessSession: + def __init__(self) -> None: + self.call_count = 0 + + async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult: + self.call_count += 1 + if self.call_count == 1: + return _FakeExecResult("", stderr=None, exit_code=7) + return _FakeExecResult( + '{"data":{"loginAsGuest":{"token":{"accessToken":"retry_after_none_stderr"}}}}' + ) + + session = _NoneStderrThenSuccessSession() + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "retry_after_none_stderr" # noqa: S105 + assert session.call_count == 2 + + +async def test_login_as_guest_rejects_whitespace_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) + session = _MultiResultSession([ + '{"data":{"loginAsGuest":{"token":{"accessToken":" "}}}}', + '{"data":{"loginAsGuest":{"token":{"accessToken":"valid_token"}}}}', + ]) + token = await _login_as_guest( + session, # type: ignore[arg-type] + container_url="http://container", + attempts=3, + ) + assert token == "valid_token" # noqa: S105 + assert session.call_count == 2 + + +def test_caido_api_login_as_guest_http_500_with_graphql_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + err_body = b'{"errors": [{"message": "database locked"}]}' + http_err = urllib.error.HTTPError( + url="http://127.0.0.1:8080/graphql", + code=500, + msg="Internal Server Error", + hdrs={}, # type: ignore[arg-type] + fp=io.BytesIO(err_body), + ) + + monkeypatch.setattr(urllib.request, "urlopen", MagicMock(side_effect=http_err)) + + with pytest.raises(RuntimeError, match="database locked"): + caido_api._login_as_guest() + + +def test_caido_api_login_as_guest_url_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + url_err = urllib.error.URLError("Connection refused") + monkeypatch.setattr(urllib.request, "urlopen", MagicMock(side_effect=url_err)) + + with pytest.raises(RuntimeError, match="Failed to connect to Caido"): + caido_api._login_as_guest() +