diff --git a/python/lightning_sdk/api/ssh_key_api.py b/python/lightning_sdk/api/ssh_key_api.py new file mode 100644 index 000000000..a59b1612d --- /dev/null +++ b/python/lightning_sdk/api/ssh_key_api.py @@ -0,0 +1,36 @@ +import json + +from lightning_sdk.api.utils import ApiException, cached_lightning_client +from lightning_sdk.lightning_cloud.openapi import V1CreateSSHPublicKeyRequest, V1SSHPublicKey + + +class SSHKeyRegistrationError(Exception): + """The platform refused to register an SSH public key. The message is the platform's own, for the user.""" + + +def _error_message(e: ApiException) -> str: + body = e.body.decode(errors="replace") if isinstance(e.body, bytes) else e.body + try: + message = json.loads(body or "")["message"] + except (ValueError, KeyError, TypeError): + message = "" + return message or f"HTTP {e.status} {e.reason}" + + +class SSHKeyApi: + """Internal API client for the SSH public keys registered to the signed-in user.""" + + def create_public_key(self, name: str, public_key: str) -> V1SSHPublicKey: + """Register ``public_key`` under ``name``. If the account already has it, the platform returns it unchanged. + + Raises: + SSHKeyRegistrationError: With the platform's reason, e.g. that the key is held by another account. + """ + # Not retried: once retries run out, the client raises a bare Exception without the platform's message. + client = cached_lightning_client(retry=False) + try: + return client.s_sh_public_key_service_create_ssh_public_key( + V1CreateSSHPublicKeyRequest(name=name, public_key=public_key) + ) + except ApiException as e: + raise SSHKeyRegistrationError(_error_message(e)) from None diff --git a/python/lightning_sdk/cli/job/ssh.py b/python/lightning_sdk/cli/job/ssh.py index ae39ee225..3801ede29 100644 --- a/python/lightning_sdk/cli/job/ssh.py +++ b/python/lightning_sdk/cli/job/ssh.py @@ -71,15 +71,8 @@ def ssh_impl( ssh_user = _ssh_user_for_job_id(job_id) - def _run(key_path: str) -> None: - command = f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}" - subprocess.run(command.split()) - + key_path = configure_ssh_internal() try: - _run(configure_ssh_internal()) - except Exception: - # Redownload keys in case they are stale, then retry once. - try: - _run(configure_ssh_internal(force_download=True)) - except Exception: - raise click.ClickException("Failed to establish SSH connection") from None + subprocess.run(f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}".split()) + except OSError: + raise click.ClickException("Failed to establish SSH connection") from None diff --git a/python/lightning_sdk/cli/legacy/configure.py b/python/lightning_sdk/cli/legacy/configure.py index 0473b2e72..78a5c06ab 100644 --- a/python/lightning_sdk/cli/legacy/configure.py +++ b/python/lightning_sdk/cli/legacy/configure.py @@ -1,17 +1,10 @@ -import os -import platform -import uuid -from pathlib import Path -from typing import Optional, Union +from typing import Optional import click -from rich.console import Console -from lightning_sdk.api.utils import _get_cloud_url -from lightning_sdk.cli.legacy.generate import _generate_ssh_config +from lightning_sdk.cli.ssh.configure import OVERWRITE_HELP, configure_studio_ssh from lightning_sdk.cli.utils.auth import require_auth_header from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace -from lightning_sdk.lightning_cloud.login import Auth @click.group(name="configure") @@ -23,39 +16,9 @@ def _configure_ssh_internal( name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False ) -> None: """Internal function to configure SSH without Click decorators.""" - auth = Auth() require_auth_header() - if not auth.api_key and not auth.load(): - raise click.UsageError("An API key is required. Run `lightning login` first.") - console = Console() - ssh_dir = Path.home() / ".ssh" - ssh_dir.mkdir(parents=True, exist_ok=True) - - key_path = ssh_dir / "lightning_rsa" - config_path = ssh_dir / "config" - - # Check if the SSH key already exists - if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite: - console.print(f"SSH key already exists at {key_path}") - else: - _download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name="lightning_rsa", overwrite=overwrite) - console.print(f"SSH key generated and saved to {key_path}") - - # Check if the SSH config already contains the required configuration studio = resolve_studio(name, resolve_teamspace(teamspace)) - config_content = _generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name) - if config_path.exists(): - with config_path.open("r") as config_file: - # check if the host already exists in the config - if f"Host {studio.name}" in config_file.read(): - console.print("SSH config already contains the required configuration.") - return - - with config_path.open("a") as config_file: - config_file.write(os.linesep) - config_file.write(config_content) - config_file.write(os.linesep) - console.print(f"SSH config updated at {config_path}") + configure_studio_ssh(studio.name, f"s_{studio._studio.id}", overwrite=overwrite) @configure.command(name="ssh") @@ -81,52 +44,8 @@ def _configure_ssh_internal( is_flag=True, flag_value=True, default=False, - help="Whether to overwrite the SSH key and config if they already exist.", + help=OVERWRITE_HELP, ) def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None: """Get SSH config entry for a studio.""" _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=overwrite) - - -def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None: - """Download a file from a URL.""" - import requests - - if local_path.exists() and not overwrite: - raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.") - - response = requests.get(url, stream=True) - response.raise_for_status() - - with open(local_path, "wb") as file: - for chunk in response.iter_content(chunk_size=8192): - file.write(chunk) - if chmod is not None: - local_path.chmod(0o600) - - -def _download_ssh_keys( - api_key: str, - key_id: str = "", - ssh_home: Union[str, Path] = "", - ssh_key_name: str = "lightning_rsa", - overwrite: bool = False, -) -> None: - if not ssh_home: - ssh_home = Path.home() / ".ssh" - elif isinstance(ssh_home, str): - ssh_home = Path(ssh_home) - if not key_id: - key_id = str(uuid.uuid4()) - - path_key = ssh_home / ssh_key_name - path_pub = ssh_home / f"{ssh_key_name}.pub" - - # todo: consider hitting the API to get the key pair directly instead of using wget - _download_file( - f"{_get_cloud_url()}/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}", - path_key, - overwrite=overwrite, - chmod=0o600, - ) - _download_file(f"{_get_cloud_url()}/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite) diff --git a/python/lightning_sdk/cli/legacy/generate.py b/python/lightning_sdk/cli/legacy/generate.py index fbc3d9cf7..b6503a615 100644 --- a/python/lightning_sdk/cli/legacy/generate.py +++ b/python/lightning_sdk/cli/legacy/generate.py @@ -3,6 +3,7 @@ import click from rich.console import Console +from lightning_sdk.cli.ssh.common import generate_ssh_config from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace @@ -33,19 +34,6 @@ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None) -> None: """Get SSH config entry for a studio.""" studio = resolve_studio(name, resolve_teamspace(teamspace)) - conf = _generate_ssh_config(key_path="~/.ssh/lightning_rsa", user=f"s_{studio._studio.id}", host=studio.name) + conf = generate_ssh_config(user=f"s_{studio._studio.id}", host=studio.name) # Print the SSH config Console().print(f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + conf) - - -def _generate_ssh_config(key_path: str, host: str, user: str) -> str: - return f"""Host {host} - User {user} - Hostname ssh.lightning.ai - IdentityFile {key_path} - IdentitiesOnly yes - ServerAliveInterval 15 - ServerAliveCountMax 4 - StrictHostKeyChecking no - UserKnownHostsFile=/dev/null - """ diff --git a/python/lightning_sdk/cli/mmt/ssh.py b/python/lightning_sdk/cli/mmt/ssh.py index ad5cd5478..b7312639f 100644 --- a/python/lightning_sdk/cli/mmt/ssh.py +++ b/python/lightning_sdk/cli/mmt/ssh.py @@ -71,15 +71,8 @@ def ssh_impl( ssh_user = _ssh_user_for_job_id(job_id) - def _run(key_path: str) -> None: - command = f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}" - subprocess.run(command.split()) - + key_path = configure_ssh_internal() try: - _run(configure_ssh_internal()) - except Exception: - # Redownload keys in case they are stale, then retry once. - try: - _run(configure_ssh_internal(force_download=True)) - except Exception: - raise click.ClickException("Failed to establish SSH connection") from None + subprocess.run(f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}".split()) + except OSError: + raise click.ClickException("Failed to establish SSH connection") from None diff --git a/python/lightning_sdk/cli/ssh/common.py b/python/lightning_sdk/cli/ssh/common.py index 0654c6afd..cb265c4d8 100644 --- a/python/lightning_sdk/cli/ssh/common.py +++ b/python/lightning_sdk/cli/ssh/common.py @@ -1,14 +1,58 @@ """Shared SSH command helpers.""" +from pathlib import Path -def generate_ssh_config(key_path: str, host: str, user: str) -> str: - return f"""Host {host} - User {user} - Hostname ssh.lightning.ai - IdentityFile {key_path} - IdentitiesOnly yes - ServerAliveInterval 15 - ServerAliveCountMax 4 - StrictHostKeyChecking no - UserKnownHostsFile=/dev/null +from lightning_sdk.cli.utils.ssh_connection import SSH_KEY_NAME + +SSH_HOST = "ssh.lightning.ai" + + +def generate_ssh_config(host: str, user: str) -> str: + """The ``~/.ssh/config`` entry for ``host``, in the same form as the web UI's setup script writes.""" + return ( + f"Host {host}\n" + f" User {user}\n" + f" Hostname {SSH_HOST}\n" + f" IdentityFile ~/.ssh/{SSH_KEY_NAME}\n" + " IdentitiesOnly yes\n" + " ServerAliveInterval 15\n" + " ServerAliveCountMax 4\n" + " StrictHostKeyChecking accept-new\n" + " UserKnownHostsFile ~/.ssh/lightning_known_hosts\n" + ) + + +def host_uses_ssh_key(config: str, host: str) -> bool: + """Whether a ``Host`` or ``Match`` block naming ``host`` already has the Lightning key as an ``IdentityFile``.""" + in_block = False + for line in config.splitlines(): + fields = line.split() + if not fields: + continue + keyword = fields[0].lower() + if keyword in ("host", "match"): + in_block = host in fields[1:] + elif in_block and keyword.startswith("identityfile") and f".ssh/{SSH_KEY_NAME}" in line: + return True + return False + + +def add_ssh_config(host: str, user: str) -> bool: + """Append the entry for ``host`` to ``~/.ssh/config``, unless one there already uses the Lightning key. + + ssh tries every ``IdentityFile`` of every block that matches, so an older entry for ``host`` that names + only ``lightning_rsa`` offers both keys once this one is added after it. + + Returns: + Whether the file was changed. """ + config_path = Path.home() / ".ssh" / "config" + config = config_path.read_text() if config_path.exists() else "" + if host_uses_ssh_key(config, host): + return False + config_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + with config_path.open("a") as config_file: + if config and not config.endswith("\n"): + config_file.write("\n") + config_file.write("\n" + generate_ssh_config(host=host, user=user)) + return True diff --git a/python/lightning_sdk/cli/ssh/configure.py b/python/lightning_sdk/cli/ssh/configure.py index 88e6a19b2..07713991e 100644 --- a/python/lightning_sdk/cli/ssh/configure.py +++ b/python/lightning_sdk/cli/ssh/configure.py @@ -1,19 +1,22 @@ """SSH configure command.""" -from pathlib import Path from typing import Optional import rich_click as click from rich.console import Console -from lightning_sdk.api.utils import _get_cloud_url from lightning_sdk.cli.resource_completion import complete_studio -from lightning_sdk.cli.ssh.common import generate_ssh_config +from lightning_sdk.cli.ssh.common import add_ssh_config from lightning_sdk.cli.utils.auth import require_auth_header from lightning_sdk.cli.utils.logging import LightningCommand from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace -from lightning_sdk.cli.utils.ssh_connection import download_file -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.cli.utils.ssh_connection import ensure_ssh_key + +OVERWRITE_HELP = "Deprecated, and has no effect: the SSH key is never replaced." +OVERWRITE_WARNING = ( + "--overwrite no longer replaces the SSH key. To make a new one, delete {key_path} and its .pub file, " + "then run this again." +) @click.command("configure", cls=LightningCommand) @@ -37,51 +40,26 @@ is_flag=True, flag_value=True, default=False, - help="Whether to overwrite the SSH key and config if they already exist.", + help=OVERWRITE_HELP, ) def configure_ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None: """Get SSH config entry for a studio.""" - import platform - import uuid - resolved_teamspace = resolve_teamspace(teamspace) studio = resolve_studio(name, resolved_teamspace) - auth = Auth() require_auth_header() - if not auth.api_key and not auth.load(): - raise click.UsageError("An API key is required. Run `lightning login` first.") - console = Console() - ssh_dir = Path.home() / ".ssh" - ssh_dir.mkdir(parents=True, exist_ok=True) + configure_studio_ssh(studio.name, f"s_{studio._studio.id}", overwrite=overwrite) - key_path = ssh_dir / "lightning_rsa" - config_path = ssh_dir / "config" - if not key_path.exists() or not (key_path.with_suffix(".pub")).exists() or overwrite: - key_id = str(uuid.uuid4()) - download_file( - f"{_get_cloud_url()}/setup/ssh-gen?t={auth.api_key}&id={key_id}&machineName={platform.node()}", - key_path, - overwrite=overwrite, - chmod=0o600, - ) - download_file( - f"{_get_cloud_url()}/setup/ssh-public?t={auth.api_key}&id={key_id}", - key_path.with_suffix(".pub"), - overwrite=overwrite, - ) - console.print(f"SSH key generated and saved to {key_path}") - else: - console.print(f"SSH key already exists at {key_path}") - - config_content = generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name) - if config_path.exists() and f"Host {studio.name}" in config_path.read_text(): - console.print("SSH config already contains the required configuration.") - return +def configure_studio_ssh(host: str, user: str, overwrite: bool = False) -> None: + """Set up this machine's SSH key and add the ``~/.ssh/config`` entry ``host`` for ``user``.""" + console = Console() + key_path = ensure_ssh_key() + if overwrite: + console.print(OVERWRITE_WARNING.format(key_path=key_path), style="yellow") + console.print(f"Using SSH key {key_path}") - with config_path.open("a") as config_file: - config_file.write("\n") - config_file.write(config_content) - config_file.write("\n") - console.print(f"SSH config updated at {config_path}") + if add_ssh_config(host=host, user=user): + console.print(f"Added {host} to ~/.ssh/config") + else: + console.print(f"~/.ssh/config already uses this key for {host}") diff --git a/python/lightning_sdk/cli/ssh/generate.py b/python/lightning_sdk/cli/ssh/generate.py index 2090146d2..d925c533e 100644 --- a/python/lightning_sdk/cli/ssh/generate.py +++ b/python/lightning_sdk/cli/ssh/generate.py @@ -31,5 +31,5 @@ def generate_ssh(name: Optional[str] = None, teamspace: Optional[str] = None) -> """Get SSH config entry for a studio.""" resolved_teamspace = resolve_teamspace(teamspace) studio = resolve_studio(name, resolved_teamspace) - conf = generate_ssh_config(key_path="~/.ssh/lightning_rsa", user=f"s_{studio._studio.id}", host=studio.name) + conf = generate_ssh_config(user=f"s_{studio._studio.id}", host=studio.name) Console().print(f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + conf) diff --git a/python/lightning_sdk/cli/studio/ssh.py b/python/lightning_sdk/cli/studio/ssh.py index 2653bc02c..0058a4f56 100644 --- a/python/lightning_sdk/cli/studio/ssh.py +++ b/python/lightning_sdk/cli/studio/ssh.py @@ -42,16 +42,10 @@ def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[Lis ssh_private_key_path = configure_ssh_internal() - ssh_options = " -o " + " -o ".join(option) if option else "" - ssh_command = f"ssh -i {ssh_private_key_path}{ssh_options} s_{studio._studio.id}@ssh.lightning.ai" + ssh_options = [arg for value in option or () for arg in ("-o", value)] + ssh_command = ["ssh", "-i", ssh_private_key_path, *ssh_options, f"s_{studio._studio.id}@ssh.lightning.ai"] try: - subprocess.run(ssh_command.split()) - except Exception: - # redownload the keys to be sure they are up to date - ssh_private_key_path = configure_ssh_internal(force_download=True) - try: - subprocess.run(ssh_command.split()) - except Exception: - # TODO: make this a generic CLI error - raise RuntimeError("Failed to establish SSH connection") from None + subprocess.run(ssh_command) + except OSError as e: + raise click.ClickException(f"Failed to establish SSH connection: {e}") from None diff --git a/python/lightning_sdk/cli/utils/ssh_connection.py b/python/lightning_sdk/cli/utils/ssh_connection.py index 144102e70..d40f48bb0 100644 --- a/python/lightning_sdk/cli/utils/ssh_connection.py +++ b/python/lightning_sdk/cli/utils/ssh_connection.py @@ -1,65 +1,89 @@ -import os +"""This machine's Lightning SSH key: created here, and only its public half sent to the platform.""" + import platform -import uuid +import shutil +import subprocess from pathlib import Path -from typing import Optional import rich_click as click -from lightning_sdk.api.utils import _get_cloud_url +from lightning_sdk.api.ssh_key_api import SSHKeyApi, SSHKeyRegistrationError from lightning_sdk.cli.utils.auth import require_auth_header -from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH +# The web UI's setup script uses the same file, so the browser and CLI flows share one key per machine. +SSH_KEY_NAME = "lightning_ed25519" +_MAX_KEY_NAME_LENGTH = 256 + + +def ssh_key_path() -> Path: + return Path.home() / ".ssh" / SSH_KEY_NAME -def configure_ssh_internal(force_download: bool = False) -> str: - """Internal function to configure SSH without Click decorators.""" - auth = Auth() + +def configure_ssh_internal() -> str: + """Make sure this machine's SSH key exists and is registered to the signed-in user, and return its path.""" require_auth_header() - if not auth.api_key and not auth.load(): - raise click.UsageError("An API key is required. Run `lightning login` first.") - return download_ssh_keys(auth.api_key, force_download=force_download) - - -def download_ssh_keys( - api_key: Optional[str], - force_download: bool = False, - ssh_key_name: str = "lightning_rsa", -) -> str: - """Download the SSH key for a User.""" - ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name) - - os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True) - - if not os.path.isfile(ssh_private_key_path) or force_download: - key_id = str(uuid.uuid4()) - download_file( - f"{_get_cloud_url()}/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}", - Path(ssh_private_key_path), - overwrite=True, - chmod=0o600, - ) - download_file( - f"{_get_cloud_url()}/setup/ssh-public?t={api_key}&id={key_id}", - Path(ssh_private_key_path + ".pub"), - overwrite=True, - ) - - return ssh_private_key_path - - -def download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None: - """Download a file from a URL.""" - import requests - - if os.path.isfile(local_path) and not overwrite: - raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.") - - response = requests.get(url, stream=True) - response.raise_for_status() - - with open(local_path, "wb") as file: - for chunk in response.iter_content(chunk_size=8192): - file.write(chunk) - if chmod is not None: - os.chmod(local_path, 0o600) + return str(ensure_ssh_key()) + + +def ensure_ssh_key() -> Path: + """Create the SSH key if it's missing, and register its public key with the signed-in user's account. + + The private key is never replaced. The platform returns the key unchanged when the account already has it, + so this is safe to call on every connect. + """ + key_path = ssh_key_path() + public_key_path = key_path.with_suffix(".pub") + if not key_path.exists(): + _generate_key(key_path) + click.echo(f"Created SSH key {key_path}", err=True) + if not public_key_path.exists(): + _write_public_key(key_path, public_key_path) + + public_key = public_key_path.read_text().strip() + try: + SSHKeyApi().create_public_key(name=_key_name(public_key), public_key=public_key) + except SSHKeyRegistrationError as e: + raise click.ClickException( + f"Could not add the SSH key {public_key_path} to your Lightning account: {e}" + ) from None + return key_path + + +def _ssh_keygen() -> str: + path = shutil.which("ssh-keygen") + if path is None: + raise click.ClickException("ssh-keygen was not found. Install OpenSSH and run this again.") + return path + + +def _generate_key(key_path: Path) -> None: + key_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + # No passphrase, so `lightning studio ssh` works without a terminal. stdin is closed so ssh-keygen + # can't prompt: if another process creates the key first, it refuses to overwrite it. + result = subprocess.run( + [_ssh_keygen(), "-q", "-t", "ed25519", "-N", "", "-C", _default_comment(), "-f", str(key_path)], + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise click.ClickException(f"ssh-keygen could not create {key_path}: {result.stderr.strip()}") + + +def _write_public_key(key_path: Path, public_key_path: Path) -> None: + # Inherits the terminal so ssh-keygen can ask for the passphrase of a key made with one. + result = subprocess.run([_ssh_keygen(), "-y", "-f", str(key_path)], stdout=subprocess.PIPE, text=True) + if result.returncode != 0: + raise click.ClickException(f"ssh-keygen could not read the public key from {key_path}") + public_key_path.write_text(result.stdout) + + +def _default_comment() -> str: + # The web UI's script comments keys the same way, and the comment becomes the key's name. + return f"lightning@{platform.node()}"[:_MAX_KEY_NAME_LENGTH] + + +def _key_name(public_key: str) -> str: + fields = public_key.split(maxsplit=2) + comment = fields[2].strip() if len(fields) > 2 else "" + return (comment or _default_comment())[:_MAX_KEY_NAME_LENGTH] diff --git a/python/tests/api/test_ssh_key_api.py b/python/tests/api/test_ssh_key_api.py new file mode 100644 index 000000000..ef79b9222 --- /dev/null +++ b/python/tests/api/test_ssh_key_api.py @@ -0,0 +1,67 @@ +import json +from unittest.mock import MagicMock, patch + +import pytest + +from lightning_sdk.api.ssh_key_api import SSHKeyApi, SSHKeyRegistrationError + +PUBLIC_KEY = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIExample lightning@laptop" + + +def _response(status: int, body: dict) -> MagicMock: + response = MagicMock() + response.status = status + response.reason = "OK" if status < 400 else "Bad Request" + response.data = json.dumps(body).encode() + response.headers = {"Content-Type": "application/json"} + return response + + +@pytest.fixture() +def _login_token(monkeypatch): + """Credentials as `lightning login` leaves them: a CLI session token and no API key.""" + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "sess_cli-login-token") + + +def test_create_sends_the_public_key_with_the_login_token(_login_token) -> None: + with patch("urllib3.PoolManager.request", return_value=_response(200, {"id": "key-1"})) as request: + created = SSHKeyApi().create_public_key(name="lightning@laptop", public_key=PUBLIC_KEY) + + assert created.id == "key-1" + request.assert_called_once() + method, url = request.call_args.args[:2] + assert method == "POST" + assert url.endswith("/v1/ssh-keys") + assert "?" not in url + assert request.call_args.kwargs["headers"]["Authorization"] == "Bearer sess_cli-login-token" + assert json.loads(request.call_args.kwargs["body"]) == {"name": "lightning@laptop", "publicKey": PUBLIC_KEY} + + +def test_create_refusal_carries_the_platform_message(_login_token) -> None: + refusal = "This SSH public key can't be added. Please generate a new key and add that instead." + with patch( + "urllib3.PoolManager.request", return_value=_response(400, {"code": 3, "message": refusal}) + ) as request, pytest.raises(SSHKeyRegistrationError) as error: + SSHKeyApi().create_public_key(name="laptop", public_key=PUBLIC_KEY) + + assert str(error.value) == refusal + request.assert_called_once() + + +def test_create_is_not_retried_on_a_server_error(_login_token) -> None: + with patch( + "urllib3.PoolManager.request", return_value=_response(500, {"code": 13, "message": "boom"}) + ) as request, pytest.raises(SSHKeyRegistrationError, match="boom"): + SSHKeyApi().create_public_key(name="laptop", public_key=PUBLIC_KEY) + + request.assert_called_once() + + +def test_create_refusal_without_a_json_body_names_the_status(_login_token) -> None: + response = _response(403, {}) + response.data = b"blocked" + with patch("urllib3.PoolManager.request", return_value=response), pytest.raises( + SSHKeyRegistrationError, match="HTTP 403" + ): + SSHKeyApi().create_public_key(name="laptop", public_key=PUBLIC_KEY) diff --git a/python/tests/cli/job/test_ssh.py b/python/tests/cli/job/test_ssh.py index cc09d6379..c45410073 100644 --- a/python/tests/cli/job/test_ssh.py +++ b/python/tests/cli/job/test_ssh.py @@ -113,24 +113,21 @@ def test_ssh_selects_mmt_rank() -> None: run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_rank1@ssh.lightning.ai"]) -def test_ssh_retries_with_fresh_keys_on_failure() -> None: +def test_ssh_reports_a_missing_ssh_client() -> None: job = MagicMock() job.name = "train" job.status = Status.Running job.id = "job_01abc" - configure = MagicMock(side_effect=["/tmp/old_key", "/tmp/new_key"]) - run = MagicMock(side_effect=[OSError("ssh missing"), None]) + configure = MagicMock(return_value="/tmp/key") + run = MagicMock(side_effect=OSError("ssh missing")) with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.job.ssh.resolve_job", return_value=job ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal", configure), patch( "lightning_sdk.cli.job.ssh.subprocess.run", run - ): + ), pytest.raises(click.ClickException, match="Failed to establish SSH connection"): ssh_impl(name="train", teamspace=None) - assert configure.call_count == 2 - configure.assert_any_call(force_download=True) - assert run.call_args_list[0].args[0][2] == "/tmp/old_key" - assert run.call_args_list[1].args[0][2] == "/tmp/new_key" - assert run.call_args_list[1].args[0][-1] == "j_01abc@ssh.lightning.ai" + configure.assert_called_once_with() + run.assert_called_once() diff --git a/python/tests/cli/mmt/test_ssh.py b/python/tests/cli/mmt/test_ssh.py index a15105ef8..ac9ba5acc 100644 --- a/python/tests/cli/mmt/test_ssh.py +++ b/python/tests/cli/mmt/test_ssh.py @@ -180,7 +180,7 @@ def test_ssh_selects_rank_when_names_lack_rank_suffix() -> None: run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_r1@ssh.lightning.ai"]) -def test_ssh_retries_with_fresh_keys_on_failure() -> None: +def test_ssh_reports_a_missing_ssh_client() -> None: rank0 = MagicMock() rank0.name = "train-0" rank0.rank = 0 @@ -191,18 +191,15 @@ def test_ssh_retries_with_fresh_keys_on_failure() -> None: mmt.name = "train" mmt.machines = (rank0,) - configure = MagicMock(side_effect=["/tmp/old_key", "/tmp/new_key"]) - run = MagicMock(side_effect=[OSError("ssh missing"), None]) + configure = MagicMock(return_value="/tmp/key") + run = MagicMock(side_effect=OSError("ssh missing")) with patch("lightning_sdk.cli.mmt.ssh.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.mmt.ssh.resolve_mmt", return_value=mmt ), patch("lightning_sdk.cli.mmt.ssh.configure_ssh_internal", configure), patch( "lightning_sdk.cli.mmt.ssh.subprocess.run", run - ): + ), pytest.raises(click.ClickException, match="Failed to establish SSH connection"): ssh_impl(name="train", teamspace=None) - assert configure.call_count == 2 - configure.assert_any_call(force_download=True) - assert run.call_args_list[0].args[0][2] == "/tmp/old_key" - assert run.call_args_list[1].args[0][2] == "/tmp/new_key" - assert run.call_args_list[1].args[0][-1] == "j_01abc@ssh.lightning.ai" + configure.assert_called_once_with() + run.assert_called_once() diff --git a/python/tests/cli/ssh/test_configure.py b/python/tests/cli/ssh/test_configure.py index 3fd08dede..f310b060a 100644 --- a/python/tests/cli/ssh/test_configure.py +++ b/python/tests/cli/ssh/test_configure.py @@ -1,3 +1,13 @@ +import shutil +import subprocess +from pathlib import Path +from unittest.mock import MagicMock, patch + +import click +import pytest + +from lightning_sdk.api.ssh_key_api import SSHKeyRegistrationError +from lightning_sdk.cli.ssh.common import add_ssh_config, host_uses_ssh_key from tests.cli.help import assert_help_contains, mock_command_logging @@ -10,17 +20,16 @@ def test_legacy_configure_requires_login_without_authentication_fallback() -> No from lightning_sdk.cli.legacy.configure import _configure_ssh_internal - auth = MagicMock() - auth.authenticate.side_effect = AssertionError("callback authentication must not start") + authenticate = MagicMock(side_effect=AssertionError("callback authentication must not start")) with patch( "lightning_sdk.cli.legacy.configure.require_auth_header", side_effect=click.UsageError("Run `lightning login` first."), - ), patch("lightning_sdk.cli.legacy.configure.Auth", return_value=auth), pytest.raises( + ), patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", authenticate), pytest.raises( click.UsageError, match="lightning login" ): _configure_ssh_internal() - auth.authenticate.assert_not_called() + authenticate.assert_not_called() def test_legacy_connect_reaches_non_browser_credential_boundary() -> None: @@ -32,22 +41,21 @@ def test_legacy_connect_reaches_non_browser_credential_boundary() -> None: from lightning_sdk.cli.legacy.connect import studio - auth = MagicMock() - auth.authenticate.side_effect = AssertionError("callback authentication must not start") + authenticate = MagicMock(side_effect=AssertionError("callback authentication must not start")) with patch( "lightning_sdk.cli.legacy.configure.require_auth_header", side_effect=click.UsageError("Run `lightning login` first."), - ), patch("lightning_sdk.cli.legacy.configure.Auth", return_value=auth), patch( + ), patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", authenticate), patch( "lightning_sdk.cli.legacy.connect.subprocess.run", side_effect=AssertionError("SSH must not start without credentials"), ), pytest.raises(click.UsageError, match="lightning login"): studio.callback(name=None, teamspace=None) - auth.authenticate.assert_not_called() + authenticate.assert_not_called() def test_configure_resolves_studio_before_authentication() -> None: - """An unresolved studio does not start authentication or key downloads.""" + """An unresolved studio does not check credentials or create a key.""" from unittest.mock import MagicMock, patch import pytest @@ -55,17 +63,21 @@ def test_configure_resolves_studio_before_authentication() -> None: from lightning_sdk.cli.ssh.configure import configure_ssh - auth = MagicMock() + require_auth_header = MagicMock() + ensure_ssh_key = MagicMock() with patch( "lightning_sdk.cli.ssh.configure.resolve_teamspace", return_value=MagicMock(), ), patch( "lightning_sdk.cli.ssh.configure.resolve_studio", side_effect=click.UsageError("Unknown studio"), - ), patch("lightning_sdk.cli.ssh.configure.Auth", auth), pytest.raises(click.UsageError, match="Unknown studio"): + ), patch("lightning_sdk.cli.ssh.configure.require_auth_header", require_auth_header), patch( + "lightning_sdk.cli.ssh.configure.ensure_ssh_key", ensure_ssh_key + ), pytest.raises(click.UsageError, match="Unknown studio"): configure_ssh.callback(name="missing") - auth.assert_not_called() + require_auth_header.assert_not_called() + ensure_ssh_key.assert_not_called() @mock_command_logging @@ -93,3 +105,194 @@ def test_configure_ssh_legacy_help() -> None: "Use `lightning ssh configure` instead of `lightning configure ssh`.", "Usage: lightning configure ssh [OPTIONS]", ) + + +needs_ssh_keygen = pytest.mark.skipif(shutil.which("ssh-keygen") is None, reason="needs OpenSSH's ssh-keygen") + +OLD_RSA_ENTRY = """Host dev + User s_studio-id + Hostname ssh.lightning.ai + IdentityFile /home/me/.ssh/lightning_rsa + IdentitiesOnly yes + ServerAliveInterval 15 + ServerAliveCountMax 4 + StrictHostKeyChecking no + UserKnownHostsFile=/dev/null +""" + + +@pytest.fixture() +def home(tmp_path, monkeypatch) -> Path: + monkeypatch.setattr(Path, "home", lambda: tmp_path) + return tmp_path + + +@pytest.fixture() +def ssh_key_api(): + api = MagicMock() + with patch("lightning_sdk.cli.utils.ssh_connection.SSHKeyApi", return_value=api), patch( + "requests.get", side_effect=AssertionError("SSH setup must not download anything") + ): + yield api + + +@pytest.fixture() +def studio(): + studio = MagicMock() + studio.name = "dev" + studio._studio.id = "studio-id" + return studio + + +def _configure(studio, *args: str): + from click.testing import CliRunner + + from lightning_sdk.cli.ssh.configure import configure_ssh + + with patch("lightning_sdk.cli.ssh.configure.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.ssh.configure.resolve_studio", return_value=studio + ), patch("lightning_sdk.cli.ssh.configure.require_auth_header", return_value="Bearer token"): + return CliRunner().invoke(configure_ssh, ["--name", "dev", *args]) + + +@needs_ssh_keygen +def test_configure_creates_a_local_key_registers_it_and_adds_the_entry(home, ssh_key_api, studio) -> None: + result = _configure(studio) + + assert result.exit_code == 0, result.output + key_path = home / ".ssh" / "lightning_ed25519" + public_key = key_path.with_suffix(".pub").read_text().strip() + assert public_key.startswith("ssh-ed25519 ") + ssh_key_api.create_public_key.assert_called_once() + assert ssh_key_api.create_public_key.call_args.kwargs["public_key"] == public_key + assert (home / ".ssh" / "config").read_text() == ( + "\n" + "Host dev\n" + " User s_studio-id\n" + " Hostname ssh.lightning.ai\n" + " IdentityFile ~/.ssh/lightning_ed25519\n" + " IdentitiesOnly yes\n" + " ServerAliveInterval 15\n" + " ServerAliveCountMax 4\n" + " StrictHostKeyChecking accept-new\n" + " UserKnownHostsFile ~/.ssh/lightning_known_hosts\n" + ) + assert "Added dev to ~/.ssh/config" in result.output + + +@needs_ssh_keygen +def test_configure_again_reuses_the_key_and_leaves_the_config(home, ssh_key_api, studio) -> None: + assert _configure(studio).exit_code == 0 + key_path = home / ".ssh" / "lightning_ed25519" + private_key = key_path.read_bytes() + config = (home / ".ssh" / "config").read_text() + first_registration = ssh_key_api.create_public_key.call_args + + result = _configure(studio, "--overwrite") + + assert result.exit_code == 0, result.output + assert ssh_key_api.create_public_key.call_count == 2 + assert ssh_key_api.create_public_key.call_args == first_registration + assert key_path.read_bytes() == private_key + assert (home / ".ssh" / "config").read_text() == config + assert "already uses this key for dev" in result.output + assert "--overwrite no longer replaces the SSH key" in result.output + + +@needs_ssh_keygen +def test_configure_shows_the_platform_refusal(home, ssh_key_api, studio) -> None: + refusal = "You can have at most 50 SSH keys. Delete one you no longer use, then add this one again." + ssh_key_api.create_public_key.side_effect = SSHKeyRegistrationError(refusal) + + result = _configure(studio) + + assert result.exit_code == 1 + # rich wraps the error in a box, so drop its border before matching. + assert refusal in " ".join(click.unstyle(result.output).replace("│", " ").split()) + assert not (home / ".ssh" / "config").exists() + + +@needs_ssh_keygen +def test_legacy_configure_creates_and_registers_the_same_key(home, ssh_key_api, studio) -> None: + from lightning_sdk.cli.legacy.configure import _configure_ssh_internal + + with patch("lightning_sdk.cli.legacy.configure.require_auth_header", return_value="Bearer token"), patch( + "lightning_sdk.cli.legacy.configure.resolve_teamspace", return_value=MagicMock() + ), patch("lightning_sdk.cli.legacy.configure.resolve_studio", return_value=studio): + _configure_ssh_internal(name="dev") + + public_key = (home / ".ssh" / "lightning_ed25519.pub").read_text().strip() + ssh_key_api.create_public_key.assert_called_once() + assert ssh_key_api.create_public_key.call_args.kwargs["public_key"] == public_key + assert host_uses_ssh_key((home / ".ssh" / "config").read_text(), "dev") + + +def test_entry_is_added_after_one_that_only_has_the_old_key(home) -> None: + config_path = home / ".ssh" / "config" + config_path.parent.mkdir() + config_path.write_text(OLD_RSA_ENTRY) + + assert add_ssh_config(host="dev", user="s_studio-id") + + config = config_path.read_text() + assert config.startswith(OLD_RSA_ENTRY) + assert host_uses_ssh_key(config, "dev") + assert not add_ssh_config(host="dev", user="s_studio-id") + assert config_path.read_text() == config + + +needs_ssh = pytest.mark.skipif(shutil.which("ssh") is None, reason="needs OpenSSH's ssh") + + +def _ssh_settings(config_path: Path, host: str) -> list: + """The settings ssh would use for ``host``, from ``ssh -G``, which doesn't connect.""" + return subprocess.run( + ["ssh", "-G", "-F", str(config_path), host], + capture_output=True, + text=True, + check=True, + stdin=subprocess.DEVNULL, + ).stdout.splitlines() + + +@needs_ssh +def test_ssh_checks_the_host_key_with_a_new_entry(home) -> None: + config_path = home / ".ssh" / "config" + add_ssh_config(host="dev", user="s_studio-id") + + resolved = _ssh_settings(config_path, "dev") + + assert "stricthostkeychecking accept-new" in resolved + # ssh expands ~ from the user database rather than $HOME, so only the file name is checked. + known_hosts = [line for line in resolved if line.startswith("userknownhostsfile ")] + assert len(known_hosts) == 1 + assert known_hosts[0].endswith("/.ssh/lightning_known_hosts") + + +@needs_ssh +def test_ssh_offers_both_keys_for_a_host_with_an_old_entry(home) -> None: + config_path = home / ".ssh" / "config" + config_path.parent.mkdir() + config_path.write_text(OLD_RSA_ENTRY) + add_ssh_config(host="dev", user="s_studio-id") + + resolved = _ssh_settings(config_path, "dev") + + identity_files = [line.split(maxsplit=1)[1] for line in resolved if line.startswith("identityfile ")] + assert identity_files[0] == "/home/me/.ssh/lightning_rsa" + assert identity_files[1].endswith(".ssh/lightning_ed25519") + assert "identitiesonly yes" in resolved + + +@pytest.mark.parametrize( + ("config", "uses_key"), + [ + ("Host dev\n IdentityFile ~/.ssh/lightning_ed25519\n", True), + ("host other dev\n identityfile=~/.ssh/lightning_ed25519\n", True), + ("Host dev-2\n IdentityFile ~/.ssh/lightning_ed25519\n", False), + ("Host dev\n IdentityFile ~/.ssh/lightning_rsa\nHost x\n IdentityFile ~/.ssh/lightning_ed25519\n", False), + ("", False), + ], +) +def test_host_uses_ssh_key(config: str, uses_key: bool) -> None: + assert host_uses_ssh_key(config, "dev") is uses_key diff --git a/python/tests/cli/studio/test_ssh.py b/python/tests/cli/studio/test_ssh.py index 53b1dc562..77fb8a8b8 100644 --- a/python/tests/cli/studio/test_ssh.py +++ b/python/tests/cli/studio/test_ssh.py @@ -1,3 +1,11 @@ +import shutil +import subprocess +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import rich_click as click + from tests.cli.help import assert_help_contains, command_text, mock_command_logging @@ -26,6 +34,62 @@ def test_ssh_resolves_before_downloading_keys() -> None: configure.assert_not_called() +@pytest.mark.skipif(shutil.which("ssh-keygen") is None, reason="needs OpenSSH's ssh-keygen") +def test_ssh_uses_the_local_key_and_registers_it(tmp_path, monkeypatch) -> None: + """The key is made on this machine once, registered on each run, and passed to ssh as is.""" + from lightning_sdk.cli.studio.ssh import ssh_impl + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + studio = MagicMock() + studio._studio.id = "studio-id" + api = MagicMock() + run = MagicMock() + real_run = subprocess.run + + def run_all_but_ssh(command, **kwargs): + return run(command, **kwargs) if command[0] == "ssh" else real_run(command, **kwargs) + + with patch("lightning_sdk.cli.studio.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.studio.ssh.resolve_studio", return_value=studio + ), patch("lightning_sdk.cli.studio.ssh.save_studio_to_config"), patch( + "lightning_sdk.cli.utils.ssh_connection.require_auth_header", return_value="Bearer token" + ), patch("lightning_sdk.cli.utils.ssh_connection.SSHKeyApi", return_value=api), patch( + "lightning_sdk.cli.studio.ssh.subprocess.run", run_all_but_ssh + ): + ssh_impl(name="dev", teamspace=None, option=["StrictHostKeyChecking no", "LogLevel=ERROR"]) + ssh_impl(name="dev", teamspace=None, option=None) + + key_path = str(tmp_path / ".ssh" / "lightning_ed25519") + public_key = (tmp_path / ".ssh" / "lightning_ed25519.pub").read_text().strip() + assert [c.kwargs["public_key"] for c in api.create_public_key.call_args_list] == [public_key, public_key] + assert run.call_args_list[0].args[0] == [ + "ssh", + "-i", + key_path, + "-o", + "StrictHostKeyChecking no", + "-o", + "LogLevel=ERROR", + "s_studio-id@ssh.lightning.ai", + ] + assert run.call_args_list[1].args[0] == ["ssh", "-i", key_path, "s_studio-id@ssh.lightning.ai"] + + +def test_ssh_reports_a_missing_ssh_client() -> None: + from lightning_sdk.cli.studio.ssh import ssh_impl + + with patch("lightning_sdk.cli.studio.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.studio.ssh.resolve_studio", return_value=MagicMock() + ), patch("lightning_sdk.cli.studio.ssh.save_studio_to_config"), patch( + "lightning_sdk.cli.studio.ssh.configure_ssh_internal", return_value="/tmp/key" + ) as configure, patch( + "lightning_sdk.cli.studio.ssh.subprocess.run", side_effect=FileNotFoundError("ssh") + ), pytest.raises(click.ClickException, match="Failed to establish SSH connection"): + ssh_impl(name="dev", teamspace=None, option=None) + + configure.assert_called_once_with() + + @mock_command_logging def test_ssh_studio(): result_text = command_text("lightning studio ssh --help") diff --git a/python/tests/cli/utils/test_ssh_connection.py b/python/tests/cli/utils/test_ssh_connection.py new file mode 100644 index 000000000..2f806c044 --- /dev/null +++ b/python/tests/cli/utils/test_ssh_connection.py @@ -0,0 +1,138 @@ +import os +import platform +import shutil +import stat +import subprocess +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import rich_click as click + +from lightning_sdk.api.ssh_key_api import SSHKeyRegistrationError +from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal, ensure_ssh_key + +needs_ssh_keygen = pytest.mark.skipif(shutil.which("ssh-keygen") is None, reason="needs OpenSSH's ssh-keygen") + + +@pytest.fixture() +def home(tmp_path, monkeypatch) -> Path: + monkeypatch.setattr(Path, "home", lambda: tmp_path) + return tmp_path + + +@pytest.fixture() +def api(): + api = MagicMock() + with patch("lightning_sdk.cli.utils.ssh_connection.SSHKeyApi", return_value=api): + yield api + + +def _make_key(path: Path, comment: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + subprocess.run( + ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-C", comment, "-f", str(path)], + check=True, + stdin=subprocess.DEVNULL, + ) + + +@needs_ssh_keygen +def test_creates_an_ed25519_key_without_a_passphrase_and_registers_it(home, api) -> None: + key_path = ensure_ssh_key() + + assert key_path == home / ".ssh" / "lightning_ed25519" + public_key = key_path.with_suffix(".pub").read_text().strip() + key_type, _, comment = public_key.split(maxsplit=2) + assert key_type == "ssh-ed25519" + assert comment == f"lightning@{platform.node()}" + if os.name == "posix": + assert stat.S_IMODE(key_path.stat().st_mode) == 0o600 + # No passphrase: reading the key back needs no terminal. + subprocess.run(["ssh-keygen", "-y", "-f", str(key_path)], check=True, stdin=subprocess.DEVNULL, capture_output=True) + api.create_public_key.assert_called_once_with(name=comment, public_key=public_key) + + +@needs_ssh_keygen +def test_a_second_run_registers_the_same_key_again(home, api) -> None: + """The platform returns a key the account already has, so a repeat needs no check here.""" + first = ensure_ssh_key() + private_key = first.read_bytes() + + assert ensure_ssh_key() == first + + assert first.read_bytes() == private_key + assert api.create_public_key.call_count == 2 + assert api.create_public_key.call_args_list[0] == api.create_public_key.call_args_list[1] + + +@needs_ssh_keygen +def test_registers_an_existing_key_under_its_comment_without_replacing_it(home, api) -> None: + key_path = home / ".ssh" / "lightning_ed25519" + _make_key(key_path, "lightning@made-by-the-web-ui") + private_key = key_path.read_bytes() + + ensure_ssh_key() + + assert key_path.read_bytes() == private_key + api.create_public_key.assert_called_once_with( + name="lightning@made-by-the-web-ui", public_key=key_path.with_suffix(".pub").read_text().strip() + ) + + +@needs_ssh_keygen +def test_recreates_a_missing_public_key_from_the_private_key(home, api) -> None: + key_path = home / ".ssh" / "lightning_ed25519" + _make_key(key_path, "lightning@laptop") + public_key = key_path.with_suffix(".pub").read_text() + private_key = key_path.read_bytes() + key_path.with_suffix(".pub").unlink() + + ensure_ssh_key() + + assert key_path.read_bytes() == private_key + assert key_path.with_suffix(".pub").read_text().split()[:2] == public_key.split()[:2] + api.create_public_key.assert_called_once() + + +@needs_ssh_keygen +def test_shows_the_platform_refusal(home, api) -> None: + refusal = "This SSH public key can't be added. Please generate a new key and add that instead." + api.create_public_key.side_effect = SSHKeyRegistrationError(refusal) + + with pytest.raises(click.ClickException) as error: + ensure_ssh_key() + + assert refusal in error.value.message + assert "lightning_ed25519.pub" in error.value.message + + +def test_explains_a_missing_ssh_keygen(home, api) -> None: + with patch("lightning_sdk.cli.utils.ssh_connection.shutil.which", return_value=None), pytest.raises( + click.ClickException, match="ssh-keygen was not found" + ): + ensure_ssh_key() + + api.create_public_key.assert_not_called() + + +def test_configure_requires_login_first(home, api) -> None: + with patch( + "lightning_sdk.cli.utils.ssh_connection.require_auth_header", + side_effect=click.UsageError("Run `lightning login` first."), + ), pytest.raises(click.UsageError, match="lightning login"): + configure_ssh_internal() + + assert not (home / ".ssh").exists() + api.create_public_key.assert_not_called() + + +def test_cli_has_no_setup_token_routes() -> None: + """Keys are made locally: nothing asks the server for one, or puts a credential in a URL.""" + import lightning_sdk.cli + + cli_dir = Path(lightning_sdk.cli.__file__).parent + for source in cli_dir.rglob("*.py"): + text = source.read_text() + assert "/setup/ssh" not in text, source + assert "?t=" not in text, source