diff --git a/README.md b/README.md index f72ea2eb..dff41c4d 100644 --- a/README.md +++ b/README.md @@ -277,8 +277,8 @@ For the Python SDK and CLI, run `lightning login` or export Inside a Lightning Studio, this value is usually already available. On a machine without a browser, such as a remote server, run -`lightning login --device`. It prints a code to enter at the address it shows, -from a browser on any other device. +`lightning login --no-browser`. It prints an address to approve the login at, +from a browser on any other device, then asks for the code that page shows. Sandbox-only API keys can also be passed through `LIGHTNING_SANDBOX_API_KEY` or configured directly in code with `Sandbox.configure(...)`. diff --git a/python/docs/cli/index.rst b/python/docs/cli/index.rst index d5d7d3f7..d1e5977f 100644 --- a/python/docs/cli/index.rst +++ b/python/docs/cli/index.rst @@ -31,12 +31,12 @@ For interactive use, sign in with: lightning login -On a machine without a browser, such as a remote server, sign in with a code -that you enter from a browser on any other device: +On a machine without a browser, such as a remote server, approve the login +from a browser on any other device, then paste the code it shows: .. code-block:: bash - lightning login --device + lightning login --no-browser For non-interactive environments, configure credentials through environment variables instead: diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 501109a4..92c03848 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -44,7 +44,7 @@ ) from lightning_sdk.cli.ls import ls from lightning_sdk.cli.rm import rm -from lightning_sdk.cli.utils.auth import browser_authentication, login_with_device_code, logout_session +from lightning_sdk.cli.utils.auth import browser_authentication, login_with_pasted_code, logout_session from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.utils.resolve import _get_authed_user, in_studio @@ -96,11 +96,11 @@ def main_cli(ctx: click.Context) -> None: @main_cli.command @click.option( - "--device", + "--no-browser", is_flag=True, - help="Sign in from a machine without a browser, such as a remote server, by approving a code on another device.", + help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) -def login(device: bool) -> None: +def login(no_browser: bool) -> None: """Sign in to Lightning AI.""" auth = Auth() if auth.api_key or auth.auth_token: @@ -114,8 +114,8 @@ def login(device: bool) -> None: # always start a fresh session: a saved one may have expired auth.clear() try: - if device: - login_with_device_code(auth) + if no_browser: + login_with_pasted_code(auth) else: with browser_authentication(True): auth.authenticate() diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 3ecc6f64..64b76be4 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,14 +1,15 @@ """Non-interactive authentication helpers for CLI commands.""" import base64 +import hashlib import hmac import json import os -import time +import secrets from contextlib import contextmanager, suppress from contextvars import ContextVar from pathlib import Path -from typing import Generator, Optional +from typing import Generator, Optional, Tuple from urllib.parse import urlencode import requests @@ -20,14 +21,67 @@ _BROWSER_AUTH_ALLOWED = ContextVar("lightning_browser_auth_allowed", default=True) _ORIGINAL_AUTH_SAVE = Auth.save -_DEVICE_CLIENT_ID = "lightning-cli" -_DEVICE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code" +_CLIENT_ID = "lightning-cli" +_PASTE_ATTEMPTS = 3 def _generate_state() -> str: return base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("ascii") +def _origin() -> str: + """Platform origin without a trailing slash, so paths are not `//cli-login`.""" + return env.LIGHTNING_CLOUD_URL.rstrip("/") + + +def _pkce_pair() -> Tuple[str, str]: + """Return a PKCE `(verifier, S256 challenge)` pair.""" + verifier = secrets.token_urlsafe(64) + challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=") + return verifier, challenge.decode("ascii") + + +def _authorize_url(redirect_uri: str, state: str, challenge: str) -> str: + query = urlencode( + { + "response_type": "code", + "client_id": _CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": "openid", + "state": state, + "code_challenge": challenge, + "code_challenge_method": "S256", + } + ) + return f"{_origin()}/oidc/authorize?{query}" + + +class _CodeRejectedError(Exception): + """The token endpoint refused the code; `str(e)` is its OAuth error.""" + + +def _exchange_code(code: str, verifier: str, redirect_uri: str) -> str: + """Trade an approved code and its PKCE verifier for a CLI session token.""" + try: + response = requests.post( + f"{_origin()}/oidc/token", + data={ + "grant_type": "authorization_code", + "client_id": _CLIENT_ID, + "code": code, + "code_verifier": verifier, + "redirect_uri": redirect_uri, + }, + timeout=10, + ) + except requests.ConnectionError as e: + raise ConnectionError(str(e)) from e + body = _json_body(response) + if response.ok and body.get("access_token"): + return body["access_token"] + raise _CodeRejectedError(body.get("error") or f"HTTP {response.status_code}") + + def _save_with_session_token( self: Auth, token: str = "", @@ -52,26 +106,22 @@ def _save_with_session_token( def _login_with_browser(self: AuthServer, auth: Auth) -> None: - """Complete browser authentication without forwarding credentials back to `/me/apps`.""" + """Sign in with an authorization code and PKCE, delivered to a callback on 127.0.0.1.""" import webbrowser import uvicorn - from fastapi import FastAPI, Query + from fastapi import FastAPI from starlette.background import BackgroundTask from starlette.responses import RedirectResponse from lightning_sdk.lightning_cloud.utils.network import find_free_network_port expected_state = _generate_state() + verifier, challenge = _pkce_pair() app = FastAPI() port = find_free_network_port() - base_url = self.get_auth_url(port) - redirect_uri = f"http://localhost:{port}/login-complete" - if urlencode({"redirectTo": redirect_uri}) in base_url: - stateful_redirect = f"{redirect_uri}?{urlencode({'state': expected_state})}" - url = base_url.replace(urlencode({"redirectTo": redirect_uri}), urlencode({"redirectTo": stateful_redirect}), 1) - else: - url = base_url + redirect_uri = f"http://127.0.0.1:{port}/callback" + url = _authorize_url(redirect_uri, expected_state, challenge) try: requests.head(url, timeout=10) @@ -86,36 +136,34 @@ def _login_with_browser(self: AuthServer, auth: Auth) -> None: logger.info(f"login started for lightning.ai, opening {url}") if not webbrowser.open(url): - deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID") - if deployment_id: + if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): raise RuntimeError( "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " "'Environment' section of the deployment settings." ) - raise RuntimeError( - "Failed to authenticate to Lightning: no browser could be opened on this machine. " - "Run `lightning login --device` to sign in from another device." - ) + click.echo("No browser could be opened on this machine.") + login_with_pasted_code(auth) + return + + @app.get("/callback") + def save_token(code: str = "", state: str = "", error: str = "") -> RedirectResponse: + failure = "" + if not state or not hmac.compare_digest(state, expected_state): + failure = "state mismatch" + elif error or not code: + failure = error or "no code in the callback" + else: + try: + auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) + except (_CodeRejectedError, ConnectionError) as e: + failure = str(e) - @app.get("/login-complete") - async def save_token( - token: str = "", - key: str = "", - state: str = "", - user_id: str = Query("", alias="userID"), - ) -> RedirectResponse: - state_valid = bool(state) and hmac.compare_digest(state, expected_state) - if token and state_valid: - auth.save(token=token, user_id=user_id, api_key=key) - logger.info("Authentication Successful") - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-successful" + if failure: + logger.warning(f"Authentication Failed ({failure}). Run `lightning login` again, or update the CLI.") + target = f"{_origin()}/cli-login-failed" else: - reason = "no token in the callback" if not token else "state mismatch" - logger.warning( - f"Authentication Failed ({reason}). If the browser was already signed in, the sign-in page " - "could not mint a CLI token. Run `lightning login` again, or update the CLI." - ) - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-failed" + logger.info("Authentication Successful") + target = f"{_origin()}/cli-login-successful" return RedirectResponse( url=target, @@ -125,7 +173,7 @@ async def save_token( def stop_server() -> None: server.should_exit = True - server = uvicorn.Server(config=uvicorn.Config(app, port=port, log_level="error")) + server = uvicorn.Server(config=uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) server.run() @@ -133,63 +181,27 @@ def stop_server() -> None: AuthServer.login_with_browser = _login_with_browser -def login_with_device_code(auth: Auth) -> None: - """Sign in by approving a code in a browser on any other device (RFC 8628 device authorization).""" - try: - response = requests.post( - f"{env.LIGHTNING_CLOUD_URL}/oidc/device_authorization", - data={"client_id": _DEVICE_CLIENT_ID}, - timeout=10, - ) - except requests.ConnectionError as e: - raise ConnectionError(str(e)) from e - grant = _json_body(response) - if "device_code" not in grant: - reason = grant.get("error") or f"HTTP {response.status_code}" - raise RuntimeError( - f"{env.LIGHTNING_CLOUD_URL} did not start a device login ({reason}). " - "It may not support `lightning login --device` yet." - ) - +def login_with_pasted_code(auth: Auth) -> None: + """Sign in on a machine without a browser: approve on any device, then paste the code shown there.""" + redirect_uri = f"{_origin()}/cli-login" + verifier, challenge = _pkce_pair() click.echo( - f"To sign in, open {grant['verification_uri']} on any device and enter this code:\n\n" - f" {grant['user_code']}\n\n" - f"Waiting for approval. The code expires in {grant['expires_in'] // 60} minutes." + "To sign in, open this address in a browser on any device and approve the request:\n\n" + f" {_authorize_url(redirect_uri, _generate_state(), challenge)}\n\n" + "The page then shows a code. Only paste a code from a page you opened this way; " + "never give it to anyone, as it gives full access to your Lightning account." ) - - interval = grant.get("interval", 5) - deadline = time.monotonic() + grant["expires_in"] - while time.monotonic() < deadline: - time.sleep(interval) + for attempt in range(_PASTE_ATTEMPTS): + code = click.prompt("Code").strip() try: - response = requests.post( - f"{env.LIGHTNING_CLOUD_URL}/oidc/token", - data={ - "grant_type": _DEVICE_GRANT_TYPE, - "device_code": grant["device_code"], - "client_id": _DEVICE_CLIENT_ID, - }, - timeout=10, - ) - except requests.RequestException: - # the code stays valid until it expires, so a lost poll is retried - continue - body = _json_body(response) - if response.ok and body.get("access_token"): - auth.save(auth_token=body["access_token"]) + auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) return - error = body.get("error") - if error == "authorization_pending": - continue - if error == "slow_down": - interval += 5 - continue - if error == "access_denied": - raise RuntimeError("Login failed: the request was denied in the browser.") - if error == "expired_token": - break - raise RuntimeError(f"Login failed: {error or f'HTTP {response.status_code}'}.") - raise RuntimeError("Login failed: the code expired before it was approved. Run `lightning login --device` again.") + except _CodeRejectedError as e: + if str(e) != "invalid_grant": + raise RuntimeError(f"Login failed: {e}.") from None + if attempt < _PASTE_ATTEMPTS - 1: + click.echo("That code was not accepted. It works once and expires after a minute. Try again.") + raise RuntimeError("Login failed: the code was not accepted. Run `lightning login --no-browser` again.") def _json_body(response: requests.Response) -> dict: @@ -217,7 +229,7 @@ def logout_session(auth: Auth) -> None: if isinstance(token, str) and token: with suppress(requests.RequestException): requests.post( - f"{env.LIGHTNING_CLOUD_URL}/v1/auth/logout", + f"{_origin()}/v1/auth/logout", headers={"Authorization": f"Bearer {token}"}, timeout=10, ) diff --git a/python/tests/cli/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index 3153dc01..8dc174cb 100644 --- a/python/tests/cli/test_entrypoint.py +++ b/python/tests/cli/test_entrypoint.py @@ -84,7 +84,7 @@ def _save_token(auth: Auth) -> None: assert "Login failed" in str(failed.exception) -def test_login_device_starts_a_fresh_device_login_instead_of_the_browser(monkeypatch, tmp_path) -> None: +def test_login_no_browser_starts_a_fresh_paste_login_instead_of_the_browser(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) Auth().save(auth_token="sess_expired") user = MagicMock() @@ -100,9 +100,9 @@ def _approve(auth: Auth) -> None: assert not Auth.secrets_file.exists() auth.save(auth_token="sess_new") - monkeypatch.setattr("lightning_sdk.cli.entrypoint.login_with_device_code", _approve) + monkeypatch.setattr("lightning_sdk.cli.entrypoint.login_with_pasted_code", _approve) - result = CliRunner().invoke(main_cli, ["login", "--device"]) + result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) assert result.exit_code == 0, result.output assert 'Successfully logged in as "Test User".' in result.output @@ -297,7 +297,7 @@ def test_login_help(): result_text = assert_help_contains("lightning login --help", "Usage: lightning login [OPTIONS]") assert "Sign in to Lightning AI." in result_text - assert "--device" in result_text + assert "--no-browser" in result_text @mock_command_logging diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index 3adbde3a..bc5488d1 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -1,5 +1,8 @@ +import base64 +import hashlib import json from unittest.mock import MagicMock, PropertyMock, patch +from urllib.parse import parse_qs, urlparse import pytest import rich_click as click @@ -7,39 +10,52 @@ from lightning_sdk.cli.legacy.deploy._auth import _AuthLitServe, _AuthMode from lightning_sdk.cli.utils.auth import ( browser_authentication, - login_with_device_code, + login_with_pasted_code, logout_session, require_auth_header, ) from lightning_sdk.lightning_cloud.login import Auth -_DEVICE_GRANT = { - "device_code": "device-code", - "user_code": "BCDF-GHJK", - "verification_uri": "https://lightning.ai/device", - "expires_in": 600, - "interval": 5, -} - -def _response(status_code: int, body=None, json_error=None) -> MagicMock: +def _response(status_code: int, body=None) -> MagicMock: response = MagicMock(status_code=status_code, ok=status_code < 400) - response.json.side_effect = json_error response.json.return_value = body return response -def _device_login_env(monkeypatch, tmp_path, responses: list) -> tuple: - """Point credentials at tmp_path, answer HTTP calls from `responses`, and record sleeps instead of sleeping.""" +def _login_env(monkeypatch, tmp_path): + """Point credentials at tmp_path and return the secrets file.""" secrets_file = tmp_path / "credentials.json" monkeypatch.setattr(Auth, "secrets_file", secrets_file) - for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): + for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN", "LIGHTNING_DEPLOYMENT_ID"): monkeypatch.delenv(key, raising=False) - post = MagicMock(side_effect=responses) - sleeps: list = [] + return secrets_file + + +def _challenge(verifier: str) -> str: + return base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() + + +def _capture_browser_login(monkeypatch, tmp_path, token_response) -> dict: + """Run the loopback login with a fake server and browser; return what was captured.""" + captured: dict = {"secrets_file": _login_env(monkeypatch, tmp_path)} + + class FakeServer: + def __init__(self, config) -> None: + captured["config"] = config + captured["route"] = next(r for r in config.app.routes if getattr(r, "path", None) == "/callback") + + def run(self) -> None: + captured["run"](captured["route"].endpoint) + + post = MagicMock(side_effect=token_response) + monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.time.sleep", sleeps.append) - return secrets_file, post, sleeps + monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) and True) + monkeypatch.setattr("uvicorn.Server", FakeServer) + captured["post"] = post + return captured def test_require_auth_header_uses_environment_credentials() -> None: @@ -159,158 +175,182 @@ def test_logout_session_revokes_server_session_and_clears_credentials(monkeypatc assert not secrets_file.exists() -def test_login_with_browser_redirects_without_credentials_in_url(monkeypatch, tmp_path) -> None: - """The `/login-complete` callback redirects to `/cli-login-successful` without forwarding tokens in the URL.""" - import asyncio - +def test_login_with_browser_exchanges_the_code_with_pkce_on_loopback(monkeypatch, tmp_path) -> None: + """The callback trades the code and the verifier matching the sent challenge for a session.""" from lightning_sdk.lightning_cloud.login import AuthServer - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + captured = _capture_browser_login( + monkeypatch, tmp_path, [_response(200, {"access_token": "sess_user-1.secret", "token_type": "Bearer"})] + ) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="the-code", state="fixed-state")) - captured: dict = {} + auth = Auth() + AuthServer().login_with_browser(auth) - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app + url = urlparse(captured["auth_url"]) + query = {k: v[0] for k, v in parse_qs(url.query).items()} + assert url.path == "/oidc/authorize" + assert query["client_id"] == "lightning-cli" + assert query["response_type"] == "code" + assert query["scope"] == "openid" + assert query["state"] == "fixed-state" + assert query["code_challenge_method"] == "S256" + assert query["redirect_uri"] == f"http://127.0.0.1:{captured['config'].port}/callback" + assert captured["config"].host == "127.0.0.1" + + data = captured["post"].call_args.kwargs["data"] + assert captured["post"].call_args.args[0].endswith("/oidc/token") + assert data["grant_type"] == "authorization_code" + assert data["code"] == "the-code" + assert data["redirect_uri"] == query["redirect_uri"] + assert _challenge(data["code_verifier"]) == query["code_challenge"] + + assert captured["location"].headers["location"].endswith("/cli-login-successful") + assert auth.auth_token == "sess_user-1.secret" + assert json.loads(captured["secrets_file"].read_text())["auth_token"] == "sess_user-1.secret" + assert (captured["secrets_file"].stat().st_mode & 0o777) == 0o600 - def run(self) -> None: - route = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - resp = asyncio.run( - route.endpoint(token="sess_user-1.secret", key="", state="fixed-state", user_id="user-1") - ) - captured["location"] = resp.headers["location"] - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) or True) - monkeypatch.setattr("uvicorn.Server", FakeServer) +@pytest.mark.parametrize( + "params", + [ + {"code": "the-code", "state": ""}, + {"code": "the-code", "state": "forged-state"}, + {"code": "", "state": "fixed-state", "error": "access_denied"}, + {"code": "", "state": "fixed-state"}, + ], + ids=["no state", "forged state", "denied", "no code"], +) +def test_login_with_browser_refuses_a_bad_callback_without_exchanging(monkeypatch, tmp_path, params) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + captured = _capture_browser_login(monkeypatch, tmp_path, []) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(**params)) auth = Auth() AuthServer().login_with_browser(auth) - assert "state=fixed-state" in captured["auth_url"] or "state%3Dfixed-state" in captured["auth_url"] - assert captured["location"].endswith("/cli-login-successful") - assert "token=" not in captured["location"] - assert auth.auth_token == "sess_user-1.secret" + assert captured["location"].headers["location"].endswith("/cli-login-failed") + captured["post"].assert_not_called() + assert not captured["secrets_file"].exists() + assert not auth.auth_token + + +def test_login_with_browser_reports_a_refused_exchange(monkeypatch, tmp_path) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + captured = _capture_browser_login(monkeypatch, tmp_path, [_response(400, {"error": "invalid_grant"})]) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="used", state="fixed-state")) + AuthServer().login_with_browser(Auth()) -def test_login_with_browser_rejects_callback_without_matching_state(monkeypatch, tmp_path) -> None: - """Omitting or forging `state` must not store the callback token or key.""" - import asyncio + assert captured["location"].headers["location"].endswith("/cli-login-failed") + assert not captured["secrets_file"].exists() + +def test_login_with_browser_without_a_browser_falls_back_to_paste(monkeypatch, tmp_path) -> None: from lightning_sdk.lightning_cloud.login import AuthServer - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + secrets_file = _login_env(monkeypatch, tmp_path) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("webbrowser.open", lambda url: False) + monkeypatch.setattr("uvicorn.Server", lambda config: (_ for _ in ()).throw(AssertionError("server started"))) + monkeypatch.setattr( + "lightning_sdk.cli.utils.auth.login_with_pasted_code", lambda auth: auth.save(auth_token="sess_x") + ) - captured: dict = {} + AuthServer().login_with_browser(Auth()) - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app + assert json.loads(secrets_file.read_text())["auth_token"] == "sess_x" - def run(self) -> None: - captured["route"] = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") +def test_login_with_browser_in_a_deployment_without_a_browser_raises(monkeypatch, tmp_path) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + _login_env(monkeypatch, tmp_path) + monkeypatch.setenv("LIGHTNING_DEPLOYMENT_ID", "dep") monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: True) - monkeypatch.setattr("uvicorn.Server", FakeServer) + monkeypatch.setattr("webbrowser.open", lambda url: False) - auth = Auth() - AuthServer().login_with_browser(auth) + with pytest.raises(RuntimeError, match="Include SDK credentials"): + AuthServer().login_with_browser(Auth()) - for state in ("", "forged-state"): - resp = asyncio.run( - captured["route"].endpoint( - token="sess_attacker.secret", key="attacker-key", state=state, user_id="attacker" - ) - ) - assert resp.headers["location"].endswith("/cli-login-failed") - assert not secrets_file.exists() - assert not auth.auth_token - assert not auth.api_key +def _paste_login(monkeypatch, tmp_path, codes: list, responses: list) -> tuple: + secrets_file = _login_env(monkeypatch, tmp_path) + post = MagicMock(side_effect=responses) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.click.prompt", MagicMock(side_effect=codes)) + return secrets_file, post -def test_login_with_device_code_saves_the_approved_session(monkeypatch, tmp_path, capsys) -> None: - """The CLI polls at the server's interval, backs off on `slow_down`, and saves the session it is given.""" - secrets_file, post, sleeps = _device_login_env( - monkeypatch, - tmp_path, - [ - _response(200, _DEVICE_GRANT), - _response(400, {"error": "authorization_pending"}), - _response(400, {"error": "slow_down"}), - _response(200, {"access_token": "sess_user-1.secret", "token_type": "Bearer", "expires_in": 604800}), - ], +def test_login_with_pasted_code_strips_a_trailing_slash_from_the_origin(monkeypatch, tmp_path, capsys) -> None: + monkeypatch.setattr("lightning_sdk.lightning_cloud.env.LIGHTNING_CLOUD_URL", "http://localhost:9800/") + _paste_login(monkeypatch, tmp_path, ["the-code"], [_response(200, {"access_token": "sess_user-1.secret"})]) + + login_with_pasted_code(Auth()) + + url = next(word for word in capsys.readouterr().out.split() if "/oidc/authorize" in word) + query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} + assert url.startswith("http://localhost:9800/oidc/authorize?") + assert query["redirect_uri"] == "http://localhost:9800/cli-login" + assert query["scope"] == "openid" + + +def test_login_with_pasted_code_prints_the_authorize_url_and_saves_the_session(monkeypatch, tmp_path, capsys) -> None: + secrets_file, post = _paste_login( + monkeypatch, tmp_path, [" the-code \n"], [_response(200, {"access_token": "sess_user-1.secret"})] ) - login_with_device_code(Auth()) - - assert sleeps == [5, 5, 10] - assert post.call_args_list[0].args[0].endswith("/oidc/device_authorization") - assert post.call_args_list[0].kwargs["data"] == {"client_id": "lightning-cli"} - for poll in post.call_args_list[1:]: - assert poll.args[0].endswith("/oidc/token") - assert poll.kwargs["data"] == { - "grant_type": "urn:ietf:params:oauth:grant-type:device_code", - "device_code": "device-code", - "client_id": "lightning-cli", - } + login_with_pasted_code(Auth()) + output = capsys.readouterr().out - assert "https://lightning.ai/device" in output - assert "BCDF-GHJK" in output + url = next(word for word in output.split() if "/oidc/authorize" in word) + query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} + assert query["client_id"] == "lightning-cli" + assert query["scope"] == "openid" + assert query["redirect_uri"].endswith("/cli-login") + assert "//cli-login" not in query["redirect_uri"] + assert query["code_challenge_method"] == "S256" + assert "never give it to anyone" in output + + data = post.call_args.kwargs["data"] + assert data["code"] == "the-code" + assert data["redirect_uri"] == query["redirect_uri"] + assert _challenge(data["code_verifier"]) == query["code_challenge"] assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" - assert (secrets_file.stat().st_mode & 0o777) == 0o600 -@pytest.mark.parametrize( - ("poll", "message"), - [ - (_response(400, {"error": "access_denied"}), "denied in the browser"), - (_response(400, {"error": "expired_token"}), "code expired"), - (_response(400, {"error": "invalid_grant"}), "invalid_grant"), - ], -) -def test_login_with_device_code_stops_when_the_server_refuses(monkeypatch, tmp_path, poll, message) -> None: - secrets_file, _, _ = _device_login_env(monkeypatch, tmp_path, [_response(200, _DEVICE_GRANT), poll]) +def test_login_with_pasted_code_prompts_again_on_invalid_grant(monkeypatch, tmp_path) -> None: + secrets_file, post = _paste_login( + monkeypatch, + tmp_path, + ["typo", "the-code"], + [_response(400, {"error": "invalid_grant"}), _response(200, {"access_token": "sess_user-1.secret"})], + ) - with pytest.raises(RuntimeError, match=message): - login_with_device_code(Auth()) + login_with_pasted_code(Auth()) - assert not secrets_file.exists() + assert [c.kwargs["data"]["code"] for c in post.call_args_list] == ["typo", "the-code"] + assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" -def test_login_with_device_code_gives_up_when_the_code_expires(monkeypatch, tmp_path) -> None: - """Polling stops at `expires_in` even if the server never says the code expired.""" - secrets_file, post, _ = _device_login_env( - monkeypatch, tmp_path, [_response(200, _DEVICE_GRANT), _response(400, {"error": "authorization_pending"})] - ) - clock = iter([0, 300]) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.time.monotonic", lambda: next(clock, 600)) +def test_login_with_pasted_code_gives_up_after_three_bad_codes(monkeypatch, tmp_path) -> None: + refused = [_response(400, {"error": "invalid_grant"})] * 3 + secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a", "b", "c"], refused) - with pytest.raises(RuntimeError, match="code expired"): - login_with_device_code(Auth()) + with pytest.raises(RuntimeError, match="not accepted"): + login_with_pasted_code(Auth()) - assert post.call_count == 2 + assert post.call_count == 3 assert not secrets_file.exists() -@pytest.mark.parametrize( - "start", - [_response(404, {"code": 5}), _response(200, json_error=ValueError("an HTML page"))], - ids=["not found", "web app page"], -) -def test_login_with_device_code_explains_a_server_without_it(monkeypatch, tmp_path, start) -> None: - _device_login_env(monkeypatch, tmp_path, [start]) +def test_login_with_pasted_code_stops_on_any_other_error(monkeypatch, tmp_path) -> None: + secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a"], [_response(404, {})]) - with pytest.raises(RuntimeError, match="may not support `lightning login --device` yet"): - login_with_device_code(Auth()) + with pytest.raises(RuntimeError, match="HTTP 404"): + login_with_pasted_code(Auth()) + + assert post.call_count == 1 + assert not secrets_file.exists()