Initial commit

tests/conftest.py

@@ -0,0 +1,170 @@
+"""
+Shared pytest fixtures for the MCP Privileged Access test suite.
+
+──────────────────────────────────────────────────────────────────────────────
+HOW MCP TOOLS WORK  (read this to understand what the tests are testing)
+──────────────────────────────────────────────────────────────────────────────
+
+An MCP tool is just an async Python function decorated with @mcp.tool().
+The decorator registers the function in FastMCP's tool registry — it does NOT
+change how the function itself is called.  This means tests can call tool
+functions directly as plain async functions:
+
+    result = await ssh_execute(host="...", command="...", ...)
+
+The MCP framework wraps tool calls in a JSON-RPC envelope when running for
+real, but for unit tests we skip the envelope entirely.
+
+FastMCP injects a Context object as the `ctx` parameter.  The Context carries:
+  • ctx.info(msg)  — progress notification sent back to the caller
+  • ctx.error(msg) — error notification
+  • ctx.request_context.request — the raw HTTP request (for IP extraction etc.)
+
+In tests we pass a MagicMock for Context so we can assert what was logged
+without making any real network calls.
+
+SECRET HANDLE LIFECYCLE
+  1. CyberArk MCP calls secret_store.store(username, password) → "secret://abc…"
+  2. Handle is returned to Claude (only the handle token, never the password).
+  3. SSH / PowerShell / DB tool calls secret_store.resolve(handle) → (user, pass).
+  4. If handle_single_use=True (default), the handle is deleted after step 3.
+  5. The password is used for the connection and then deleted from local scope.
+
+This means:
+  • Each test that needs to resolve a credential must create its OWN fresh handle.
+  • Attempting to resolve the same handle twice raises KeyError.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Must be set before any mcp_privileged import triggers Settings() at module level.
+os.environ.setdefault("MCP_API_KEYS", "test-key-for-pytest")
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from mcp_privileged.secret_store import secret_store
+
+
+# ── Context mock ──────────────────────────────────────────────────────────────
+
+@pytest.fixture
+def mock_ctx() -> MagicMock:
+    """
+    Minimal mock of the FastMCP Context object.
+
+    ctx.info() and ctx.error() are AsyncMocks so tests can await them and
+    also assert what messages were emitted:
+
+        ctx.error.assert_awaited_once()
+        assert "expired" in str(ctx.error.call_args)
+    """
+    ctx = MagicMock()
+    ctx.info = AsyncMock()
+    ctx.error = AsyncMock()
+    # _extract_client_ip reads these — plain dict works fine
+    ctx.request_context.request.headers = {}
+    ctx.request_context.request.client = None
+    return ctx
+
+
+# ── Credential handle factory ─────────────────────────────────────────────────
+
+@pytest.fixture
+async def credential_handle() -> str:
+    """
+    Store a test credential and return a fresh secret handle.
+
+    Because handle_single_use=True (default), each test fixture invocation
+    creates a NEW handle so tests don't step on each other.
+
+    Usage:
+        async def test_something(credential_handle, mock_ctx):
+            result = await ssh_execute(..., secret_handle=credential_handle, ctx=mock_ctx)
+    """
+    return await secret_store.store("svc_user", "P@ssw0rd!")
+
+
+@pytest.fixture
+async def credential_handle_with_details() -> tuple[str, str, str]:
+    """
+    Return (handle, username, password) so tests can assert on the values.
+    The password is exposed here ONLY for test assertions — never in prod code.
+    """
+    username = "admin_user"
+    password = "S3cr3tP@ss123"
+    handle = await secret_store.store(username, password)
+    return handle, username, password
+
+
+# ── asyncssh mock helpers ─────────────────────────────────────────────────────
+
+def make_ssh_cm(
+    stdout: str = "",
+    stderr: str = "",
+    exit_status: int = 0,
+) -> tuple[AsyncMock, AsyncMock]:
+    """
+    Build a mock for asyncssh.connect used as an async context manager.
+
+    asyncssh.connect() is called as:
+        async with asyncssh.connect(host, port=..., ...) as conn:
+            result = await conn.run(command, timeout=...)
+
+    The mock chain:
+        asyncssh.connect(...)         → returns mock_cm
+        async with mock_cm as conn:   → calls mock_cm.__aenter__() → mock_conn
+        await conn.run(...)           → returns MagicMock(stdout, stderr, exit_status)
+
+    Returns (mock_cm, mock_conn) so tests can inspect call_args on mock_conn.run.
+    """
+    mock_conn = AsyncMock()
+    mock_conn.run = AsyncMock(
+        return_value=MagicMock(stdout=stdout, stderr=stderr, exit_status=exit_status)
+    )
+    mock_cm = AsyncMock()
+    mock_cm.__aenter__ = AsyncMock(return_value=mock_conn)
+    mock_cm.__aexit__ = AsyncMock(return_value=False)
+    return mock_cm, mock_conn
+
+
+# ── pypsrp mock helpers ────────────────────────────────────────────────────────
+
+def make_ps_result(
+    output: list[str] | None = None,
+    had_errors: bool = False,
+    errors: list[str] | None = None,
+) -> tuple[list[str], bool, list[str]]:
+    """
+    Build the tuple returned by _run_ps_sync so tests can patch it directly.
+
+    Usage:
+        with patch(
+            "mcp_privileged.powershell.server._run_ps_sync",
+            return_value=make_ps_result(output=["Hello"]),
+        ):
+            ...
+    """
+    return (output or [], had_errors, errors or [])
+
+
+# ── asyncpg / aiomysql mock helpers ───────────────────────────────────────────
+
+def make_db_result(
+    columns: list[str],
+    rows: list[list],
+) -> tuple[list[str], list[list]]:
+    """
+    Build the (columns, rows) tuple returned by _dispatch_query.
+
+    Usage:
+        with patch(
+            "mcp_privileged.database.server._dispatch_query",
+            new=AsyncMock(return_value=make_db_result(["id", "name"], [[1, "Alice"]])),
+        ):
+            ...
+    """
+    return columns, rows

tests/test_auth.py

@@ -0,0 +1,72 @@
+"""
+Tests for the API key middleware.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from fastapi.testclient import TestClient
+
+import mcp_privileged.auth as auth_module
+from mcp_privileged.auth import ApiKeyMiddleware
+
+_FAKE_SETTINGS = SimpleNamespace(mcp_api_keys={"valid-key-1", "valid-key-2"})
+
+
+def _make_app() -> FastAPI:
+    app = FastAPI()
+    app.add_middleware(ApiKeyMiddleware)
+
+    @app.get("/mcp/test")
+    async def protected() -> JSONResponse:
+        return JSONResponse({"ok": True})
+
+    @app.get("/health")
+    async def health() -> JSONResponse:
+        return JSONResponse({"status": "ok"})
+
+    return app
+
+
+@pytest.fixture
+def client(monkeypatch) -> TestClient:
+    monkeypatch.setattr(auth_module, "settings", _FAKE_SETTINGS)
+    return TestClient(_make_app(), raise_server_exceptions=True)
+
+
+def test_health_requires_no_auth(client: TestClient) -> None:
+    response = client.get("/health")
+    assert response.status_code == 200
+
+
+def test_missing_key_returns_401(client: TestClient) -> None:
+    response = client.get("/mcp/test")
+    assert response.status_code == 401
+
+
+def test_invalid_key_returns_401(client: TestClient) -> None:
+    response = client.get("/mcp/test", headers={"X-API-Key": "wrong-key"})
+    assert response.status_code == 401
+
+
+def test_valid_x_api_key_header(client: TestClient) -> None:
+    response = client.get("/mcp/test", headers={"X-API-Key": "valid-key-1"})
+    assert response.status_code == 200
+
+
+def test_valid_bearer_token(client: TestClient) -> None:
+    response = client.get(
+        "/mcp/test", headers={"Authorization": "Bearer valid-key-2"}
+    )
+    assert response.status_code == 200
+
+
+def test_bearer_case_insensitive(client: TestClient) -> None:
+    response = client.get(
+        "/mcp/test", headers={"Authorization": "bearer valid-key-1"}
+    )
+    assert response.status_code == 200

tests/test_cyberark_client.py

@@ -0,0 +1,165 @@
+"""
+Tests for the CyberArk CCP client.
+
+All tests use httpx.MockTransport to avoid real network calls.
+"""
+
+from __future__ import annotations
+
+import json
+
+import httpx
+import pytest
+
+from mcp_privileged.cyberark.client import (
+    CyberArkCCPClient,
+    CyberArkError,
+    Credential,
+)
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _ok_response(username: str = "svc_account", password: str = "S3cr3tP@ss") -> dict:
+    return {
+        "Content": password,
+        "UserName": username,
+        "Address": "db.internal",
+        "Safe": "PROD-DB",
+        "Folder": "Root",
+        "Name": "PROD-DB-svc_account",
+        "PlatformID": "Oracle",
+        "PasswordChangeInProcess": "False",
+    }
+
+
+def _error_response(code: str, msg: str) -> dict:
+    return {"ErrorCode": code, "ErrorMsg": msg}
+
+
+class _MockTransport(httpx.AsyncBaseTransport):
+    """Simple mock transport that returns a pre-set response."""
+
+    def __init__(self, status_code: int, body: dict) -> None:
+        self._status = status_code
+        self._body = body
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            self._status,
+            headers={"content-type": "application/json"},
+            content=json.dumps(self._body).encode(),
+            request=request,
+        )
+
+
+def _client_with_transport(transport: httpx.AsyncBaseTransport) -> CyberArkCCPClient:
+    """Create a CyberArkCCPClient with a mock transport pre-injected."""
+    client = CyberArkCCPClient()
+    client._http = httpx.AsyncClient(transport=transport)
+    return client
+
+
+# ── Tests ─────────────────────────────────────────────────────────────────────
+
+async def test_get_credential_success() -> None:
+    transport = _MockTransport(200, _ok_response())
+    client = _client_with_transport(transport)
+
+    cred = await client.get_credential(
+        app_id="MyApp", safe="PROD-DB", object_name="PROD-DB-svc_account"
+    )
+
+    assert isinstance(cred, Credential)
+    assert cred.username == "svc_account"
+    assert cred.password == "S3cr3tP@ss"
+    assert cred.address == "db.internal"
+    assert cred.platform_id == "Oracle"
+    assert cred.password_change_in_process is False
+
+
+async def test_get_credential_not_found_raises() -> None:
+    transport = _MockTransport(404, _error_response("APPAP007E", "Credential object not found"))
+    client = _client_with_transport(transport)
+
+    with pytest.raises(CyberArkError) as exc_info:
+        await client.get_credential(app_id="MyApp", safe="PROD-DB", object_name="missing")
+
+    assert exc_info.value.error_code == "APPAP007E"
+    assert exc_info.value.status_code == 404
+
+
+async def test_get_credential_auth_failure_raises() -> None:
+    transport = _MockTransport(403, _error_response("APPAP006E", "Authentication failure"))
+    client = _client_with_transport(transport)
+
+    with pytest.raises(CyberArkError) as exc_info:
+        await client.get_credential(app_id="BadApp", safe="PROD-DB", object_name="obj")
+
+    assert exc_info.value.error_code == "APPAP006E"
+    assert exc_info.value.status_code == 403
+
+
+async def test_get_credential_unknown_error_code() -> None:
+    """Unknown error codes should still raise CyberArkError with the raw message."""
+    transport = _MockTransport(500, _error_response("ZZZZZ999E", "Unexpected internal error"))
+    client = _client_with_transport(transport)
+
+    with pytest.raises(CyberArkError) as exc_info:
+        await client.get_credential(app_id="MyApp", safe="S", object_name="O")
+
+    assert "Unexpected internal error" in str(exc_info.value)
+
+
+async def test_get_credential_non_json_body() -> None:
+    """Non-JSON 500 responses should still raise a CyberArkError."""
+    class _HtmlTransport(httpx.AsyncBaseTransport):
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            return httpx.Response(500, content=b"<html>Internal Server Error</html>", request=request)
+
+    client = _client_with_transport(_HtmlTransport())
+    with pytest.raises(CyberArkError) as exc_info:
+        await client.get_credential(app_id="MyApp", safe="S", object_name="O")
+    assert exc_info.value.status_code == 500
+
+
+async def test_connect_error_raises() -> None:
+    class _FailTransport(httpx.AsyncBaseTransport):
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            raise httpx.ConnectError("Connection refused")
+
+    client = _client_with_transport(_FailTransport())
+    with pytest.raises(CyberArkError, match="Cannot reach CCP"):
+        await client.get_credential(app_id="MyApp", safe="S", object_name="O")
+
+
+async def test_timeout_raises() -> None:
+    class _TimeoutTransport(httpx.AsyncBaseTransport):
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            raise httpx.ReadTimeout("Timed out")
+
+    client = _client_with_transport(_TimeoutTransport())
+    with pytest.raises(CyberArkError, match="timed out"):
+        await client.get_credential(app_id="MyApp", safe="S", object_name="O")
+
+
+async def test_assert_started_raises_if_not_started() -> None:
+    client = CyberArkCCPClient()
+    with pytest.raises(RuntimeError, match="not been started"):
+        await client.get_credential(app_id="A", safe="S", object_name="O")
+
+
+async def test_list_safes_raises_not_implemented() -> None:
+    client = _client_with_transport(_MockTransport(200, {}))
+    with pytest.raises(NotImplementedError):
+        await client.list_safes("MyApp")
+
+
+async def test_password_not_in_error_message() -> None:
+    """Ensure passwords are never leaked into exception messages."""
+    transport = _MockTransport(200, _ok_response(password="SuperSecret123"))
+    client = _client_with_transport(transport)
+    cred = await client.get_credential(app_id="A", safe="S", object_name="O")
+    assert cred.password == "SuperSecret123"
+    # The Credential dataclass itself is fine, but error paths must not include it
+    # (no error raised here — just confirming the happy path returns it correctly
+    #  and the password doesn't appear in repr of the transport or request)

tests/test_database_server.py

@@ -0,0 +1,303 @@
+"""
+Tests for the Database MCP tool (db_query).
+
+We patch _dispatch_query (the internal router) rather than individual drivers
+so the tests stay driver-agnostic.  Driver-specific tests (asyncpg / aiomysql /
+pyodbc) are covered in the integration section at the bottom.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from mcp_privileged.config import settings
+from mcp_privileged.database.server import (
+    _cell_str,
+    _format_result,
+    db_query,
+)
+from mcp_privileged.secret_store import secret_store
+from tests.conftest import make_db_result
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+async def _handle(username: str = "db_svc", password: str = "DbP@ss!") -> str:
+    return await secret_store.store(username, password)
+
+
+def _patch_dispatch(columns: list[str], rows: list[list]):
+    """Patch _dispatch_query to return a pre-built result without hitting a DB."""
+    return patch(
+        "mcp_privileged.database.server._dispatch_query",
+        new=AsyncMock(return_value=make_db_result(columns, rows)),
+    )
+
+
+# ── Tests ─────────────────────────────────────────────────────────────────────
+
+async def test_db_query_success_postgres(mock_ctx) -> None:
+    """Happy path: postgres query returns columns + rows."""
+    handle = await _handle()
+    cols = ["id", "name", "email"]
+    rows = [[1, "Alice", "alice@example.com"], [2, "Bob", "bob@example.com"]]
+
+    with _patch_dispatch(cols, rows):
+        result = await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT id, name, email FROM users",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+        )
+
+    assert "Rows returned: 2" in result
+    assert "id" in result and "name" in result and "email" in result
+    assert "Alice" in result
+    assert "Database: mydb (postgres)" in result
+
+
+async def test_db_query_success_mysql(mock_ctx) -> None:
+    """MySQL variant — db_type routing and label are correct."""
+    handle = await _handle()
+
+    with _patch_dispatch(["host_name"], [["mysql-server-01"]]):
+        result = await db_query(
+            host="mysql.internal",
+            database="ops",
+            query="SELECT @@hostname",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="mysql",
+        )
+
+    assert "mysql" in result
+    assert "mysql-server-01" in result
+
+
+async def test_db_query_success_mssql(mock_ctx) -> None:
+    """SQL Server variant — db_type routing and label are correct."""
+    handle = await _handle()
+
+    with _patch_dispatch(["name"], [["SQLSERVER01"]]):
+        result = await db_query(
+            host="sql.internal",
+            database="master",
+            query="SELECT @@SERVERNAME",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="mssql",
+        )
+
+    assert "mssql" in result
+    assert "SQLSERVER01" in result
+
+
+async def test_db_query_default_port_resolved(mock_ctx) -> None:
+    """port=0 triggers the default port for the db_type."""
+    handle = await _handle()
+
+    with _patch_dispatch(["v"], [[42]]) as mock_dispatch:
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 42",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+            port=0,
+        )
+
+    _, kwargs = mock_dispatch.call_args
+    assert kwargs["port"] == 5432
+
+
+async def test_db_query_custom_port_forwarded(mock_ctx) -> None:
+    """Explicit port is forwarded unchanged."""
+    handle = await _handle()
+
+    with _patch_dispatch(["v"], [[1]]) as mock_dispatch:
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+            port=15432,
+        )
+
+    _, kwargs = mock_dispatch.call_args
+    assert kwargs["port"] == 15432
+
+
+async def test_db_query_username_override(mock_ctx) -> None:
+    """username_override replaces the credential username."""
+    handle = await _handle(username="readonly_user")
+
+    with _patch_dispatch(["v"], [[1]]) as mock_dispatch:
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+            username_override="dba_user",
+        )
+
+    _, kwargs = mock_dispatch.call_args
+    assert kwargs["username"] == "dba_user"
+
+
+async def test_db_query_invalid_db_type(mock_ctx) -> None:
+    """Unknown db_type raises ValueError before touching the credential store."""
+    handle = await _handle()
+
+    with pytest.raises(ValueError, match="Unsupported db_type"):
+        await db_query(
+            host="db.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="oracle",
+        )
+
+
+async def test_db_query_invalid_handle(mock_ctx) -> None:
+    """Unknown handle raises KeyError and calls ctx.error."""
+    with pytest.raises(KeyError):
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle="secret://doesnotexist0000000000000000",
+            ctx=mock_ctx,
+            db_type="postgres",
+        )
+
+    mock_ctx.error.assert_awaited_once()
+
+
+async def test_db_query_driver_exception_propagates(mock_ctx) -> None:
+    """Exceptions from _dispatch_query propagate and call ctx.error."""
+    handle = await _handle()
+
+    with patch(
+        "mcp_privileged.database.server._dispatch_query",
+        new=AsyncMock(side_effect=ConnectionRefusedError("DB port closed")),
+    ):
+        with pytest.raises(ConnectionRefusedError):
+            await db_query(
+                host="pg.internal",
+                database="mydb",
+                query="SELECT 1",
+                secret_handle=handle,
+                ctx=mock_ctx,
+                db_type="postgres",
+            )
+
+    mock_ctx.error.assert_awaited_once()
+
+
+async def test_db_query_rows_capped(mock_ctx) -> None:
+    """Rows exceeding db_max_rows are truncated and the result says so."""
+    handle = await _handle()
+    many_rows = [[i, f"user_{i}"] for i in range(2000)]
+
+    with _patch_dispatch(["id", "name"], many_rows):
+        with patch.object(settings, "db_max_rows", 10):
+            result = await db_query(
+                host="pg.internal",
+                database="mydb",
+                query="SELECT id, name FROM big_table",
+                secret_handle=handle,
+                ctx=mock_ctx,
+                db_type="postgres",
+            )
+
+    assert "Rows returned: 10" in result
+    assert "more rows exist" in result
+
+
+async def test_db_query_empty_result(mock_ctx) -> None:
+    """An empty result set is handled gracefully."""
+    handle = await _handle()
+
+    with _patch_dispatch([], []):
+        result = await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1 WHERE false",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+        )
+
+    assert "Rows returned: 0" in result
+    assert "No rows returned" in result
+
+
+async def test_db_query_password_not_in_ctx_messages(mock_ctx) -> None:
+    """The credential password must never leak into ctx.info or ctx.error."""
+    secret_password = "DB$ecretPass99"
+    handle = await secret_store.store("db_user", secret_password)
+
+    with _patch_dispatch(["v"], [[1]]):
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+        )
+
+    all_calls = mock_ctx.info.await_args_list + mock_ctx.error.await_args_list
+    for call in all_calls:
+        assert secret_password not in str(call)
+
+
+# ── Unit tests for helpers ────────────────────────────────────────────────────
+
+def test_cell_str_none() -> None:
+    assert _cell_str(None) == ""
+
+
+def test_cell_str_normal() -> None:
+    assert _cell_str(42) == "42"
+    assert _cell_str("hello") == "hello"
+
+
+def test_cell_str_truncated() -> None:
+    long_val = "x" * 10_000
+    with patch.object(settings, "db_max_cell_bytes", 10):
+        result = _cell_str(long_val)
+    assert "…" in result
+    assert len(result) < 20
+
+
+def test_format_result_no_rows() -> None:
+    result = _format_result("host", "db", "postgres", "SELECT 1", [], [], False, 5.0)
+    assert "No rows returned" in result
+
+
+def test_format_result_with_rows() -> None:
+    cols = ["id", "name"]
+    rows = [[1, "Alice"], [2, "Bob"]]
+    result = _format_result("host", "db", "postgres", "SELECT ...", cols, rows, False, 12.3)
+    assert "id" in result
+    assert "Alice" in result
+    assert "Bob" in result
+    assert "Rows returned: 2" in result
+
+
+def test_format_result_truncated_flag() -> None:
+    cols = ["id"]
+    rows = [[i] for i in range(5)]
+    result = _format_result("host", "db", "postgres", "SELECT ...", cols, rows, True, 1.0)
+    assert "capped" in result

tests/test_integration.py

@@ -0,0 +1,330 @@
+"""
+Integration tests — end-to-end flows across multiple MCP tools.
+
+These tests verify that the FULL PIPELINE works:
+  CyberArk MCP  →  (handle)  →  SSH / PowerShell / DB MCP
+
+They also serve as a learning resource for how MCP tools compose:
+
+  ┌─────────────────────────────────────────────────────────────────┐
+  │  Claude (LLM)                                                   │
+  │    1. Calls get_credential(safe, object_name)                   │
+  │       → receives "secret://abc123..."  (handle only)            │
+  │    2. Calls ssh_execute(host, command, secret_handle=handle)    │
+  │       → receives command output                                 │
+  └─────────────────────────────────────────────────────────────────┘
+
+At no point does Claude see the actual password.
+The handle is an opaque token that binds a short TTL credential to one use.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from mcp_privileged.cyberark.client import CyberArkCCPClient, Credential
+from mcp_privileged.cyberark.server import get_credential
+from mcp_privileged.database.server import db_query
+from mcp_privileged.powershell.server import ps_execute
+from mcp_privileged.secret_store import secret_store
+from mcp_privileged.ssh.server import ssh_execute
+from tests.conftest import make_db_result, make_ps_result, make_ssh_cm
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _make_ctx(client_ip: str = "10.0.0.1") -> MagicMock:
+    ctx = MagicMock()
+    ctx.info = AsyncMock()
+    ctx.error = AsyncMock()
+    ctx.request_context.request.headers = {"X-Forwarded-For": client_ip}
+    ctx.request_context.request.client = None
+    return ctx
+
+
+def _mock_cyberark_client(username: str, password: str, address: str = "db.internal"):
+    """Patch the CyberArk CCP client to return a fixed credential."""
+    cred = Credential(
+        username=username,
+        password=password,
+        address=address,
+        safe="PROD-SAFE",
+        folder="Root",
+        object_name="PROD-DB-svc",
+        platform_id="UnixSSH",
+        password_change_in_process=False,
+    )
+    mock_client = MagicMock(spec=CyberArkCCPClient)
+    mock_client.get_credential = AsyncMock(return_value=cred)
+    mock_client._settings_app_id = lambda: "MCP-Privileged-Service"
+    return patch("mcp_privileged.cyberark.server.cyberark_client", mock_client)
+
+
+# ── Full pipeline: CyberArk → SSH ─────────────────────────────────────────────
+
+async def test_cyberark_to_ssh_full_pipeline() -> None:
+    """
+    Simulate the complete CyberArk → SSH pipeline:
+
+    1. get_credential() fetches from CyberArk, stores in secret_store, returns handle.
+    2. ssh_execute() resolves the handle, uses the password to connect, returns output.
+    3. The password never appears in either tool's return value.
+
+    This is the primary privileged-access use case:
+        Claude: "Run `df -h` on linux01 using the PROD-LINUX credential"
+    """
+    ctx_cyberark = _make_ctx("192.168.1.10")
+    ctx_ssh = _make_ctx("192.168.1.10")
+
+    # Step 1: Claude calls get_credential
+    with _mock_cyberark_client(username="root", password="SshSecret!"):
+        handle_response = await get_credential(
+            safe="PROD-SAFE",
+            object_name="PROD-LINUX-root",
+            ctx=ctx_cyberark,
+        )
+
+    # The LLM receives a handle string — NOT the password
+    assert "secret://" in handle_response
+    assert "SshSecret!" not in handle_response
+
+    # Extract the handle token from the formatted response text
+    handle = next(
+        line.split("Handle: ")[1]
+        for line in handle_response.splitlines()
+        if line.startswith("Handle: ")
+    )
+
+    # Step 2: Claude calls ssh_execute with the handle
+    mock_cm, _ = make_ssh_cm(stdout="/dev/sda1  50G  10G  40G  20% /\n", exit_status=0)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        ssh_result = await ssh_execute(
+            host="linux01.internal",
+            command="df -h",
+            secret_handle=handle,
+            ctx=ctx_ssh,
+        )
+
+    assert "Exit code: 0" in ssh_result
+    assert "/dev/sda1" in ssh_result
+    assert "SshSecret!" not in ssh_result
+
+
+async def test_cyberark_to_powershell_full_pipeline() -> None:
+    """Simulate CyberArk → PowerShell pipeline."""
+    ctx_ca = _make_ctx()
+    ctx_ps = _make_ctx()
+
+    with _mock_cyberark_client(username="domain\\svc_ps", password="WinSecret!"):
+        handle_response = await get_credential(
+            safe="WIN-SAFE",
+            object_name="WIN-svc_ps",
+            ctx=ctx_ca,
+        )
+
+    assert "WinSecret!" not in handle_response
+    handle = next(
+        line.split("Handle: ")[1]
+        for line in handle_response.splitlines()
+        if line.startswith("Handle: ")
+    )
+
+    ps_result = make_ps_result(output=["WIN-SERVER-01"], had_errors=False)
+
+    with patch("mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result):
+        ps_out = await ps_execute(
+            host="win01.internal",
+            script="hostname",
+            secret_handle=handle,
+            ctx=ctx_ps,
+        )
+
+    assert "Had errors: False" in ps_out
+    assert "WIN-SERVER-01" in ps_out
+    assert "WinSecret!" not in ps_out
+
+
+async def test_cyberark_to_database_full_pipeline() -> None:
+    """Simulate CyberArk → Database pipeline."""
+    ctx_ca = _make_ctx()
+    ctx_db = _make_ctx()
+
+    with _mock_cyberark_client(username="db_reader", password="DbSecret!"):
+        handle_response = await get_credential(
+            safe="DB-SAFE",
+            object_name="PROD-PG-reader",
+            ctx=ctx_ca,
+        )
+
+    assert "DbSecret!" not in handle_response
+    handle = next(
+        line.split("Handle: ")[1]
+        for line in handle_response.splitlines()
+        if line.startswith("Handle: ")
+    )
+
+    with patch(
+        "mcp_privileged.database.server._dispatch_query",
+        new=AsyncMock(return_value=make_db_result(["count"], [[42]])),
+    ):
+        db_out = await db_query(
+            host="pg.internal",
+            database="prod",
+            query="SELECT COUNT(*) FROM users",
+            secret_handle=handle,
+            ctx=ctx_db,
+            db_type="postgres",
+        )
+
+    assert "42" in db_out
+    assert "DbSecret!" not in db_out
+
+
+# ── Handle lifecycle ──────────────────────────────────────────────────────────
+
+async def test_handle_single_use_enforced() -> None:
+    """
+    A handle issued by get_credential can only be resolved ONCE
+    (when handle_single_use=True, which is the default).
+
+    This prevents credential replay attacks:
+    if an attacker intercepts the handle, it's already been consumed.
+    """
+    ctx = _make_ctx()
+    mock_cm, _ = make_ssh_cm(stdout="ok\n", exit_status=0)
+
+    with _mock_cyberark_client(username="user", password="pass"):
+        handle_response = await get_credential(
+            safe="S", object_name="O", ctx=ctx
+        )
+
+    handle = next(
+        line.split("Handle: ")[1]
+        for line in handle_response.splitlines()
+        if line.startswith("Handle: ")
+    )
+
+    # First use — succeeds
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        await ssh_execute(
+            host="host1", command="id", secret_handle=handle, ctx=ctx
+        )
+
+    # Second use — same handle, should fail
+    with pytest.raises(KeyError, match="consumed|not found"):
+        await ssh_execute(
+            host="host1", command="id", secret_handle=handle, ctx=ctx
+        )
+
+
+async def test_handle_cannot_be_shared_across_tools() -> None:
+    """
+    A handle resolved by ssh_execute cannot then be reused by db_query.
+    One credential fetch = one privileged operation.
+    """
+    ctx = _make_ctx()
+    mock_cm, _ = make_ssh_cm(stdout="ok\n", exit_status=0)
+
+    # Issue one handle
+    handle = await secret_store.store("user", "pass")
+
+    # SSH consumes it
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        await ssh_execute(
+            host="host1", command="id", secret_handle=handle, ctx=ctx
+        )
+
+    # DB tries to reuse it — must fail
+    with pytest.raises(KeyError):
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=ctx,
+            db_type="postgres",
+        )
+
+
+async def test_expired_handle_rejected() -> None:
+    """
+    A handle past its TTL is rejected even if not yet consumed.
+    We simulate expiry by manually backdating the entry's created_at.
+    """
+    import time
+
+    handle = await secret_store.store("user", "pass")
+    handle_id = handle.split("://")[1]
+
+    # Backdate the entry so it looks expired
+    async with secret_store._lock:
+        entry = secret_store._store[handle_id]
+        entry.created_at = time.monotonic() - 99999  # very old
+
+    with pytest.raises(KeyError, match="expired"):
+        await secret_store.resolve(handle, resolved_by="test")
+
+
+# ── Concurrent handle isolation ───────────────────────────────────────────────
+
+async def test_concurrent_handles_are_independent() -> None:
+    """
+    Multiple handles issued at the same time are independent.
+    Resolving one does not affect the others.
+    """
+    handles = [await secret_store.store(f"user_{i}", f"pass_{i}") for i in range(5)]
+
+    # Resolve them in reverse order
+    results = []
+    for handle in reversed(handles):
+        username, password = await secret_store.resolve(handle, resolved_by="test")
+        results.append((username, password))
+
+    assert len(results) == 5
+    # Each (username, password) pair is unique
+    assert len(set(results)) == 5
+
+
+# ── Audit trail ───────────────────────────────────────────────────────────────
+
+async def test_audit_events_fired_for_ssh(mock_ctx) -> None:
+    """
+    ssh_execute must call ctx.info() at least twice:
+    once for connection start, once for completion.
+    ctx.error must NOT be called on the happy path.
+    """
+    handle = await secret_store.store("user", "pass")
+    mock_cm, _ = make_ssh_cm(stdout="ok\n", exit_status=0)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        await ssh_execute(
+            host="host1", command="id", secret_handle=handle, ctx=mock_ctx
+        )
+
+    assert mock_ctx.info.await_count >= 2
+    mock_ctx.error.assert_not_awaited()
+
+
+async def test_audit_events_fired_for_db(mock_ctx) -> None:
+    """db_query must emit ctx.info on the happy path, not ctx.error."""
+    handle = await secret_store.store("user", "pass")
+
+    with patch(
+        "mcp_privileged.database.server._dispatch_query",
+        new=AsyncMock(return_value=make_db_result(["v"], [[1]])),
+    ):
+        await db_query(
+            host="pg.internal",
+            database="mydb",
+            query="SELECT 1",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            db_type="postgres",
+        )
+
+    assert mock_ctx.info.await_count >= 2
+    mock_ctx.error.assert_not_awaited()

tests/test_powershell_server.py

@@ -0,0 +1,227 @@
+"""
+Tests for the PowerShell MCP tool (ps_execute).
+
+pypsrp is a synchronous library.  The server wraps _run_ps_sync() in
+asyncio.run_in_executor so we patch _run_ps_sync directly — no real WinRM
+connections are made.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+from mcp_privileged.powershell.server import _format_result, _truncate, ps_execute
+from mcp_privileged.secret_store import secret_store
+from tests.conftest import make_ps_result
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+async def _handle(username: str = "svc_user", password: str = "P@ss!") -> str:
+    return await secret_store.store(username, password)
+
+
+# ── Tests ─────────────────────────────────────────────────────────────────────
+
+async def test_ps_execute_success(mock_ctx) -> None:
+    """Happy path: script runs, output is returned, had_errors=False."""
+    handle = await _handle()
+    ps_result = make_ps_result(output=["Win2022", "Server"], had_errors=False)
+
+    with patch("mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result):
+        result = await ps_execute(
+            host="win01.internal",
+            script="$PSVersionTable.OS; hostname",
+            secret_handle=handle,
+            ctx=mock_ctx,
+        )
+
+    assert "Had errors: False" in result
+    assert "Win2022" in result
+    assert "Server" in result
+    assert "Host: win01.internal" in result
+
+
+async def test_ps_execute_with_errors(mock_ctx) -> None:
+    """Script produces errors — had_errors=True and error records are included."""
+    handle = await _handle()
+    ps_result = make_ps_result(
+        output=[],
+        had_errors=True,
+        errors=["Get-Item : Cannot find path 'C:\\missing'"],
+    )
+
+    with patch("mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result):
+        result = await ps_execute(
+            host="win01.internal",
+            script="Get-Item C:\\missing",
+            secret_handle=handle,
+            ctx=mock_ctx,
+        )
+
+    assert "Had errors: True" in result
+    assert "Cannot find path" in result
+    assert "--- errors ---" in result
+
+
+async def test_ps_execute_no_output(mock_ctx) -> None:
+    """Script runs but produces no output (e.g. Set-* cmdlets)."""
+    handle = await _handle()
+    ps_result = make_ps_result(output=[], had_errors=False)
+
+    with patch("mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result):
+        result = await ps_execute(
+            host="win01.internal",
+            script="Set-TimeZone -Id 'UTC'",
+            secret_handle=handle,
+            ctx=mock_ctx,
+        )
+
+    assert "Had errors: False" in result
+    assert "(no output)" in result
+
+
+async def test_ps_execute_username_override(mock_ctx) -> None:
+    """username_override is forwarded to _run_ps_sync."""
+    handle = await _handle(username="domain\\original")
+    ps_result = make_ps_result(output=["ok"])
+
+    with patch(
+        "mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result
+    ) as mock_run:
+        await ps_execute(
+            host="win01.internal",
+            script="whoami",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            username_override="domain\\admin",
+        )
+
+    # Third positional arg to _run_ps_sync is username
+    _args, _ = mock_run.call_args
+    assert _args[2] == "domain\\admin"
+
+
+async def test_ps_execute_credential_username_used_by_default(mock_ctx) -> None:
+    """Without username_override, the credential username is forwarded."""
+    handle = await _handle(username="domain\\svc_ps")
+    ps_result = make_ps_result(output=["ok"])
+
+    with patch(
+        "mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result
+    ) as mock_run:
+        await ps_execute(
+            host="win01.internal",
+            script="whoami",
+            secret_handle=handle,
+            ctx=mock_ctx,
+        )
+
+    _args, _ = mock_run.call_args
+    assert _args[2] == "domain\\svc_ps"
+
+
+async def test_ps_execute_invalid_handle(mock_ctx) -> None:
+    """Unknown handle raises KeyError before any WinRM connection is attempted."""
+    with pytest.raises(KeyError):
+        await ps_execute(
+            host="win01.internal",
+            script="hostname",
+            secret_handle="secret://doesnotexist0000000000000000",
+            ctx=mock_ctx,
+        )
+
+    mock_ctx.error.assert_awaited_once()
+
+
+async def test_ps_execute_winrm_exception_propagates(mock_ctx) -> None:
+    """Exceptions from _run_ps_sync propagate and call ctx.error."""
+    handle = await _handle()
+
+    with patch(
+        "mcp_privileged.powershell.server._run_ps_sync",
+        side_effect=ConnectionRefusedError("WinRM port closed"),
+    ):
+        with pytest.raises(ConnectionRefusedError):
+            await ps_execute(
+                host="dead.host",
+                script="hostname",
+                secret_handle=handle,
+                ctx=mock_ctx,
+            )
+
+    mock_ctx.error.assert_awaited_once()
+
+
+async def test_ps_execute_password_not_in_ctx_messages(mock_ctx) -> None:
+    """The password must never appear in any ctx.info or ctx.error call."""
+    secret_password = "WinRM$ecret99"
+    handle = await secret_store.store("user", secret_password)
+    ps_result = make_ps_result(output=["ok"])
+
+    with patch("mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result):
+        await ps_execute(
+            host="win01.internal",
+            script="hostname",
+            secret_handle=handle,
+            ctx=mock_ctx,
+        )
+
+    all_calls = mock_ctx.info.await_args_list + mock_ctx.error.await_args_list
+    for call in all_calls:
+        assert secret_password not in str(call), "Password leaked into MCP context log"
+
+
+async def test_ps_execute_ssl_and_port_forwarded(mock_ctx) -> None:
+    """use_ssl=True and custom port are forwarded to _run_ps_sync."""
+    handle = await _handle()
+    ps_result = make_ps_result(output=["ok"])
+
+    with patch(
+        "mcp_privileged.powershell.server._run_ps_sync", return_value=ps_result
+    ) as mock_run:
+        await ps_execute(
+            host="win01.internal",
+            script="hostname",
+            secret_handle=handle,
+            ctx=mock_ctx,
+            port=5986,
+            use_ssl=True,
+        )
+
+    _args, _ = mock_run.call_args
+    assert _args[1] == 5986   # port
+    assert _args[5] is True   # use_ssl
+
+
+# ── Unit tests for helpers ─────────────────────────────────────────────────────
+
+def test_truncate_passthrough() -> None:
+    assert _truncate("hello", 1024, "output") == "hello"
+
+
+def test_truncate_applies_limit() -> None:
+    result = _truncate("x" * 10_000, 100, "output")
+    assert "truncated" in result
+    assert len(result.encode()) < 300
+
+
+def test_format_result_no_errors() -> None:
+    result = _format_result("win01", "Get-Process", False, ["proc1", "proc2"], [])
+    assert "Had errors: False" in result
+    assert "proc1" in result
+    assert "--- errors ---" not in result
+
+
+def test_format_result_with_errors() -> None:
+    result = _format_result("win01", "bad_cmd", True, [], ["Error: not found"])
+    assert "Had errors: True" in result
+    assert "--- errors ---" in result
+    assert "not found" in result
+
+
+def test_format_result_empty_output() -> None:
+    result = _format_result("win01", "Set-X", False, [], [])
+    assert "(no output)" in result

tests/test_secret_store.py

@@ -0,0 +1,100 @@
+"""
+Tests for the secret handle store.
+Covers: store, resolve, single-use, TTL expiry, revoke, and sweeper.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+import pytest
+
+from mcp_privileged.secret_store import SecretStore, HANDLE_PREFIX
+
+
+@pytest.fixture
+def store() -> SecretStore:
+    return SecretStore()
+
+
+async def test_store_returns_handle(store: SecretStore) -> None:
+    handle = await store.store("user1", "s3cr3t")
+    assert handle.startswith(HANDLE_PREFIX)
+
+
+async def test_resolve_returns_credentials(store: SecretStore) -> None:
+    handle = await store.store("user1", "s3cr3t")
+    username, password = await store.resolve(handle, resolved_by="test")
+    assert username == "user1"
+    assert password == "s3cr3t"
+
+
+async def test_single_use_invalidates_after_first_resolve(
+    store: SecretStore, monkeypatch
+) -> None:
+    monkeypatch.setattr("mcp_privileged.secret_store.settings.handle_single_use", True)
+    handle = await store.store("user1", "s3cr3t")
+    await store.resolve(handle, resolved_by="test")
+    with pytest.raises(KeyError, match="already_consumed|not found"):
+        await store.resolve(handle, resolved_by="test")
+
+
+async def test_multi_use_allows_repeated_resolve(
+    store: SecretStore, monkeypatch
+) -> None:
+    monkeypatch.setattr("mcp_privileged.secret_store.settings.handle_single_use", False)
+    handle = await store.store("user1", "s3cr3t")
+    for _ in range(3):
+        username, password = await store.resolve(handle, resolved_by="test")
+        assert password == "s3cr3t"
+
+
+async def test_expired_handle_raises(store: SecretStore, monkeypatch) -> None:
+    monkeypatch.setattr("mcp_privileged.secret_store.settings.handle_ttl_seconds", 1)
+    handle = await store.store("user1", "s3cr3t")
+    # Manually backdate the entry's creation time
+    handle_id = handle[len(HANDLE_PREFIX):]
+    store._store[handle_id].created_at = time.monotonic() - 5
+    with pytest.raises(KeyError, match="expired"):
+        await store.resolve(handle, resolved_by="test")
+
+
+async def test_unknown_handle_raises(store: SecretStore) -> None:
+    with pytest.raises(KeyError):
+        await store.resolve(f"{HANDLE_PREFIX}nonexistent", resolved_by="test")
+
+
+async def test_invalid_handle_format_raises(store: SecretStore) -> None:
+    with pytest.raises(ValueError, match="Invalid handle format"):
+        await store.resolve("not-a-handle", resolved_by="test")
+
+
+async def test_revoke_removes_handle(store: SecretStore) -> None:
+    handle = await store.store("user1", "s3cr3t")
+    assert await store.revoke(handle) is True
+    with pytest.raises(KeyError):
+        await store.resolve(handle, resolved_by="test")
+
+
+async def test_revoke_nonexistent_returns_false(store: SecretStore) -> None:
+    assert await store.revoke(f"{HANDLE_PREFIX}nonexistent") is False
+
+
+async def test_purge_expired_removes_stale(store: SecretStore, monkeypatch) -> None:
+    monkeypatch.setattr("mcp_privileged.secret_store.settings.handle_ttl_seconds", 1)
+    handle = await store.store("user1", "s3cr3t")
+    handle_id = handle[len(HANDLE_PREFIX):]
+    store._store[handle_id].created_at = time.monotonic() - 5
+    count = await store.purge_expired()
+    assert count == 1
+    assert handle_id not in store._store
+
+
+async def test_password_not_in_repr(store: SecretStore) -> None:
+    """SecretStr must not leak the password in string representations."""
+    handle = await store.store("user1", "topsecret")
+    handle_id = handle[len(HANDLE_PREFIX):]
+    entry = store._store[handle_id]
+    assert "topsecret" not in repr(entry)
+    assert "topsecret" not in str(entry.password)

tests/test_ssh_server.py

@@ -0,0 +1,291 @@
+"""
+Tests for the SSH MCP tool (ssh_execute).
+
+All tests mock asyncssh.connect — no real SSH connections are made.
+The secret_store is used directly so handle issuance/resolution is tested
+end-to-end through the real store.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import asyncssh
+import pytest
+
+from mcp_privileged.secret_store import secret_store
+from mcp_privileged.ssh.server import _truncate, _format_result, ssh_execute
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _make_ctx() -> MagicMock:
+    """Return a minimal mock MCP Context."""
+    ctx = MagicMock()
+    ctx.info = AsyncMock()
+    ctx.error = AsyncMock()
+    # _extract_client_ip uses these
+    ctx.request_context.request.headers = {}
+    ctx.request_context.request.client = None
+    return ctx
+
+
+def _make_ssh_cm(
+    stdout: str = "",
+    stderr: str = "",
+    exit_status: int = 0,
+) -> tuple[AsyncMock, AsyncMock]:
+    """
+    Build a mock for asyncssh.connect used as an async context manager.
+
+    Returns (context_manager_mock, conn_mock).
+    Patch asyncssh.connect with return_value=context_manager_mock.
+    """
+    mock_conn = AsyncMock()
+    mock_conn.run = AsyncMock(
+        return_value=MagicMock(stdout=stdout, stderr=stderr, exit_status=exit_status)
+    )
+    mock_cm = AsyncMock()
+    mock_cm.__aenter__ = AsyncMock(return_value=mock_conn)
+    mock_cm.__aexit__ = AsyncMock(return_value=False)
+    return mock_cm, mock_conn
+
+
+async def _fresh_handle(username: str = "svc_user", password: str = "P@ssw0rd!") -> str:
+    """Store a credential and return a fresh (unconsumed) handle."""
+    return await secret_store.store(username, password)
+
+
+# ── Tests ─────────────────────────────────────────────────────────────────────
+
+async def test_ssh_execute_success() -> None:
+    """Happy path: command runs, stdout is returned, exit code is 0."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+    mock_cm, _ = _make_ssh_cm(stdout="hello world\n", exit_status=0)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        result = await ssh_execute(
+            host="linux01.internal",
+            command="echo hello world",
+            secret_handle=handle,
+            ctx=ctx,
+        )
+
+    assert "Exit code: 0" in result
+    assert "hello world" in result
+    assert "Host: linux01.internal" in result
+    assert "Command: echo hello world" in result
+
+
+async def test_ssh_execute_nonzero_exit_not_raised() -> None:
+    """A non-zero exit code is returned in the result, not raised as an exception."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+    mock_cm, _ = _make_ssh_cm(stdout="", stderr="command not found\n", exit_status=127)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        result = await ssh_execute(
+            host="linux01.internal",
+            command="notacommand",
+            secret_handle=handle,
+            ctx=ctx,
+        )
+
+    assert "Exit code: 127" in result
+    assert "command not found" in result
+
+
+async def test_ssh_execute_stderr_included() -> None:
+    """Both stdout and stderr appear in the result when both are non-empty."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+    mock_cm, _ = _make_ssh_cm(stdout="result\n", stderr="warning: low disk\n", exit_status=0)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        result = await ssh_execute(
+            host="host1",
+            command="df -h",
+            secret_handle=handle,
+            ctx=ctx,
+        )
+
+    assert "result" in result
+    assert "warning: low disk" in result
+
+
+async def test_ssh_execute_username_override() -> None:
+    """username_override replaces the credential's username in the connect call."""
+    handle = await _fresh_handle(username="original_user")
+    ctx = _make_ctx()
+    mock_cm, _ = _make_ssh_cm(stdout="uid=0(root)\n", exit_status=0)
+
+    with patch(
+        "mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm
+    ) as mock_connect:
+        await ssh_execute(
+            host="host1",
+            command="id",
+            secret_handle=handle,
+            ctx=ctx,
+            username_override="root",
+        )
+
+    _args, _kwargs = mock_connect.call_args
+    assert _kwargs["username"] == "root"
+
+
+async def test_ssh_execute_credential_username_used_by_default() -> None:
+    """Without username_override, the credential's username is passed to connect."""
+    handle = await _fresh_handle(username="db_admin")
+    ctx = _make_ctx()
+    mock_cm, _ = _make_ssh_cm(stdout="ok\n", exit_status=0)
+
+    with patch(
+        "mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm
+    ) as mock_connect:
+        await ssh_execute(
+            host="host1",
+            command="whoami",
+            secret_handle=handle,
+            ctx=ctx,
+        )
+
+    _args, _kwargs = mock_connect.call_args
+    assert _kwargs["username"] == "db_admin"
+
+
+async def test_ssh_execute_invalid_handle_raises() -> None:
+    """An unknown handle raises KeyError and calls ctx.error."""
+    ctx = _make_ctx()
+
+    with pytest.raises(KeyError):
+        await ssh_execute(
+            host="host1",
+            command="id",
+            secret_handle="secret://doesnotexist0000000000000000",
+            ctx=ctx,
+        )
+
+    ctx.error.assert_awaited_once()
+
+
+async def test_ssh_execute_connect_os_error_propagates() -> None:
+    """An OSError (e.g. connection refused) propagates and calls ctx.error."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+
+    with patch(
+        "mcp_privileged.ssh.server.asyncssh.connect",
+        side_effect=OSError("Connection refused"),
+    ):
+        with pytest.raises(OSError):
+            await ssh_execute(
+                host="dead.host",
+                command="id",
+                secret_handle=handle,
+                ctx=ctx,
+            )
+
+    ctx.error.assert_awaited_once()
+
+
+async def test_ssh_execute_permission_denied_propagates() -> None:
+    """asyncssh.PermissionDenied propagates and calls ctx.error."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+
+    with patch(
+        "mcp_privileged.ssh.server.asyncssh.connect",
+        side_effect=asyncssh.PermissionDenied("Permission denied"),
+    ):
+        with pytest.raises(asyncssh.PermissionDenied):
+            await ssh_execute(
+                host="host1",
+                command="id",
+                secret_handle=handle,
+                ctx=ctx,
+            )
+
+    ctx.error.assert_awaited_once()
+
+
+async def test_ssh_execute_command_timeout_propagates() -> None:
+    """asyncio.TimeoutError from conn.run propagates and calls ctx.error."""
+    handle = await _fresh_handle()
+    ctx = _make_ctx()
+
+    mock_conn = AsyncMock()
+    mock_conn.run = AsyncMock(side_effect=asyncio.TimeoutError())
+    mock_cm = AsyncMock()
+    mock_cm.__aenter__ = AsyncMock(return_value=mock_conn)
+    mock_cm.__aexit__ = AsyncMock(return_value=False)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        with pytest.raises(asyncio.TimeoutError):
+            await ssh_execute(
+                host="slow.host",
+                command="sleep 999",
+                secret_handle=handle,
+                ctx=ctx,
+                timeout_seconds=1,
+            )
+
+    ctx.error.assert_awaited_once()
+
+
+async def test_ssh_execute_password_not_in_result() -> None:
+    """The credential password must never appear in the tool's return value."""
+    secret_password = "SuperSecret!123"
+    handle = await _fresh_handle(password=secret_password)
+    ctx = _make_ctx()
+    # Simulate a misconfigured command that echoes env vars containing the password
+    mock_cm, _ = _make_ssh_cm(stdout=f"PASSWORD={secret_password}\n", exit_status=0)
+
+    with patch("mcp_privileged.ssh.server.asyncssh.connect", return_value=mock_cm):
+        result = await ssh_execute(
+            host="host1",
+            command="env",
+            secret_handle=handle,
+            ctx=ctx,
+        )
+
+    # The password leaking from stdout is the application's problem, not ours —
+    # what we must guarantee is that the *handle resolution* never injects it.
+    # Verify it doesn't appear in any ctx.error/ctx.info call from our code:
+    for call in ctx.error.await_args_list + ctx.info.await_args_list:
+        assert secret_password not in str(call), "Password leaked into MCP context log"
+
+
+# ── Unit tests for helpers ────────────────────────────────────────────────────
+
+def test_truncate_short_text_unchanged() -> None:
+    text = "hello world"
+    assert _truncate(text, 1024, "stdout") == text
+
+
+def test_truncate_long_text_truncated() -> None:
+    text = "x" * 10_000
+    result = _truncate(text, 100, "stdout")
+    assert "truncated" in result
+    assert len(result.encode("utf-8")) <= 200   # marker adds a short suffix
+
+
+def test_format_result_no_stderr() -> None:
+    result = _format_result("myhost", "ls /", 0, "bin\nlib\n", "")
+    assert "--- stderr ---" not in result
+    assert "Exit code: 0" in result
+    assert "bin" in result
+
+
+def test_format_result_with_stderr() -> None:
+    result = _format_result("myhost", "bad_cmd", 1, "", "not found\n")
+    assert "--- stderr ---" in result
+    assert "not found" in result
+    assert "Exit code: 1" in result
+
+
+def test_format_result_empty_stdout_shows_empty_marker() -> None:
+    result = _format_result("myhost", "true", 0, "", "")
+    assert "(empty)" in result