diff --git a/apps/api/src/api/tool_instances.py b/apps/api/src/api/tool_instances.py index e9d68d4..eca5437 100644 --- a/apps/api/src/api/tool_instances.py +++ b/apps/api/src/api/tool_instances.py @@ -1611,38 +1611,14 @@ async def start_instance( # Mount selected SSH keys into container home dir if instance.ssh_key_ids: + from src.services.ssh_keys import write_ssh_config, _sanitize_filename + + # Collect all valid keys first + ssh_keys_to_mount = [] for key_id in instance.ssh_key_ids: ssh_key = await session.get(SSHKey, uuid.UUID(key_id)) if ssh_key and ssh_key.user_id == user_id: - try: - ssh_dir = prepare_ssh_key_files( - instance_dir, - ssh_key, - subdir=f"mounts/ssh/{key_id}/.ssh", - uid=container_uid, - gid=container_gid, - ) - ssh_target = os.path.join(home_dir, ".ssh") - extra_volumes.append( - { - "source": ssh_dir, - "target": ssh_target, - "type": "bind", - } - ) - logger.debug( - "Mounted SSH key %s for instance %s to %s", - ssh_key.name, - instance.id, - ssh_target, - ) - except Exception as exc: - logger.error( - "Failed to prepare SSH key %s for instance %s: %s", - key_id, - instance.id, - exc, - ) + ssh_keys_to_mount.append(ssh_key) else: logger.warning( "SSH key %s not found or not authorized for user %s", @@ -1650,6 +1626,79 @@ async def start_instance( user_id, ) + if ssh_keys_to_mount: + # Use a single shared .ssh directory so all keys are visible + ssh_dir = os.path.join(instance_dir, "mounts", "ssh", ".ssh") + os.makedirs(ssh_dir, exist_ok=True) + + key_filenames = [] + for ssh_key in ssh_keys_to_mount: + # Use sanitized key name as filename prefix to avoid collisions + key_name = _sanitize_filename(ssh_key.name) + # If multiple keys have the same name, append a short hash + base_filename = f"id_ed25519_{key_name}" + filename = base_filename + counter = 1 + while filename in key_filenames: + filename = f"{base_filename}_{counter}" + counter += 1 + key_filenames.append(filename) + + try: + prepare_ssh_key_files( + instance_dir, + ssh_key, + subdir="mounts/ssh/.ssh", + uid=container_uid, + gid=container_gid, + key_filename=filename, + write_config=False, + ) + logger.debug( + "Prepared SSH key %s as %s for instance %s", + ssh_key.name, + filename, + instance.id, + ) + except Exception as exc: + logger.error( + "Failed to prepare SSH key %s for instance %s: %s", + ssh_key.id, + instance.id, + exc, + ) + + # Write combined SSH config with all keys + try: + write_ssh_config( + ssh_dir, + key_filenames, + uid=container_uid, + gid=container_gid, + ) + except Exception as exc: + logger.error( + "Failed to write SSH config for instance %s: %s", + instance.id, + exc, + ) + + # Mount the single .ssh directory into container home + ssh_target = os.path.join(home_dir, ".ssh") + extra_volumes.append( + { + "source": ssh_dir, + "target": ssh_target, + "type": "bind", + } + ) + logger.debug( + "Mounted %d SSH key(s) for instance %s to %s", + len(ssh_keys_to_mount), + instance.id, + ssh_target, + ) + # ── MANIFEST-BASED FLOW ────────────────────────────────────── resolved_manifest = None diff --git a/apps/api/src/services/ssh_keys.py b/apps/api/src/services/ssh_keys.py index 2ab4c67..cd750d4 100644 --- a/apps/api/src/services/ssh_keys.py +++ b/apps/api/src/services/ssh_keys.py @@ -2,6 +2,7 @@ import logging import os +import re from pathlib import Path from cryptography.fernet import Fernet @@ -22,12 +23,28 @@ def _get_fernet() -> Fernet: return Fernet(key) +def _sanitize_filename(name: str) -> str: + """Sanitize a string for use as a filename. + + Replaces non-alphanumeric characters with underscores and strips + leading/trailing underscores. + """ + sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", name) + sanitized = sanitized.strip("_") + # Ensure it's not empty + if not sanitized: + sanitized = "key" + return sanitized + + def prepare_ssh_key_files( instance_dir: str, ssh_key, subdir: str = ".ssh", uid: int | None = None, gid: int | None = None, + key_filename: str = "id_ed25519", + write_config: bool = True, ) -> str: """Decrypt and write SSH key files to instance directory for container mounting. @@ -37,6 +54,12 @@ def prepare_ssh_key_files( subdir: Subdirectory within instance_dir to write to (default: ".ssh") uid: Optional UID to own the files (for bind-mount into non-root container) gid: Optional GID to own the files + key_filename: Base filename for the key pair (default: "id_ed25519"). + The private key will be named "{key_filename}" and the public key + "{key_filename}.pub". + write_config: Whether to write an SSH config file (default: True). + Set to False when combining multiple keys into one directory, + then call write_ssh_config() separately. Returns: Path to the .ssh directory @@ -49,51 +72,101 @@ def prepare_ssh_key_files( private_key = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode() # Write private key with restricted permissions - private_key_path = ssh_dir / "id_ed25519" + private_key_path = ssh_dir / key_filename private_key_path.write_text(private_key) os.chmod(private_key_path, 0o600) # Write public key - public_key_path = ssh_dir / "id_ed25519.pub" + public_key_path = ssh_dir / f"{key_filename}.pub" public_key_path.write_text(ssh_key.public_key) os.chmod(public_key_path, 0o644) - # Write SSH config - config_path = ssh_dir / "config" - config_content = """Host * + # Write SSH config (only if requested) + if write_config: + config_path = ssh_dir / "config" + config_content = f"""Host * StrictHostKeyChecking no UserKnownHostsFile /dev/null - IdentityFile ~/.ssh/id_ed25519 + IdentityFile ~/.ssh/{key_filename} IdentitiesOnly yes """ + config_path.write_text(config_content) + os.chmod(config_path, 0o644) + + # Set ownership to target container user if requested + if uid is not None or gid is not None: + effective_uid = uid if uid is not None else -1 + effective_gid = gid if gid is not None else -1 + try: + os.chown(ssh_dir, effective_uid, effective_gid) + os.chown(private_key_path, effective_uid, effective_gid) + os.chown(public_key_path, effective_uid, effective_gid) + os.chown(config_path, effective_uid, effective_gid) + logger.debug( + "Set SSH key ownership to uid=%s gid=%s for %s", + effective_uid, + effective_gid, + ssh_dir, + ) + except PermissionError as exc: + logger.warning( + "Cannot chown SSH keys to uid=%s gid=%s (running as uid=%s): %s", + effective_uid, + effective_gid, + os.getuid(), + exc, + ) + else: + # Still chown the key files even if we didn't write config + if uid is not None or gid is not None: + effective_uid = uid if uid is not None else -1 + effective_gid = gid if gid is not None else -1 + try: + os.chown(private_key_path, effective_uid, effective_gid) + os.chown(public_key_path, effective_uid, effective_gid) + except PermissionError: + pass + + return str(ssh_dir) + + +def write_ssh_config( + ssh_dir: str, + key_filenames: list[str], + uid: int | None = None, + gid: int | None = None, +) -> None: + """Write an SSH config file that includes multiple IdentityFile entries. + + Args: + ssh_dir: Path to the .ssh directory + key_filenames: List of key filenames (without .pub extension) + uid: Optional UID to own the config file + gid: Optional GID to own the config file + """ + ssh_dir_path = Path(ssh_dir) + ssh_dir_path.mkdir(parents=True, exist_ok=True) + + config_path = ssh_dir_path / "config" + lines = ["Host *"] + lines.append(" StrictHostKeyChecking no") + lines.append(" UserKnownHostsFile /dev/null") + lines.append(" IdentitiesOnly yes") + for filename in key_filenames: + lines.append(f" IdentityFile ~/.ssh/{filename}") + lines.append("") + + config_content = "\n".join(lines) config_path.write_text(config_content) os.chmod(config_path, 0o644) - # Set ownership to target container user if requested if uid is not None or gid is not None: effective_uid = uid if uid is not None else -1 effective_gid = gid if gid is not None else -1 try: - os.chown(ssh_dir, effective_uid, effective_gid) - os.chown(private_key_path, effective_uid, effective_gid) - os.chown(public_key_path, effective_uid, effective_gid) os.chown(config_path, effective_uid, effective_gid) - logger.debug( - "Set SSH key ownership to uid=%s gid=%s for %s", - effective_uid, - effective_gid, - ssh_dir, - ) - except PermissionError as exc: - logger.warning( - "Cannot chown SSH keys to uid=%s gid=%s (running as uid=%s): %s", - effective_uid, - effective_gid, - os.getuid(), - exc, - ) - - return str(ssh_dir) + except PermissionError: + pass def cleanup_ssh_key_files(instance_dir: str) -> None: diff --git a/apps/api/src/services/terminal_session.py b/apps/api/src/services/terminal_session.py index 01ab59c..e2609c0 100644 --- a/apps/api/src/services/terminal_session.py +++ b/apps/api/src/services/terminal_session.py @@ -164,7 +164,9 @@ class TerminalSession: self._read_handler_set = True logger.debug("Started event-driven reading for session %s", self.session_id) except Exception as exc: - logger.error("Failed to start reading for session %s: %s", self.session_id, exc) + logger.error( + "Failed to start reading for session %s: %s", self.session_id, exc + ) def _stop_reading(self) -> None: """Unregister PTY master fd from asyncio event loop.""" @@ -263,9 +265,7 @@ class TerminalSession: if self._ack_timeout_handle: self._ack_timeout_handle.cancel() loop = asyncio.get_event_loop() - self._ack_timeout_handle = loop.call_later( - 5.0, self._ack_timeout_fallback - ) + self._ack_timeout_handle = loop.call_later(5.0, self._ack_timeout_fallback) def _ack_timeout_fallback(self) -> None: """If no ack received for 5s, assume client is dead and resume.""" @@ -281,7 +281,11 @@ class TerminalSession: """Pause reading from PTY due to flow control.""" self._paused = True self._stop_reading() - logger.debug("Paused output for session %s (%d unacked)", self.session_id, self._unacknowledged_bytes) + logger.debug( + "Paused output for session %s (%d unacked)", + self.session_id, + self._unacknowledged_bytes, + ) def _resume_output(self) -> None: """Resume reading from PTY.""" @@ -327,7 +331,9 @@ class TerminalSession: self._cols = cols self._rows = rows - logger.debug("resize() called for session %s: %sx%s", self.session_id, cols, rows) + logger.debug( + "resize() called for session %s: %sx%s", self.session_id, cols, rows + ) self._set_terminal_size(cols, rows) # Send SIGWINCH to docker exec process diff --git a/apps/api/tests/unit/test_tool_instances_legacy.py b/apps/api/tests/unit/test_tool_instances_legacy.py index a2a1d86..2bcc1e2 100644 --- a/apps/api/tests/unit/test_tool_instances_legacy.py +++ b/apps/api/tests/unit/test_tool_instances_legacy.py @@ -837,23 +837,24 @@ class TestStartInstanceSshPermissions: mock_session.get.side_effect = _get with patch("os.path.exists", return_value=True): - with patch( - "src.api.tool_instances._prepare_manifest_instance" - ) as mock_prepare: - mock_prepare.return_value = ( - "headquarter/test:latest", - "services:\n app:\n image: test", - {"name": "test-manifest", "user": {"name": "user"}}, - "/home/user", - ) - result = await start_instance( - project_id=fake_project_id, - repo_id=fake_repo_id, - instance_id=fake_instance_id, - data=None, - user_id=fake_user_id, - session=mock_session, - ) + with patch("os.makedirs"): + with patch( + "src.api.tool_instances._prepare_manifest_instance" + ) as mock_prepare: + mock_prepare.return_value = ( + "headquarter/test:latest", + "services:\n app:\n image: test", + {"name": "test-manifest", "user": {"name": "user"}}, + "/home/user", + ) + result = await start_instance( + project_id=fake_project_id, + repo_id=fake_repo_id, + instance_id=fake_instance_id, + data=None, + user_id=fake_user_id, + session=mock_session, + ) assert result["status"] == "running" mock_apply_ssh.assert_called_once_with("abc123", "/home/user/.ssh", "user") @@ -960,15 +961,16 @@ class TestStartInstanceSshPermissions: mock_session.get.side_effect = _get with patch("os.path.exists", return_value=True): - with patch("src.api.tool_instances._modify_compose_file"): - result = await start_instance( - project_id=fake_project_id, - repo_id=fake_repo_id, - instance_id=fake_instance_id, - data=None, - user_id=fake_user_id, - session=mock_session, - ) + with patch("os.makedirs"): + with patch("src.api.tool_instances._modify_compose_file"): + result = await start_instance( + project_id=fake_project_id, + repo_id=fake_repo_id, + instance_id=fake_instance_id, + data=None, + user_id=fake_user_id, + session=mock_session, + ) assert result["status"] == "running" mock_apply_ssh.assert_called_once_with("abc123", "/root/.ssh", "root")