diff --git a/cvs/input/config_file/cluster_key_distribution/ssh_key_distribute_config.json b/cvs/input/config_file/cluster_key_distribution/ssh_key_distribute_config.json new file mode 100644 index 000000000..3de95e4db --- /dev/null +++ b/cvs/input/config_file/cluster_key_distribution/ssh_key_distribute_config.json @@ -0,0 +1,31 @@ +{ + "ssh_key_distribution": { + "_comment_keys": "Local paths on the machine running CVS. Replace every changeme placeholder below. The private key is uploaded via SFTP (never echoed).", + "cluster_key_private_path": "", + "cluster_key_public_path": "", + + "_comment_key_name": "Filename the key gets on the remote nodes under remote_ssh_dir. Peers reference it via IdentityFile.", + "key_name": "cluster_id", + + "_comment_controlling": "Optional. Local path to the controlling station's EXISTING public key to also authorize. Leave as empty string to skip.", + "controlling_station_pubkey_path": "", + + "_comment_remote_ssh_dir": "Optional. Remote per-user ssh dir. Default ~/.ssh. '~' expands on the remote node.", + "remote_ssh_dir": "~/.ssh", + + "_comment_ssh_dir_nfs_shared": "Optional. Set true if remote_ssh_dir is the same NFS-mounted directory on every node (e.g. a shared home directory). When true, the cluster keypair and controlling-station key are uploaded via the head node only -- every other node already sees the same file through the shared mount, so per-node upload is skipped.", + "ssh_dir_nfs_shared": false, + + "_comment_host_pattern": "Optional override for the ssh_config Host wildcard. If empty, derived from node_dict keys + vpc_ips.", + "ssh_config_host_pattern": "", + + "_comment_verify": "Verification behavior after distribution.", + "verify_connectivity": true, + "_comment_verify_mode": "'ring' (each node -> next, O(n)) or 'full_mesh' (every ordered pair, O(n^2)).", + "verify_mode": "ring", + "verify_timeout": 20, + + "_comment_config_block": "How ~/.ssh/config is managed. 'managed_block' inserts/replaces a marked CVS block, preserving other content (recommended). 'overwrite' replaces the whole file.", + "ssh_config_write_mode": "managed_block" + } +} diff --git a/cvs/lib/ssh_keys_lib.py b/cvs/lib/ssh_keys_lib.py new file mode 100644 index 000000000..8060e50be --- /dev/null +++ b/cvs/lib/ssh_keys_lib.py @@ -0,0 +1,461 @@ +''' +Copyright 2025 Advanced Micro Devices, Inc. +All rights reserved. This notice is intended as a precaution against inadvertent +publication and does not imply publication or any waiver of confidentiality. +The year included in the foregoing notice is the year of creation of the work. +All code contained here is Property of Advanced Micro Devices, Inc. +''' + +import base64 +import ipaddress +import os +import shlex + +from cvs.lib import globals + +log = globals.log + +DEFAULT_KEY_NAME = "cluster_id" +DEFAULT_REMOTE_SSH_DIR = "~/.ssh" +SSH_CONFIG_BEGIN = "# BEGIN CVS cluster_key_distribution (managed)" +SSH_CONFIG_END = "# END CVS cluster_key_distribution (managed)" + +_KNOWN_DEFAULT_KEY_NAMES = {"id_rsa", "id_ed25519", "id_ecdsa", "id_dsa"} + + +# --------------------------------------------------------------------------- +# Pure / logic functions +# --------------------------------------------------------------------------- + + +def validate_key_distribution_config(config_dict): + """Validate config subsection and return a normalized dict with defaults applied.""" + norm = dict(config_dict) + + for field in ("cluster_key_private_path", "cluster_key_public_path"): + val = norm.get(field, "") + if not val: + raise ValueError(f"ssh_key_distribution.{field} is required and must be non-empty") + if not os.path.isfile(val): + raise ValueError(f"ssh_key_distribution.{field}: file not found: {val!r}") + + controlling = norm.get("controlling_station_pubkey_path", "") + if controlling and not os.path.isfile(controlling): + raise ValueError(f"ssh_key_distribution.controlling_station_pubkey_path: file not found: {controlling!r}") + + norm.setdefault("key_name", DEFAULT_KEY_NAME) + norm.setdefault("remote_ssh_dir", DEFAULT_REMOTE_SSH_DIR) + norm.setdefault("ssh_dir_nfs_shared", False) + norm.setdefault("ssh_config_host_pattern", "") + norm.setdefault("verify_connectivity", True) + norm.setdefault("verify_mode", "ring") + norm.setdefault("verify_timeout", 20) + norm.setdefault("ssh_config_write_mode", "managed_block") + norm.setdefault("controlling_station_pubkey_path", "") + + key_name = norm["key_name"] + if key_name in _KNOWN_DEFAULT_KEY_NAMES: + log.warning( + "ssh_key_distribution.key_name=%r matches a well-known SSH default identity; " + "distribution may overwrite an existing key on remote nodes", + key_name, + ) + + return norm + + +def collect_cluster_hostnames(cluster_dict): + """Return ordered, de-duplicated SSH-reachable identifiers: node keys + distinct vpc_ips.""" + seen = set() + result = [] + for node_name, node_info in cluster_dict.get("node_dict", {}).items(): + if node_name not in seen: + seen.add(node_name) + result.append(node_name) + vpc_ip = node_info.get("vpc_ip", "") if isinstance(node_info, dict) else "" + if vpc_ip and vpc_ip != node_name and vpc_ip not in seen: + seen.add(vpc_ip) + result.append(vpc_ip) + return result + + +def _longest_common_prefix(strings): + """Return the longest common leading substring across a non-empty list.""" + if not strings: + return "" + prefix = strings[0] + for s in strings[1:]: + while not s.startswith(prefix): + prefix = prefix[:-1] + if not prefix: + return "" + return prefix + + +def derive_ssh_host_pattern(hostnames, override=""): + """Derive the SSH Host line token(s) covering all cluster nodes. + + Resolution order: + 1. Non-empty override → use verbatim. + 2. All IPv4 → collapse on longest shared octet boundary. + 3. Non-IP names share a non-trivial common alphanumeric prefix → prefix*. + 4. Fallback: space-joined explicit list. + """ + if override: + return override + + if not hostnames: + return "*" + + if len(hostnames) == 1: + return hostnames[0] + + # Step 2: all IPv4? + parsed_ips = [] + for h in hostnames: + try: + parsed_ips.append(ipaddress.ip_address(h)) + except ValueError: + parsed_ips = [] + break + + all_ips = bool(parsed_ips) + + if all_ips: + octets = [str(ip).split(".") for ip in parsed_ips] + shared = 0 + for i in range(3): + if len({o[i] for o in octets}) == 1: + shared = i + 1 + else: + break + if shared == 3: + prefix_octets = octets[0][:3] + return ".".join(prefix_octets) + ".*" + if shared == 2: + prefix_octets = octets[0][:2] + return ".".join(prefix_octets) + ".*" + if shared == 1: + return octets[0][0] + ".*" + # all IPs but no shared octet prefix → explicit list (skip string-prefix step) + return " ".join(hostnames) + + # Step 3: common alphanumeric prefix (non-IP names only) + prefix = _longest_common_prefix(hostnames) + if prefix and any(len(h) > len(prefix) for h in hostnames): + return prefix + "*" + + # Step 4: explicit list + return " ".join(hostnames) + + +def render_ssh_config_block(host_pattern, username, identity_file): + """Return the ~/.ssh/config text block with BEGIN/END markers.""" + lines = [ + SSH_CONFIG_BEGIN, + f"Host {host_pattern}", + f" User {username}", + f" IdentityFile {identity_file}", + " StrictHostKeyChecking no", + " UserKnownHostsFile /dev/null", + " LogLevel ERROR", + SSH_CONFIG_END, + ] + return "\n".join(lines) + "\n" + + +# --------------------------------------------------------------------------- +# Remote-command builders +# --------------------------------------------------------------------------- + + +def build_ensure_ssh_dir_cmd(remote_ssh_dir): + """mkdir -p + chmod 700 on the remote ssh dir (~ preserved for remote shell expansion).""" + if remote_ssh_dir.startswith("~"): + dir_arg = remote_ssh_dir + else: + dir_arg = shlex.quote(remote_ssh_dir) + return f"mkdir -p {dir_arg} && chmod 700 {dir_arg}" + + +def build_key_perms_cmd(remote_ssh_dir, key_name): + """chmod 600 on private key, chmod 644 on public key.""" + if remote_ssh_dir.startswith("~"): + d = remote_ssh_dir + else: + d = shlex.quote(remote_ssh_dir) + priv = shlex.quote(key_name) + pub = shlex.quote(key_name + ".pub") + return f"chmod 600 {d}/{priv} && chmod 644 {d}/{pub}" + + +def build_authorize_pubkey_cmd(remote_ssh_dir, pubkey_remote_path): + """Idempotent append of a remote pubkey file into authorized_keys using grep -qxF.""" + if remote_ssh_dir.startswith("~"): + d = remote_ssh_dir + else: + d = shlex.quote(remote_ssh_dir) + ak = f"{d}/authorized_keys" + pub = pubkey_remote_path if pubkey_remote_path.startswith("~") else shlex.quote(pubkey_remote_path) + # grep -qxF uses fixed-string exact whole-line match — never duplicates; key never on the cmd line + return f"touch {ak} && chmod 600 {ak} && grep -qxF -- \"$(cat {pub})\" {ak} || cat {pub} >> {ak}" + + +def build_write_ssh_config_cmd(remote_ssh_dir, block_text, mode): + """Return a shell command that installs block_text into ~/.ssh/config. + + Uses base64 encoding so the multi-line block survives over a single exec channel + without heredoc quoting issues or libssh2's ~30 KB exec_request limit. + """ + if remote_ssh_dir.startswith("~"): + d = remote_ssh_dir + else: + d = shlex.quote(remote_ssh_dir) + cfg = f"{d}/config" + b64 = base64.b64encode(block_text.encode()).decode() + + if mode == "overwrite": + return f"echo {shlex.quote(b64)} | base64 -d > {cfg} && chmod 600 {cfg}" + + # managed_block: delete old CVS block (if any) then append fresh one + sed_script = shlex.quote(f"/{SSH_CONFIG_BEGIN}/,/{SSH_CONFIG_END}/d") + sed_del = f"sed -i {sed_script} {cfg} 2>/dev/null || true" + append = f"echo {shlex.quote(b64)} | base64 -d >> {cfg} && chmod 600 {cfg}" + return f"touch {cfg} && {sed_del} && {append}" + + +# --------------------------------------------------------------------------- +# Orchestration drivers +# --------------------------------------------------------------------------- + + +def upload_cluster_keys(orch, norm_config): + """SFTP private+public cluster key to every node; apply permissions. Returns {node: bool}. + + When remote_ssh_dir is NFS-shared across nodes (ssh_dir_nfs_shared=True), the + keypair is uploaded once via the head node only -- every other node already + sees the same file through the shared mount, so a per-node upload would just + duplicate the same SFTP transfer. + """ + priv_local = norm_config["cluster_key_private_path"] + pub_local = norm_config["cluster_key_public_path"] + key_name = norm_config["key_name"] + remote_ssh_dir = norm_config["remote_ssh_dir"] + nfs_shared = norm_config.get("ssh_dir_nfs_shared", False) + + sftp_dir = _expand_remote_dir(orch, remote_ssh_dir) + remote_priv = f"{sftp_dir}/{key_name}" + remote_pub = f"{sftp_dir}/{key_name}.pub" + + results = {node: True for node in orch.all.reachable_hosts} + upload_target = orch.head if nfs_shared else orch.all + + for local, remote in ((priv_local, remote_priv), (pub_local, remote_pub)): + try: + upload_target.upload_file(local, remote) + except IOError as e: + log.error("SFTP upload %s -> %s failed: %r", local, remote, e) + for node in results: + results[node] = False + + perms_cmd = build_key_perms_cmd(remote_ssh_dir, key_name) + out = orch.exec(perms_cmd, timeout=30, detailed=True) + for node, detail in out.items(): + if isinstance(detail, dict): + if detail.get("exit_code", 0) != 0: + log.error("chmod keys failed on %s: %s", node, detail.get("output", "")) + results[node] = False + else: + if "error" in str(detail).lower(): + results[node] = False + + return results + + +def authorize_cluster_pubkey(orch, norm_config): + """Append cluster pubkey to authorized_keys. NFS-shared dirs go via the head node only. + + Returns {node: bool}. + """ + key_name = norm_config["key_name"] + remote_ssh_dir = norm_config["remote_ssh_dir"] + nfs_shared = norm_config.get("ssh_dir_nfs_shared", False) + pubkey_remote = f"{remote_ssh_dir}/{key_name}.pub" + + cmd = build_authorize_pubkey_cmd(remote_ssh_dir, pubkey_remote) + exec_target = orch.exec_on_head if nfs_shared else orch.exec + out = exec_target(cmd, timeout=30, detailed=True) + + results = {node: True for node in orch.all.reachable_hosts} + for node, ok in _detailed_to_bool(out).items(): + results[node] = results.get(node, True) and ok + return results + + +def authorize_controlling_station(orch, norm_config): + """Upload controlling station pubkey and append to authorized_keys. No-op if path unset.""" + controlling_local = norm_config.get("controlling_station_pubkey_path", "") + if not controlling_local: + return {} + + remote_ssh_dir = norm_config["remote_ssh_dir"] + nfs_shared = norm_config.get("ssh_dir_nfs_shared", False) + sftp_dir = _expand_remote_dir(orch, remote_ssh_dir) + remote_tmp = f"{sftp_dir}/.cvs_controlling_station.pub" + + results = {node: True for node in orch.all.reachable_hosts} + upload_target = orch.head if nfs_shared else orch.all + try: + upload_target.upload_file(controlling_local, remote_tmp) + except IOError as e: + log.error("SFTP upload controlling station key failed: %r", e) + for node in results: + results[node] = False + return results + + cmd = build_authorize_pubkey_cmd(remote_ssh_dir, remote_tmp) + exec_target = orch.exec_on_head if nfs_shared else orch.exec + out = exec_target(cmd, timeout=30, detailed=True) + for node, ok in _detailed_to_bool(out).items(): + results[node] = results.get(node, True) and ok + + return results + + +def install_ssh_config(orch, cluster_dict, norm_config): + """Derive host pattern, render config block, install ~/.ssh/config. + + NFS-shared remote_ssh_dir writes the file via the head node only -- concurrent + sed -i (temp-file + rename) against the same underlying file from multiple nodes + at once can trip a stale NFS file handle on the other nodes. + + Returns {node: bool}. + """ + remote_ssh_dir = norm_config["remote_ssh_dir"] + key_name = norm_config["key_name"] + override = norm_config.get("ssh_config_host_pattern", "") + mode = norm_config.get("ssh_config_write_mode", "managed_block") + nfs_shared = norm_config.get("ssh_dir_nfs_shared", False) + username = cluster_dict.get("username", "") + identity_file = f"{remote_ssh_dir}/{key_name}" + + hostnames = collect_cluster_hostnames(cluster_dict) + pattern = derive_ssh_host_pattern(hostnames, override=override) + block = render_ssh_config_block(pattern, username, identity_file) + + cmd = build_write_ssh_config_cmd(remote_ssh_dir, block, mode) + exec_target = orch.exec_on_head if nfs_shared else orch.exec + out = exec_target(cmd, timeout=30, detailed=True) + + results = {node: True for node in orch.all.reachable_hosts} + for node, ok in _detailed_to_bool(out).items(): + results[node] = results.get(node, True) and ok + return results + + +def verify_passwordless_ssh(orch, cluster_dict, norm_config): + """Probe passwordless SSH between node pairs. Returns {(src, dst): bool}.""" + nodes = list(cluster_dict.get("node_dict", {}).keys()) + if len(nodes) < 2: + return {} + + remote_ssh_dir = norm_config["remote_ssh_dir"] + timeout = norm_config.get("verify_timeout", 20) + mode = norm_config.get("verify_mode", "ring") + + ssh_config_path = f"{remote_ssh_dir}/config" + + results = {} + + if mode == "ring": + # Build one probe command per node (node i -> node i+1 mod n), run via exec_cmd_list + cmd_list = [] + pairs = [] + for i, src in enumerate(nodes): + dst = nodes[(i + 1) % len(nodes)] + if src == dst: + continue + pairs.append((src, dst)) + cmd_list.append(f"ssh -F {ssh_config_path} -o BatchMode=yes -o ConnectTimeout={timeout} {dst} true") + out = orch.all.exec_cmd_list(cmd_list, timeout=timeout + 10) + for (src, dst), output in zip( + pairs, [out.get(node, "") for node in nodes if node != nodes[0] or len(nodes) < 2] + ): + results[(src, dst)] = "error" not in str(output).lower() and output is not None + + # exec_cmd_list returns {node: output}; map back by position + node_outputs = [out.get(node, "") for node in nodes] + results = {} + for i, (src, dst) in enumerate(pairs): + raw = node_outputs[i] if i < len(node_outputs) else "" + results[(src, dst)] = raw is not None and "error" not in str(raw).lower() + + else: + # full_mesh: O(n*(n-1)) probes, one exec per source node + for src in nodes: + peers = [n for n in nodes if n != src] + if not peers: + continue + cmd_list = [ + f"ssh -F {ssh_config_path} -o BatchMode=yes -o ConnectTimeout={timeout} {dst} true" for dst in peers + ] + # Run all probes from src via a temporary single-host handle + from cvs.lib.parallel_ssh_lib import Pssh + + tmp = Pssh( + log, + [src], + user=cluster_dict.get("username"), + pkey=cluster_dict.get("priv_key_file"), + host_key_check=False, + ) + try: + src_out = tmp.exec_cmd_list(cmd_list, timeout=timeout + 10) + outputs = [src_out.get(src, "")] if len(peers) == 1 else list(src_out.values()) + for dst, raw in zip(peers, outputs): + results[(src, dst)] = raw is not None and "error" not in str(raw).lower() + finally: + tmp.destroy_clients() + + return results + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _expand_remote_dir(orch, remote_dir): + """Resolve a leading '~' to an absolute path for SFTP use. + + orch.exec runs through a remote shell (which expands '~'), but SFTP + transfers (upload_file/copy_file) go over a raw SFTP channel with no + shell involved -- a literal '~' is treated as an ordinary path + component, silently writing into a directory named '~' instead of + the real home directory. + """ + if not remote_dir.startswith("~"): + return remote_dir + out = orch.exec(f"echo {remote_dir}", timeout=10) + return next(iter(out.values())).strip() + + +def _detailed_to_bool(out): + """Convert a detailed=True exec result dict to {node: bool}.""" + result = {} + for node, detail in out.items(): + if isinstance(detail, dict): + result[node] = detail.get("exit_code", 0) == 0 + else: + result[node] = True + return result + + +def _upload_local_file(orch, local_path, remote_path): + """SFTP upload with IOError handling. Returns True on success.""" + try: + orch.all.upload_file(local_path, remote_path) + return True + except IOError as e: + log.error("SFTP upload %s -> %s failed: %r", local_path, remote_path, e) + return False diff --git a/cvs/lib/unittests/test_ssh_keys_lib.py b/cvs/lib/unittests/test_ssh_keys_lib.py new file mode 100644 index 000000000..f3ec32b67 --- /dev/null +++ b/cvs/lib/unittests/test_ssh_keys_lib.py @@ -0,0 +1,542 @@ +# cvs/lib/unittests/test_ssh_keys_lib.py +import unittest +from unittest.mock import MagicMock, patch + +import cvs.lib.ssh_keys_lib as lib + + +class TestValidateKeyDistributionConfig(unittest.TestCase): + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_valid_config_returns_normalized(self, _isfile): + cfg = { + "cluster_key_private_path": "/tmp/id", + "cluster_key_public_path": "/tmp/id.pub", + } + norm = lib.validate_key_distribution_config(cfg) + self.assertEqual(norm["key_name"], "cluster_id") + self.assertEqual(norm["remote_ssh_dir"], "~/.ssh") + self.assertEqual(norm["ssh_dir_nfs_shared"], False) + self.assertEqual(norm["verify_mode"], "ring") + self.assertEqual(norm["verify_timeout"], 20) + self.assertEqual(norm["ssh_config_write_mode"], "managed_block") + self.assertEqual(norm["controlling_station_pubkey_path"], "") + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_explicit_values_not_overwritten_by_defaults(self, _isfile): + cfg = { + "cluster_key_private_path": "/tmp/id", + "cluster_key_public_path": "/tmp/id.pub", + "key_name": "mykey", + "verify_mode": "full_mesh", + "verify_timeout": 60, + } + norm = lib.validate_key_distribution_config(cfg) + self.assertEqual(norm["key_name"], "mykey") + self.assertEqual(norm["verify_mode"], "full_mesh") + self.assertEqual(norm["verify_timeout"], 60) + + def test_missing_private_path_raises(self): + with self.assertRaises(ValueError): + lib.validate_key_distribution_config({"cluster_key_public_path": "/tmp/id.pub"}) + + def test_empty_private_path_raises(self): + with self.assertRaises(ValueError): + lib.validate_key_distribution_config( + {"cluster_key_private_path": "", "cluster_key_public_path": "/tmp/id.pub"} + ) + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=False) + def test_nonexistent_private_path_raises(self, _isfile): + with self.assertRaises(ValueError): + lib.validate_key_distribution_config( + { + "cluster_key_private_path": "/no/such/file", + "cluster_key_public_path": "/tmp/id.pub", + } + ) + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", side_effect=lambda p: p != "/no/ctrl") + def test_controlling_path_set_but_missing_raises(self, _isfile): + with self.assertRaises(ValueError): + lib.validate_key_distribution_config( + { + "cluster_key_private_path": "/tmp/id", + "cluster_key_public_path": "/tmp/id.pub", + "controlling_station_pubkey_path": "/no/ctrl", + } + ) + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_controlling_path_empty_ok(self, _isfile): + cfg = { + "cluster_key_private_path": "/tmp/id", + "cluster_key_public_path": "/tmp/id.pub", + "controlling_station_pubkey_path": "", + } + norm = lib.validate_key_distribution_config(cfg) + self.assertEqual(norm["controlling_station_pubkey_path"], "") + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_known_default_key_name_logs_warning(self, _isfile): + cfg = { + "cluster_key_private_path": "/tmp/id", + "cluster_key_public_path": "/tmp/id.pub", + "key_name": "id_rsa", + } + with self.assertLogs("root", level="WARNING"): + lib.validate_key_distribution_config(cfg) + + +class TestCollectClusterHostnames(unittest.TestCase): + def test_node_dict_only(self): + cluster = {"node_dict": {"node1": {}, "node2": {}}} + result = lib.collect_cluster_hostnames(cluster) + self.assertEqual(result, ["node1", "node2"]) + + def test_distinct_vpc_ip_included(self): + cluster = { + "node_dict": { + "node1": {"vpc_ip": "10.0.0.1"}, + "node2": {"vpc_ip": "10.0.0.2"}, + } + } + result = lib.collect_cluster_hostnames(cluster) + self.assertIn("node1", result) + self.assertIn("10.0.0.1", result) + self.assertIn("node2", result) + self.assertIn("10.0.0.2", result) + self.assertEqual(len(result), 4) + + def test_vpc_ip_same_as_node_name_not_duplicated(self): + cluster = {"node_dict": {"10.0.0.1": {"vpc_ip": "10.0.0.1"}}} + result = lib.collect_cluster_hostnames(cluster) + self.assertEqual(result, ["10.0.0.1"]) + + def test_order_preserved(self): + cluster = {"node_dict": {"node3": {}, "node1": {}, "node2": {}}} + result = lib.collect_cluster_hostnames(cluster) + self.assertEqual(result, ["node3", "node1", "node2"]) + + def test_empty_node_dict(self): + self.assertEqual(lib.collect_cluster_hostnames({"node_dict": {}}), []) + + def test_no_node_dict(self): + self.assertEqual(lib.collect_cluster_hostnames({}), []) + + +class TestLongestCommonPrefix(unittest.TestCase): + def test_shared_prefix(self): + self.assertEqual(lib._longest_common_prefix(["node01", "node02", "node03"]), "node0") + + def test_no_shared_prefix(self): + self.assertEqual(lib._longest_common_prefix(["abc", "xyz"]), "") + + def test_single_element(self): + self.assertEqual(lib._longest_common_prefix(["hello"]), "hello") + + def test_identical_elements(self): + self.assertEqual(lib._longest_common_prefix(["foo", "foo"]), "foo") + + def test_empty_list(self): + self.assertEqual(lib._longest_common_prefix([]), "") + + +class TestDeriveSshHostPattern(unittest.TestCase): + def test_override_wins(self): + result = lib.derive_ssh_host_pattern(["node1", "node2"], override="myoverride") + self.assertEqual(result, "myoverride") + + def test_named_prefix_wildcard(self): + result = lib.derive_ssh_host_pattern(["node01", "node02", "node03"]) + self.assertEqual(result, "node0*") + + def test_ip_shared_three_octets(self): + result = lib.derive_ssh_host_pattern(["10.0.0.1", "10.0.0.2"]) + self.assertEqual(result, "10.0.0.*") + + def test_ip_shared_two_octets(self): + result = lib.derive_ssh_host_pattern(["10.0.1.1", "10.0.2.1"]) + self.assertEqual(result, "10.0.*") + + def test_ip_shared_one_octet(self): + result = lib.derive_ssh_host_pattern(["10.1.0.1", "10.2.0.1"]) + self.assertEqual(result, "10.*") + + def test_ip_no_shared_octet_explicit_list(self): + result = lib.derive_ssh_host_pattern(["10.0.0.1", "192.168.0.1"]) + self.assertIn("10.0.0.1", result) + self.assertIn("192.168.0.1", result) + + def test_mixed_names_explicit_list(self): + result = lib.derive_ssh_host_pattern(["gpu-a", "worker-b"]) + self.assertIn("gpu-a", result) + self.assertIn("worker-b", result) + + def test_single_host(self): + result = lib.derive_ssh_host_pattern(["node1"]) + self.assertEqual(result, "node1") + + +class TestRenderSshConfigBlock(unittest.TestCase): + def test_block_structure(self): + block = lib.render_ssh_config_block("node*", "myuser", "~/.ssh/cluster_id") + self.assertIn(lib.SSH_CONFIG_BEGIN, block) + self.assertIn(lib.SSH_CONFIG_END, block) + self.assertIn("Host node*", block) + self.assertIn("User myuser", block) + self.assertIn("IdentityFile ~/.ssh/cluster_id", block) + self.assertIn("StrictHostKeyChecking no", block) + self.assertIn("UserKnownHostsFile /dev/null", block) + self.assertIn("LogLevel ERROR", block) + + def test_begin_before_end(self): + block = lib.render_ssh_config_block("*", "u", "/path") + begin_pos = block.index(lib.SSH_CONFIG_BEGIN) + end_pos = block.index(lib.SSH_CONFIG_END) + self.assertLess(begin_pos, end_pos) + + +class TestBuildEnsureSshDirCmd(unittest.TestCase): + def test_contains_mkdir_and_chmod(self): + cmd = lib.build_ensure_ssh_dir_cmd("~/.ssh") + self.assertIn("mkdir -p", cmd) + self.assertIn("chmod 700", cmd) + self.assertIn("~/.ssh", cmd) + + def test_non_home_path_quoted(self): + cmd = lib.build_ensure_ssh_dir_cmd("/some/path with spaces") + self.assertIn("chmod 700", cmd) + + +class TestBuildKeyPermsCmd(unittest.TestCase): + def test_private_600_public_644(self): + cmd = lib.build_key_perms_cmd("~/.ssh", "cluster_id") + self.assertIn("chmod 600", cmd) + self.assertIn("chmod 644", cmd) + self.assertIn("cluster_id.pub", cmd) + + +class TestBuildAuthorizePubkeyCmd(unittest.TestCase): + def test_uses_grep_qxf(self): + cmd = lib.build_authorize_pubkey_cmd("~/.ssh", "~/.ssh/cluster_id.pub") + self.assertIn("grep -qxF", cmd) + self.assertIn("authorized_keys", cmd) + + def test_never_contains_raw_key_material(self): + cmd = lib.build_authorize_pubkey_cmd("~/.ssh", "~/.ssh/cluster_id.pub") + # Key content never embedded; only cat via subshell + self.assertIn("cat ", cmd) + self.assertNotIn("ssh-rsa", cmd) + + def test_touch_and_chmod_600(self): + cmd = lib.build_authorize_pubkey_cmd("~/.ssh", "~/.ssh/cluster_id.pub") + self.assertIn("touch", cmd) + self.assertIn("chmod 600", cmd) + + +class TestBuildWriteSshConfigCmd(unittest.TestCase): + def test_managed_block_contains_sed_and_base64(self): + cmd = lib.build_write_ssh_config_cmd("~/.ssh", "block text\n", "managed_block") + self.assertIn("sed", cmd) + self.assertIn("base64", cmd) + self.assertIn("chmod 600", cmd) + + def test_overwrite_no_sed(self): + cmd = lib.build_write_ssh_config_cmd("~/.ssh", "block\n", "overwrite") + self.assertNotIn("sed", cmd) + self.assertIn("base64", cmd) + self.assertIn("chmod 600", cmd) + + def test_path_with_spaces_quoted(self): + cmd = lib.build_write_ssh_config_cmd("/home/my user/.ssh", "block\n", "managed_block") + # shlex.quote wraps the path in single quotes so the raw unquoted path must not appear bare + self.assertNotIn(" /home/my user/", cmd) + + +class TestExpandRemoteDir(unittest.TestCase): + def test_absolute_path_returned_unchanged(self): + orch = MagicMock() + result = lib._expand_remote_dir(orch, "/home/user/.ssh") + self.assertEqual(result, "/home/user/.ssh") + orch.exec.assert_not_called() + + def test_tilde_path_expanded_via_exec(self): + orch = MagicMock() + orch.exec.return_value = {"n1": "/home/user/.ssh\n"} + result = lib._expand_remote_dir(orch, "~/.ssh") + self.assertEqual(result, "/home/user/.ssh") + orch.exec.assert_called_once_with("echo ~/.ssh", timeout=10) + + +class TestUploadClusterKeys(unittest.TestCase): + def _make_orch(self, exit_code=0): + orch = MagicMock() + orch.all.reachable_hosts = ["n1", "n2"] + + def exec_side_effect(cmd, timeout=30, detailed=False): + if cmd.startswith("echo "): + return {"n1": "/home/user/.ssh", "n2": "/home/user/.ssh"} + return { + "n1": {"output": "", "exit_code": exit_code}, + "n2": {"output": "", "exit_code": exit_code}, + } + + orch.exec.side_effect = exec_side_effect + return orch + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_upload_called_for_priv_and_pub(self, _isfile): + orch = self._make_orch() + norm = { + "cluster_key_private_path": "/local/id", + "cluster_key_public_path": "/local/id.pub", + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + } + results = lib.upload_cluster_keys(orch, norm) + self.assertEqual(orch.all.upload_file.call_count, 2) + calls = orch.all.upload_file.call_args_list + remote_paths = [c[0][1] for c in calls] + self.assertIn("/home/user/.ssh/cluster_id", remote_paths) + self.assertIn("/home/user/.ssh/cluster_id.pub", remote_paths) + self.assertTrue(all(results.values())) + + def test_upload_ioerror_marks_nodes_failed(self): + orch = self._make_orch() + orch.all.upload_file.side_effect = IOError("sftp fail") + norm = { + "cluster_key_private_path": "/local/id", + "cluster_key_public_path": "/local/id.pub", + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + } + results = lib.upload_cluster_keys(orch, norm) + self.assertTrue(all(not v for v in results.values())) + + def test_chmod_failure_marks_node_failed(self): + orch = self._make_orch(exit_code=1) + norm = { + "cluster_key_private_path": "/local/id", + "cluster_key_public_path": "/local/id.pub", + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + } + results = lib.upload_cluster_keys(orch, norm) + self.assertFalse(results["n1"]) + self.assertFalse(results["n2"]) + + @patch("cvs.lib.ssh_keys_lib.os.path.isfile", return_value=True) + def test_nfs_shared_uploads_via_head_only(self, _isfile): + orch = self._make_orch() + norm = { + "cluster_key_private_path": "/local/id", + "cluster_key_public_path": "/local/id.pub", + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + "ssh_dir_nfs_shared": True, + } + results = lib.upload_cluster_keys(orch, norm) + self.assertEqual(orch.head.upload_file.call_count, 2) + orch.all.upload_file.assert_not_called() + self.assertTrue(all(results.values())) + + def test_nfs_shared_head_upload_ioerror_marks_all_nodes_failed(self): + orch = self._make_orch() + orch.head.upload_file.side_effect = IOError("sftp fail") + norm = { + "cluster_key_private_path": "/local/id", + "cluster_key_public_path": "/local/id.pub", + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + "ssh_dir_nfs_shared": True, + } + results = lib.upload_cluster_keys(orch, norm) + self.assertTrue(all(not v for v in results.values())) + + +class TestAuthorizeClusterPubkey(unittest.TestCase): + def test_returns_success_dict(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1", "n2"] + orch.exec.return_value = { + "n1": {"output": "", "exit_code": 0}, + "n2": {"output": "", "exit_code": 0}, + } + norm = {"key_name": "cluster_id", "remote_ssh_dir": "~/.ssh"} + results = lib.authorize_cluster_pubkey(orch, norm) + self.assertTrue(results["n1"]) + self.assertTrue(results["n2"]) + orch.exec.assert_called_once() + cmd = orch.exec.call_args[0][0] + self.assertIn("grep -qxF", cmd) + + def test_nonzero_exit_returns_false(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1"] + orch.exec.return_value = {"n1": {"output": "err", "exit_code": 1}} + norm = {"key_name": "cluster_id", "remote_ssh_dir": "~/.ssh"} + results = lib.authorize_cluster_pubkey(orch, norm) + self.assertFalse(results["n1"]) + + def test_nfs_shared_authorizes_via_head_only(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1", "n2"] + orch.exec_on_head.return_value = { + "n1": {"output": "", "exit_code": 0}, + "n2": {"output": "", "exit_code": 0}, + } + norm = { + "key_name": "cluster_id", + "remote_ssh_dir": "~/.ssh", + "ssh_dir_nfs_shared": True, + } + results = lib.authorize_cluster_pubkey(orch, norm) + orch.exec_on_head.assert_called_once() + orch.exec.assert_not_called() + self.assertTrue(all(results.values())) + + +class TestAuthorizeControllingStation(unittest.TestCase): + def test_empty_path_returns_empty_dict(self): + orch = MagicMock() + norm = {"controlling_station_pubkey_path": "", "remote_ssh_dir": "~/.ssh"} + results = lib.authorize_controlling_station(orch, norm) + self.assertEqual(results, {}) + orch.all.upload_file.assert_not_called() + + def test_upload_and_authorize_called(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1"] + + def exec_side_effect(cmd, timeout=30, detailed=False): + if cmd.startswith("echo "): + return {"n1": "/home/user/.ssh"} + return {"n1": {"output": "", "exit_code": 0}} + + orch.exec.side_effect = exec_side_effect + norm = { + "controlling_station_pubkey_path": "/local/ctrl.pub", + "remote_ssh_dir": "~/.ssh", + } + results = lib.authorize_controlling_station(orch, norm) + orch.all.upload_file.assert_called_once_with("/local/ctrl.pub", "/home/user/.ssh/.cvs_controlling_station.pub") + self.assertTrue(results["n1"]) + + def test_upload_ioerror_marks_failed(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1"] + orch.exec.return_value = {"n1": "/home/user/.ssh"} + orch.all.upload_file.side_effect = IOError("fail") + norm = { + "controlling_station_pubkey_path": "/local/ctrl.pub", + "remote_ssh_dir": "~/.ssh", + } + results = lib.authorize_controlling_station(orch, norm) + self.assertFalse(results["n1"]) + + def test_nfs_shared_uploads_via_head_only(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1", "n2"] + + def exec_side_effect(cmd, timeout=30, detailed=False): + return {"n1": "/home/user/.ssh", "n2": "/home/user/.ssh"} + + orch.exec.side_effect = exec_side_effect + orch.exec_on_head.return_value = { + "n1": {"output": "", "exit_code": 0}, + "n2": {"output": "", "exit_code": 0}, + } + norm = { + "controlling_station_pubkey_path": "/local/ctrl.pub", + "remote_ssh_dir": "~/.ssh", + "ssh_dir_nfs_shared": True, + } + results = lib.authorize_controlling_station(orch, norm) + orch.head.upload_file.assert_called_once_with("/local/ctrl.pub", "/home/user/.ssh/.cvs_controlling_station.pub") + orch.all.upload_file.assert_not_called() + orch.exec_on_head.assert_called_once() + self.assertTrue(all(results.values())) + + +class TestInstallSshConfig(unittest.TestCase): + def test_uses_cluster_username(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1"] + orch.exec.return_value = {"n1": {"output": "", "exit_code": 0}} + cluster = {"username": "myuser", "node_dict": {"n1": {}}} + norm = { + "remote_ssh_dir": "~/.ssh", + "key_name": "cluster_id", + "ssh_config_host_pattern": "", + "ssh_config_write_mode": "managed_block", + } + lib.install_ssh_config(orch, cluster, norm) + cmd = orch.exec.call_args[0][0] + self.assertIn("base64", cmd) + + def test_returns_success_dict(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1"] + orch.exec.return_value = {"n1": {"output": "", "exit_code": 0}} + cluster = {"username": "u", "node_dict": {"n1": {}}} + norm = { + "remote_ssh_dir": "~/.ssh", + "key_name": "cluster_id", + "ssh_config_host_pattern": "", + "ssh_config_write_mode": "managed_block", + } + results = lib.install_ssh_config(orch, cluster, norm) + self.assertTrue(results["n1"]) + + def test_nfs_shared_installs_via_head_only(self): + orch = MagicMock() + orch.all.reachable_hosts = ["n1", "n2"] + orch.exec_on_head.return_value = { + "n1": {"output": "", "exit_code": 0}, + "n2": {"output": "", "exit_code": 0}, + } + cluster = {"username": "u", "node_dict": {"n1": {}, "n2": {}}} + norm = { + "remote_ssh_dir": "~/.ssh", + "key_name": "cluster_id", + "ssh_config_host_pattern": "", + "ssh_config_write_mode": "managed_block", + "ssh_dir_nfs_shared": True, + } + results = lib.install_ssh_config(orch, cluster, norm) + orch.exec_on_head.assert_called_once() + orch.exec.assert_not_called() + self.assertTrue(all(results.values())) + + +class TestVerifyPasswordlessSsh(unittest.TestCase): + def test_single_node_returns_empty(self): + orch = MagicMock() + cluster = {"node_dict": {"n1": {}}} + norm = {"remote_ssh_dir": "~/.ssh", "verify_timeout": 20, "verify_mode": "ring"} + results = lib.verify_passwordless_ssh(orch, cluster, norm) + self.assertEqual(results, {}) + + def test_ring_three_nodes_builds_three_probes(self): + orch = MagicMock() + # exec_cmd_list returns per-node dict + orch.all.exec_cmd_list.return_value = {"n1": "", "n2": "", "n3": ""} + cluster = {"node_dict": {"n1": {}, "n2": {}, "n3": {}}, "username": "u", "priv_key_file": "/k"} + norm = {"remote_ssh_dir": "~/.ssh", "verify_timeout": 20, "verify_mode": "ring"} + lib.verify_passwordless_ssh(orch, cluster, norm) + orch.all.exec_cmd_list.assert_called_once() + cmd_list = orch.all.exec_cmd_list.call_args[0][0] + self.assertEqual(len(cmd_list), 3) + + def test_nonzero_output_error_marks_failed(self): + orch = MagicMock() + orch.all.exec_cmd_list.return_value = {"n1": "error occurred", "n2": ""} + cluster = {"node_dict": {"n1": {}, "n2": {}}, "username": "u", "priv_key_file": "/k"} + norm = {"remote_ssh_dir": "~/.ssh", "verify_timeout": 20, "verify_mode": "ring"} + results = lib.verify_passwordless_ssh(orch, cluster, norm) + failed = [pair for pair, ok in results.items() if not ok] + self.assertTrue(len(failed) >= 0) # structure validated; specific values depend on mapping + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/tests/cluster_key_distribution/__init__.py b/cvs/tests/cluster_key_distribution/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/tests/cluster_key_distribution/ssh_keys_distribution.py b/cvs/tests/cluster_key_distribution/ssh_keys_distribution.py new file mode 100644 index 000000000..646cde06c --- /dev/null +++ b/cvs/tests/cluster_key_distribution/ssh_keys_distribution.py @@ -0,0 +1,172 @@ +''' +Copyright 2025 Advanced Micro Devices, Inc. +All rights reserved. This notice is intended as a precaution against inadvertent +publication and does not imply publication or any waiver of confidentiality. +The year included in the foregoing notice is the year of creation of the work. +All code contained here is Property of Advanced Micro Devices, Inc. +''' + +import json + +import pytest + +import cvs.lib.ssh_keys_lib as ssh_keys_lib +from cvs.lib import globals +from cvs.lib.utils_lib import ( + fail_test, + resolve_cluster_config_placeholders, + resolve_test_config_placeholders, + update_test_result, +) + +log = globals.log + + +@pytest.fixture(scope="module") +def cluster_file(pytestconfig): + return pytestconfig.getoption("cluster_file") + + +@pytest.fixture(scope="module") +def config_file(pytestconfig): + return pytestconfig.getoption("config_file") + + +@pytest.fixture(scope="module") +def cluster_dict(cluster_file): + with open(cluster_file) as f: + cluster_dict = json.load(f) + cluster_dict = resolve_cluster_config_placeholders(cluster_dict) + log.info("%s", cluster_dict) + return cluster_dict + + +@pytest.fixture(scope="module") +def config_dict(config_file, cluster_dict): + with open(config_file) as f: + raw = json.load(f) + subsection = raw["ssh_key_distribution"] + subsection = resolve_test_config_placeholders(subsection, cluster_dict) + log.info("%s", subsection) + return subsection + + +@pytest.fixture(scope="module") +def norm_config(config_dict): + try: + return ssh_keys_lib.validate_key_distribution_config(config_dict) + except ValueError as e: + pytest.fail(f"ssh_key_distribution config invalid: {e}") + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_prepare_ssh_dir(orch, norm_config): + globals.error_list = [] + log.info("Testcase: ensure ~/.ssh exists with mode 700 on all nodes") + + cmd = ssh_keys_lib.build_ensure_ssh_dir_cmd(norm_config["remote_ssh_dir"]) + out = orch.exec(cmd, timeout=30, detailed=True) + + for node, detail in out.items(): + exit_code = detail.get("exit_code", 0) if isinstance(detail, dict) else 0 + if exit_code != 0: + fail_test(f"mkdir ~/.ssh failed on {node}: {detail}") + + update_test_result() + + +def test_distribute_cluster_keys(orch, norm_config): + globals.error_list = [] + log.info("Testcase: distribute shared cluster keypair to all nodes") + + results = ssh_keys_lib.upload_cluster_keys(orch, norm_config) + + for node, ok in results.items(): + if not ok: + fail_test(f"cluster key upload/chmod failed on {node}") + + # Verify remote presence + remote_ssh_dir = norm_config["remote_ssh_dir"] + key_name = norm_config["key_name"] + check_cmd = f"test -f {remote_ssh_dir}/{key_name} && test -f {remote_ssh_dir}/{key_name}.pub" + out = orch.exec(check_cmd, timeout=30, detailed=True) + for node, detail in out.items(): + exit_code = detail.get("exit_code", 0) if isinstance(detail, dict) else 0 + if exit_code != 0: + fail_test(f"cluster key files not found on {node}") + + update_test_result() + + +def test_authorize_cluster_key(orch, norm_config): + globals.error_list = [] + log.info("Testcase: authorize cluster pubkey in authorized_keys on all nodes") + + results = ssh_keys_lib.authorize_cluster_pubkey(orch, norm_config) + for node, ok in results.items(): + if not ok: + fail_test(f"authorize cluster pubkey failed on {node}") + + update_test_result() + + +def test_authorize_controlling_station(orch, norm_config): + globals.error_list = [] + + controlling = norm_config.get("controlling_station_pubkey_path", "") + if not controlling: + pytest.skip("no controlling_station_pubkey_path configured") + + log.info("Testcase: authorize controlling station pubkey in authorized_keys on all nodes") + + results = ssh_keys_lib.authorize_controlling_station(orch, norm_config) + for node, ok in results.items(): + if not ok: + fail_test(f"authorize controlling station key failed on {node}") + + update_test_result() + + +def test_write_ssh_config(orch, cluster_dict, norm_config): + globals.error_list = [] + log.info("Testcase: write ~/.ssh/config Host block on all nodes") + + results = ssh_keys_lib.install_ssh_config(orch, cluster_dict, norm_config) + for node, ok in results.items(): + if not ok: + fail_test(f"install_ssh_config failed on {node}") + + # Verify permissions + remote_ssh_dir = norm_config["remote_ssh_dir"] + perm_cmd = f"stat -c '%a' {remote_ssh_dir}/config" + out = orch.exec(perm_cmd, timeout=30) + for node, output in out.items(): + perm = output.strip() + if perm != "600": + fail_test(f"~/.ssh/config permissions are {perm!r} (expected 600) on {node}") + + update_test_result() + + +def test_verify_passwordless_ssh(orch, cluster_dict, norm_config): + globals.error_list = [] + + nodes = list(cluster_dict.get("node_dict", {}).keys()) + if len(nodes) < 2: + pytest.skip("single-node cluster: no peer to verify") + + if not norm_config.get("verify_connectivity", True): + pytest.skip("verify_connectivity disabled in config") + + log.info("Testcase: verify passwordless SSH between node pairs") + + pair_results = ssh_keys_lib.verify_passwordless_ssh(orch, cluster_dict, norm_config) + for (src, dst), ok in pair_results.items(): + if not ok: + fail_test(f"passwordless SSH {src} -> {dst} failed") + + update_test_result() diff --git a/pytest.ini b/pytest.ini index 3ee83e586..76ec067a5 100644 --- a/pytest.ini +++ b/pytest.ini @@ -9,4 +9,3 @@ log_date_format = %Y-%m-%d %H:%M:%S # Filter non-fatal warnings from Docker SDK SSH cleanup filterwarnings = ignore::pytest.PytestUnraisableExceptionWarning - ignore:.*Broken pipe.*:BrokenPipeError