diff --git a/packages/notte-sdk/README.md b/packages/notte-sdk/README.md index 037dbc079..802f267c2 100644 --- a/packages/notte-sdk/README.md +++ b/packages/notte-sdk/README.md @@ -34,6 +34,9 @@ with notte.Session(open_viewer=True) as session: ) ``` +When `api_key` and `NOTTE_API_KEY` are both unset, the client also uses the +environment-specific API key saved by `notte auth login` in the system keyring. + ## Core Components ### Session Management diff --git a/packages/notte-sdk/pyproject.toml b/packages/notte-sdk/pyproject.toml index 0aa2a3b2e..fe3a66bf6 100644 --- a/packages/notte-sdk/pyproject.toml +++ b/packages/notte-sdk/pyproject.toml @@ -13,6 +13,7 @@ packages = [ requires-python = ">=3.11" dependencies = [ "halo>=0.0.28", + "keyring>=25.6.0", "notte-core==1.4.4.dev", "websockets>=13.1", ] diff --git a/packages/notte-sdk/src/notte_sdk/auth.py b/packages/notte-sdk/src/notte_sdk/auth.py new file mode 100644 index 000000000..5f1c40e40 --- /dev/null +++ b/packages/notte-sdk/src/notte_sdk/auth.py @@ -0,0 +1,90 @@ +import base64 +import json +import os +import sys +from typing import Any +from urllib.parse import urlparse + +from notte_core.common.logging import logger + +KEYRING_SERVICE = "notte-cli" +KEYRING_KEY = "api_key" + +_HOST_TO_ENV_LABEL = { + "api.notte.cc": "prod", + "us-prod.notte.cc": "prod", + "us-staging.notte.cc": "staging", + "us-dev.notte.cc": "dev", + "us-dev-test.notte.cc": "dev", +} + + +def resolve_env_label(api_url: str) -> str: + """Return the notte-cli keyring environment label for an API URL.""" + hostname = urlparse(api_url).hostname + if hostname is None: + return "prod" + return _HOST_TO_ENV_LABEL.get(hostname, hostname) + + +def _decode_cli_secret(secret: bytes) -> str | None: + """Decode an item written by github.com/99designs/keyring.""" + try: + item: dict[str, Any] = json.loads(secret) + data = item.get("Data") + if not isinstance(data, str): + return None + return base64.b64decode(data).decode() + except (UnicodeDecodeError, ValueError, TypeError, json.JSONDecodeError): + return None + + +def _get_from_secret_service(key: str) -> str | None: + """Read the custom Secret Service collection used by notte-cli on Linux.""" + if sys.platform != "linux": + return None + + try: + import secretstorage + + bus = secretstorage.dbus_init() + for collection in secretstorage.get_all_collections(bus): + if collection.get_label() != KEYRING_SERVICE: + continue + for item in collection.search_items({"profile": key}): + value = _decode_cli_secret(item.get_secret()) + if value: + return value + except Exception as exc: + # Keyring access is a best-effort fallback. Headless Linux environments + # commonly have no Secret Service session available. + logger.debug(f"Could not read notte-cli Secret Service keyring: {exc}") + return None + + +def _get_from_system_keyring(key: str) -> str | None: + """Read backends whose notte-cli representation matches Python keyring.""" + try: + import keyring + + return keyring.get_password(KEYRING_SERVICE, key) + except Exception as exc: + logger.debug(f"Could not read notte-cli system keyring: {exc}") + return None + + +def get_keyring_api_key(api_url: str) -> str | None: + """Load the API key stored by notte-cli for the selected API environment.""" + env_key = f"{KEYRING_KEY}:{resolve_env_label(api_url)}" + for key in (env_key, KEYRING_KEY if env_key == f"{KEYRING_KEY}:prod" else None): + if key is None: + continue + value = _get_from_secret_service(key) or _get_from_system_keyring(key) + if value: + return value + return None + + +def resolve_api_key(api_key: str | None, server_url: str) -> str | None: + """Resolve an API key from code, environment, then the notte-cli keyring.""" + return api_key or os.getenv("NOTTE_API_KEY") or get_keyring_api_key(server_url) diff --git a/packages/notte-sdk/src/notte_sdk/endpoints/base.py b/packages/notte-sdk/src/notte_sdk/endpoints/base.py index 06255ee9a..3c67a5fa6 100644 --- a/packages/notte-sdk/src/notte_sdk/endpoints/base.py +++ b/packages/notte-sdk/src/notte_sdk/endpoints/base.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ValidationError from requests.exceptions import ConnectionError +from notte_sdk.auth import resolve_api_key from notte_sdk.errors import AuthenticationError, NotteAPIError, NotteAPIExecutionError if TYPE_CHECKING: @@ -100,26 +101,31 @@ def __init__( """ Initialize a new API client instance. - Sets up the client by resolving an API key from the provided parameter or the - NOTTE_API_KEY environment variable. Selects the server URL (defaulting to a - preconfigured server if none is provided), initializes a mapping of endpoints - using the implemented 'endpoints' method, and stores an optional base endpoint - path for constructing request URLs. + Sets up the client by resolving an API key from the provided parameter, the + NOTTE_API_KEY environment variable, or the notte-cli keyring. Selects the server + URL (defaulting to a preconfigured server if none is provided), initializes a + mapping of endpoints using the implemented 'endpoints' method, and stores an + optional base endpoint path for constructing request URLs. Args: base_endpoint_path: Optional base path to be prefixed to endpoint URLs. api_key: Optional API key for authentication; if not supplied, retrieved from - the NOTTE_API_KEY environment variable. + the NOTTE_API_KEY environment variable or the notte-cli keyring. Raises: - AuthenticationError: If an API key is neither provided nor available in the environment. + AuthenticationError: If an API key cannot be resolved. """ self.root_client = root_client # pyright: ignore [reportUnannotatedClassAttribute] - token = api_key or os.getenv("NOTTE_API_KEY") + self.server_url: str = server_url or os.getenv("NOTTE_API_URL") or self.DEFAULT_NOTTE_API_URL + token = api_key or os.getenv("NOTTE_API_KEY") or getattr(root_client, "_resolved_api_key", None) + if token is None: + token = resolve_api_key(api_key, self.server_url) if token is None: - raise AuthenticationError("NOTTE_API_KEY needs to be provided") + raise AuthenticationError( + "No API key found. Provide api_key, set NOTTE_API_KEY, or run 'notte auth login'." + ) + setattr(root_client, "_resolved_api_key", token) self.token: str = token - self.server_url: str = server_url or os.getenv("NOTTE_API_URL") or self.DEFAULT_NOTTE_API_URL self.base_endpoint_path: str | None = base_endpoint_path self.verbose: bool = verbose self.db_preview = (os.getenv("NOTTE_DB_PREVIEW_BRANCH") or "").strip() or None diff --git a/tests/sdk/test_client.py b/tests/sdk/test_client.py index 9f78f43fd..e6abc6964 100644 --- a/tests/sdk/test_client.py +++ b/tests/sdk/test_client.py @@ -64,7 +64,7 @@ def test_client_initialization_with_params() -> None: def test_client_initialization_without_api_key() -> None: - with patch.dict(os.environ, clear=True): + with patch.dict(os.environ, clear=True), patch("notte_sdk.endpoints.base.resolve_api_key", return_value=None): with pytest.raises(AuthenticationError): _ = NotteClient() diff --git a/tests/sdk/test_sdk_auth.py b/tests/sdk/test_sdk_auth.py new file mode 100644 index 000000000..658159755 --- /dev/null +++ b/tests/sdk/test_sdk_auth.py @@ -0,0 +1,62 @@ +import base64 +import json +import os +from unittest.mock import patch + +from notte_sdk import NotteClient +from notte_sdk.auth import _decode_cli_secret, get_keyring_api_key, resolve_api_key, resolve_env_label + + +def test_resolve_env_label_matches_cli() -> None: + assert resolve_env_label("https://api.notte.cc") == "prod" + assert resolve_env_label("https://us-staging.notte.cc") == "staging" + assert resolve_env_label("http://localhost:8000") == "localhost" + assert resolve_env_label("") == "prod" + + +def test_decode_cli_secret() -> None: + encoded_key = base64.b64encode(b"test-key").decode() + secret = json.dumps({"Key": "api_key:prod", "Data": encoded_key}).encode() + + assert _decode_cli_secret(secret) == "test-key" + assert _decode_cli_secret(b"not-json") is None + + +def test_get_keyring_api_key_uses_environment_key() -> None: + with ( + patch("notte_sdk.auth._get_from_secret_service", side_effect=lambda key: f"value-for-{key}"), + patch("notte_sdk.auth._get_from_system_keyring", return_value=None), + ): + assert get_keyring_api_key("https://us-dev.notte.cc") == "value-for-api_key:dev" + + +def test_get_keyring_api_key_supports_legacy_prod_key() -> None: + values = {"api_key": "legacy-key"} # pragma: allowlist secret + with ( + patch("notte_sdk.auth._get_from_secret_service", side_effect=lambda key: values.get(key)), + patch("notte_sdk.auth._get_from_system_keyring", return_value=None), + ): + assert get_keyring_api_key("https://api.notte.cc") == "legacy-key" + + +def test_resolve_api_key_precedence(monkeypatch) -> None: + monkeypatch.setenv("NOTTE_API_KEY", "env-key") + with patch("notte_sdk.auth.get_keyring_api_key", return_value="keyring-key") as get_keyring: + assert resolve_api_key("code-key", "https://api.notte.cc") == "code-key" + assert resolve_api_key(None, "https://api.notte.cc") == "env-key" + get_keyring.assert_not_called() + + monkeypatch.delenv("NOTTE_API_KEY") + with patch("notte_sdk.auth.get_keyring_api_key", return_value="keyring-key"): + assert resolve_api_key(None, "https://api.notte.cc") == "keyring-key" + + +def test_client_only_reads_keyring_once() -> None: + with ( + patch.dict(os.environ, {}, clear=True), + patch("notte_sdk.endpoints.base.resolve_api_key", return_value="keyring-key") as resolve, + ): + client = NotteClient() + + assert resolve.call_count == 1 + assert client.sessions.token == client.workflows.token == "keyring-key" diff --git a/uv.lock b/uv.lock index 16f2f8d58..5d64c9399 100644 --- a/uv.lock +++ b/uv.lock @@ -3055,6 +3055,7 @@ version = "1.4.4.dev0" source = { editable = "packages/notte-sdk" } dependencies = [ { name = "halo" }, + { name = "keyring" }, { name = "notte-core" }, { name = "websockets" }, ] @@ -3067,6 +3068,7 @@ playwright = [ [package.metadata] requires-dist = [ { name = "halo", specifier = ">=0.0.28" }, + { name = "keyring", specifier = ">=25.6.0" }, { name = "notte-core", editable = "packages/notte-core" }, { name = "playwright", marker = "extra == 'playwright'", specifier = "~=1.55" }, { name = "websockets", specifier = ">=13.1" },