feat: add SSH agent authentication

Author: O2eg

README.md

@@ -147,7 +147,7 @@ The deterministic precedence is `explicit CLI > CH_DIAG_* environment > TOML
 lists and can then be repeated. Supported environment names follow CLI names,
 for example `CH_DIAG_HOST`, `CH_DIAG_PORT`, `CH_DIAG_COLLECTION_MODE`,
 `CH_DIAG_TARGET_SCOPE`, `CH_DIAG_OUTPUT_FORMAT`, `CH_DIAG_SSH_HOST` and
-`CH_DIAG_DURATION`. Run `ch-diag --help` and the command help for the full CLI.
+`CH_DIAG_SSH_AGENT`. Run `ch-diag --help` and the command help for the full CLI.
 
 Literal passwords are deliberately rejected in TOML. Put only
 `password_env = "CH_DIAG_PASSWORD"` in the file and supply the secret through
@@ -324,8 +324,10 @@ Node is the default and does not require a configured cluster.
 
 ## SSH mode and known hosts
 
-SSH is key-only. `--ssh-known-hosts` is mandatory and host-key verification is
-never silently disabled. Prepare a dedicated file before the first run:
+SSH uses explicit public-key authentication through either a key file or an
+existing local agent. `--ssh-known-hosts` is mandatory and host-key
+verification is never silently disabled. Prepare a dedicated file before the
+first run:
 
 ```bash
 mkdir -p ~/.ssh
@@ -350,6 +352,31 @@ ch-diag snapshots \
   --out-dir reports/remote
 ```
 
+Alternatively, load the identity into a local agent and replace `--ssh-key`
+with `--ssh-agent`:
+
+```bash
+eval "$(ssh-agent -s)"
+ssh-add ~/.ssh/ch_diag_ed25519
+ssh-add -l
+
+ch-diag snapshots \
+  --host 127.0.0.1 --port 9000 \
+  --collection-mode remote --target-scope node \
+  --ssh-host db-host.example.net --ssh-port 22 \
+  --ssh-user chdiag --ssh-agent \
+  --ssh-known-hosts ~/.ssh/ch_diag_known_hosts \
+  --duration 10 --interval 5 \
+  --out-dir reports/remote
+```
+
+`--ssh-key` and `--ssh-agent` are mutually exclusive. Agent mode requires a
+live Unix socket in the inherited `SSH_AUTH_SOCK` for the complete collection.
+`ch-diag` does not start an agent, run `ssh-add`, or forward the agent to the
+target. Key mode explicitly disables fallback to identities from an agent.
+TOML may select agent mode with `agent = true` in `[ssh]`; the corresponding
+environment switch is `CH_DIAG_SSH_AGENT=true`.
+
 `--host` and `--port` in this mode are resolved from the SSH target. The
 ClickHouse native connection travels through the tunnel; shell scripts execute
 on the SSH target.

ch_diag/__init__.py

@@ -1,3 +1,3 @@
 """ch_diag package."""
 
-__version__ = "0.8.0"
+__version__ = "0.8.1"

ch_diag/cli.py

@@ -55,6 +55,39 @@ def __call__(
         setattr(namespace, self.dest, current)
 
 
+class _SshKeyAction(argparse.Action):
+    """Select key-file auth and override an agent default from config or env."""
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: str,
+        option_string: str | None = None,
+    ) -> None:
+        del parser, option_string
+        namespace.ssh_key = values
+        namespace.ssh_agent = False
+
+
+class _SshAgentAction(argparse.Action):
+    """Select agent auth and override a key default from config or env."""
+
+    def __init__(self, option_strings: Sequence[str], dest: str, **kwargs: Any) -> None:
+        super().__init__(option_strings, dest, nargs=0, **kwargs)
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Any,
+        option_string: str | None = None,
+    ) -> None:
+        del parser, values, option_string
+        namespace.ssh_agent = True
+        namespace.ssh_key = None
+
+
 def _parser(defaults: dict[str, Any] | None = None) -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(prog="ch-diag", description="ClickHouse diagnostics")
     parser.add_argument("--config", help="TOML configuration file (must precede command)")
@@ -152,7 +185,14 @@ def _add_collection_arguments(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--ssh-host")
     parser.add_argument("--ssh-port", type=int, default=22)
     parser.add_argument("--ssh-user")
-    parser.add_argument("--ssh-key")
+    ssh_auth = parser.add_mutually_exclusive_group()
+    ssh_auth.add_argument("--ssh-key", action=_SshKeyAction)
+    ssh_auth.add_argument(
+        "--ssh-agent",
+        action=_SshAgentAction,
+        default=False,
+        help="authenticate through the existing SSH_AUTH_SOCK agent",
+    )
     parser.add_argument("--ssh-known-hosts")
 
 
@@ -200,18 +240,35 @@ def _connection(args: argparse.Namespace) -> ConnectionConfig:
 
 
 def _ssh_config(args: argparse.Namespace) -> SshConfig | None:
-    values = (args.ssh_host, args.ssh_user, args.ssh_key, args.ssh_known_hosts)
+    values = (
+        args.ssh_host,
+        args.ssh_user,
+        args.ssh_key,
+        args.ssh_agent,
+        args.ssh_known_hosts,
+    )
     if not any(values):
         return None
-    if not all(values):
+    if not args.ssh_host or not args.ssh_user or not args.ssh_known_hosts:
         raise ValueError(
-            "SSH requires --ssh-host, --ssh-user, --ssh-key and --ssh-known-hosts"
+            "SSH requires --ssh-host, --ssh-user, --ssh-known-hosts and "
+            "exactly one of --ssh-key or --ssh-agent"
         )
+    if bool(args.ssh_key) == bool(args.ssh_agent):
+        raise ValueError("SSH requires exactly one of --ssh-key or --ssh-agent")
+    agent_path = None
+    if args.ssh_agent:
+        agent_path = os.environ.get("SSH_AUTH_SOCK")
+        if not agent_path:
+            raise ValueError(
+                "--ssh-agent requires SSH_AUTH_SOCK to reference a running agent"
+            )
     return SshConfig(
         host=args.ssh_host,
         port=args.ssh_port,
         username=args.ssh_user,
         client_key=args.ssh_key,
+        agent_path=agent_path,
         known_hosts=args.ssh_known_hosts,
     )
 

ch_diag/config.py

@@ -42,6 +42,7 @@
         "port": "ssh_port",
         "user": "ssh_user",
         "key": "ssh_key",
+        "agent": "ssh_agent",
         "known_hosts": "ssh_known_hosts",
     },
     "snapshots": {"duration": "duration", "interval": "interval"},
@@ -75,12 +76,13 @@
     "CH_DIAG_SSH_PORT": "ssh_port",
     "CH_DIAG_SSH_USER": "ssh_user",
     "CH_DIAG_SSH_KEY": "ssh_key",
+    "CH_DIAG_SSH_AGENT": "ssh_agent",
     "CH_DIAG_SSH_KNOWN_HOSTS": "ssh_known_hosts",
     "CH_DIAG_DURATION": "duration",
     "CH_DIAG_INTERVAL": "interval",
 }
 
-BOOLEAN_FIELDS = {"secure", "no_verify", "strip_meta"}
+BOOLEAN_FIELDS = {"secure", "no_verify", "strip_meta", "ssh_agent"}
 INTEGER_FIELDS = {"port", "ssh_port"}
 FLOAT_FIELDS = {"duration", "interval"}
 LIST_FIELDS = {"item_ids", "tags", "output_format"}
@@ -142,6 +144,12 @@ def resolve_cli_defaults(
             raise ValueError(
                 "unknown configuration section(s): " + ", ".join(unknown_sections)
             )
+        ssh_values = document.get("ssh") or {}
+        if isinstance(ssh_values, dict) and ssh_values.get("key") and _boolean(
+            ssh_values.get("agent", False),
+            "configuration [ssh].agent",
+        ):
+            raise ValueError("configuration [ssh] must select either key or agent")
         for section, fields in CONFIG_FIELDS.items():
             values = document.get(section) or {}
             if not isinstance(values, dict):
@@ -160,6 +168,21 @@ def resolve_cli_defaults(
                 )
 
     env = environment if environment is not None else os.environ
+    env_key_present = "CH_DIAG_SSH_KEY" in env
+    env_agent_present = "CH_DIAG_SSH_AGENT" in env
+    env_agent_enabled = (
+        _boolean(env["CH_DIAG_SSH_AGENT"], "environment variable CH_DIAG_SSH_AGENT")
+        if env_agent_present
+        else False
+    )
+    if env_key_present and env.get("CH_DIAG_SSH_KEY") and env_agent_enabled:
+        raise ValueError(
+            "environment must select either CH_DIAG_SSH_KEY or CH_DIAG_SSH_AGENT"
+        )
+    if env_key_present:
+        defaults["ssh_agent"] = False
+    if env_agent_enabled:
+        defaults["ssh_key"] = None
     for name, destination in ENV_FIELDS.items():
         if name in env:
             defaults[destination] = _coerce(

ch_diag/ssh_transport.py

@@ -1,9 +1,10 @@
-"""Key-only SSH transport with mandatory known_hosts verification."""
+"""Explicit key-or-agent SSH transport with mandatory known_hosts verification."""
 
 from __future__ import annotations
 
 from dataclasses import dataclass
 from pathlib import Path
+import stat
 from typing import Any
 
 from .errors import ChDiagError
@@ -24,18 +25,32 @@ def _asyncssh():
 class SshConfig:
     host: str
     username: str
-    client_key: str
     known_hosts: str
+    client_key: str | None = None
+    agent_path: str | None = None
     port: int = 22
     connect_timeout: float = 10.0
 
     def validate(self) -> None:
         if not self.host or not self.username:
             raise ValueError("SSH host and username are required")
-        for label, value in (("SSH private key", self.client_key), ("known_hosts", self.known_hosts)):
-            path = Path(value).expanduser()
-            if not path.is_file():
-                raise ValueError(f"{label} file does not exist: {path}")
+        if bool(self.client_key) == bool(self.agent_path):
+            raise ValueError("SSH requires exactly one private key or agent socket")
+        if self.client_key:
+            key_path = Path(self.client_key).expanduser()
+            if not key_path.is_file():
+                raise ValueError(f"SSH private key file does not exist: {key_path}")
+        if self.agent_path:
+            agent_path = Path(self.agent_path).expanduser()
+            try:
+                mode = agent_path.stat().st_mode
+            except OSError as exc:
+                raise ValueError(f"SSH agent socket is not available: {agent_path}") from exc
+            if not stat.S_ISSOCK(mode):
+                raise ValueError(f"SSH_AUTH_SOCK is not a socket: {agent_path}")
+        known_hosts = Path(self.known_hosts).expanduser()
+        if not known_hosts.is_file():
+            raise ValueError(f"known_hosts file does not exist: {known_hosts}")
 
 
 class SshSession:
@@ -48,16 +63,25 @@ def __init__(self, config: SshConfig, connection: Any) -> None:
     async def connect(cls, config: SshConfig) -> "SshSession":
         config.validate()
         asyncssh = _asyncssh()
+        client_keys = (
+            [str(Path(config.client_key).expanduser())] if config.client_key else []
+        )
+        agent_path = (
+            str(Path(config.agent_path).expanduser()) if config.agent_path else None
+        )
         try:
             connection = await asyncssh.connect(
                 config.host,
                 port=config.port,
                 username=config.username,
-                client_keys=[str(Path(config.client_key).expanduser())],
+                client_keys=client_keys,
                 known_hosts=str(Path(config.known_hosts).expanduser()),
+                config=None,
+                preferred_auth=["publickey"],
                 password_auth=False,
                 kbdint_auth=False,
-                agent_path=None,
+                agent_path=agent_path,
+                agent_forwarding=False,
                 connect_timeout=config.connect_timeout,
             )
         except Exception as exc:

docs/security.md

@@ -11,7 +11,8 @@ it redacts, and what remains sensitive in a report.
   statement and runs with ClickHouse `readonly=2`, timeout, row and byte guards.
 - Cluster names are selected from `system.clusters` and quoted by the adapter.
 - Host scripts are package content, not user-supplied commands. Remote mode
-  requires key authentication and a verified `known_hosts` file.
+  requires explicit key-file or local-agent authentication and a verified
+  `known_hosts` file. Agent forwarding is always disabled.
 - The checksum detects corruption or replacement relative to its installed
   baseline. It cannot authenticate an attacker who can replace both Python
   code and the checksum manifest; package provenance and filesystem ownership
@@ -22,8 +23,9 @@ it redacts, and what remains sensitive in a report.
 Passwords should be supplied by a prompt or an environment variable named by
 `password_env`. Literal TOML passwords are rejected. Passwords, private-key
 contents and authentication options are not written to stdout, logs or
-artifacts. CLI password arguments remain visible to local process inspection
-and are therefore discouraged.
+artifacts. The local agent socket pathname can appear in validation errors; it
+should not contain sensitive naming. CLI password arguments remain visible to
+local process inspection and are therefore discouraged.
 
 `--strip-meta` removes embedded vendor SQL/shell sources, instructions,
 manifests, provenance and the HTML metadata controls. It does not redact values

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ch-diag"
-version = "0.8.0"
+version = "0.8.1"
 description = "ClickHouse diagnostics with host metrics and portable JSON/HTML reports"
 readme = "README.md"
 requires-python = ">=3.10"