import asyncio import os import unittest from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import AsyncMock, patch, sentinel import ssh import ssh_control class SSHTests(unittest.TestCase): def setUp(self): self.temp = TemporaryDirectory() self.previous = os.environ.get("PRIME_SSH_CONTROL_DIR") os.environ["PRIME_SSH_CONTROL_DIR"] = self.temp.name def tearDown(self): if self.previous is None: os.environ.pop("PRIME_SSH_CONTROL_DIR", None) else: os.environ["PRIME_SSH_CONTROL_DIR"] = self.previous self.temp.cleanup() def test_control_path_is_stable_short_and_transport_specific(self): first = ssh.control_path_for("user@example.test") second = ssh.control_path_for("user@example.test") alternate = ssh.control_path_for("user@example.test", port=2222) self.assertEqual(first, second) self.assertNotEqual(first, alternate) self.assertLess(len(os.fsencode(first)), 100) self.assertEqual(Path(self.temp.name).stat().st_mode & 0o777, 0o700) connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) master_args = ssh._master_start_args(connection) self.assertIn("ControlMaster=yes", master_args) self.assertNotIn("-M", master_args) self.assertNotIn("-MNf", master_args) self.assertIn("-Nf", master_args) self.assertIn("ProxyCommand=false", ssh._reuse_only_args()) def test_one_shot_ssh_with_no_prior_call_succeeds(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) missing = ssh._Result(("ssh",), 255, "", "missing", connection.control_path) live = ssh._Result(("ssh",), 0, "", "", connection.control_path) started = ssh._Result(("ssh",), 0, "", "", connection.control_path) with ( patch.object(ssh, "_control", AsyncMock(side_effect=[missing, live])), patch.object(ssh, "_execute", AsyncMock(side_effect=lambda *a, **kw: (kw["control_path"].touch(), started)[1])) as execute, patch.object(ssh, "_bash", return_value=sentinel.handle) as bash_mock, ): handle = asyncio.run(ssh.ssh("user@example.test", "uptime")) self.assertIs(handle, sentinel.handle) execute.assert_awaited_once() bash_cmd = bash_mock.call_args.args[0] self.assertIn("ssh", bash_cmd) self.assertIn(str(connection.control_path), bash_cmd) self.assertIn("user@example.test", bash_cmd) self.assertIn("uptime", bash_cmd) def test_connection_reuse_across_commands(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) missing = ssh._Result(("ssh",), 255, "", "missing", connection.control_path) live = ssh._Result(("ssh",), 0, "", "", connection.control_path) started = ssh._Result(("ssh",), 0, "", "", connection.control_path) with ( patch.object(ssh, "_control", AsyncMock(side_effect=[missing, live])), patch.object(ssh, "_execute", AsyncMock(side_effect=lambda *a, **kw: (kw["control_path"].touch(), started)[1])) as execute, ): c = asyncio.run(ssh.ssh("user@example.test")) self.assertIsInstance(c, ssh.Connection) self.assertEqual(c.control_path, connection.control_path) execute.assert_awaited_once() c.control_path.touch() with patch.object(ssh, "_bash", return_value=sentinel.handle) as bash_mock: h1 = c.run("hostname") h2 = c.run("uptime") self.assertIs(h1, sentinel.handle) self.assertIs(h2, sentinel.handle) self.assertEqual(bash_mock.call_count, 2) cmd1 = bash_mock.call_args_list[0].args[0] cmd2 = bash_mock.call_args_list[1].args[0] self.assertIn(f"-S {c.control_path}", cmd1) self.assertIn(f"-S {c.control_path}", cmd2) def test_alias_imports_and_exposes_same_behavior(self): self.assertIs(ssh_control.Connection, ssh.Connection) self.assertIs(ssh_control.SSHControlError, ssh.SSHControlError) self.assertIs(ssh_control.SSHTimeoutError, ssh.SSHTimeoutError) self.assertIs(ssh_control.control_path_for, ssh.control_path_for) self.assertIs(ssh_control.close, ssh.close) self.assertIs(ssh_control.run, ssh.run) self.assertIs(ssh_control.ssh, ssh.ssh) self.assertIs(ssh_control.ssh_control, ssh.ssh) connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) missing = ssh._Result(("ssh",), 255, "", "missing", connection.control_path) live = ssh._Result(("ssh",), 0, "", "", connection.control_path) started = ssh._Result(("ssh",), 0, "", "", connection.control_path) with ( patch.object(ssh, "_control", AsyncMock(side_effect=[missing, live])), patch.object(ssh, "_execute", AsyncMock(side_effect=lambda *a, **kw: (kw["control_path"].touch(), started)[1])), patch.object(ssh, "_bash", return_value=sentinel.handle), ): handle = asyncio.run(ssh_control.run("user@example.test", "uptime")) self.assertIs(handle, sentinel.handle) with ( patch.object(ssh, "_control", AsyncMock(return_value=live)), patch.object(ssh, "_execute", AsyncMock()), ): conn = asyncio.run(ssh_control("user@example.test")) self.assertIsInstance(conn, ssh.Connection) def test_unreachable_host_fails_with_clear_error_not_ssh_255(self): connection = ssh._connection( "unreachable.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) missing = ssh._Result(("ssh",), 255, "", "missing", connection.control_path) failed_start = ssh._Result( ("ssh",), 255, "", "ssh: connect to host unreachable.test port 22: Connection refused", connection.control_path, ) with ( patch.object(ssh, "_control", AsyncMock(return_value=missing)), patch.object(ssh, "_execute", AsyncMock(return_value=failed_start)), ): with self.assertRaises(ssh.SSHControlError) as ctx: asyncio.run(ssh.ssh("unreachable.test", "uptime")) err = str(ctx.exception) self.assertIn("unreachable.test", err) self.assertIn("Connection refused", err) self.assertNotIn("exit 255", err) def test_rejects_option_and_destination_injection(self): with self.assertRaises(ValueError): ssh.control_path_for("-oProxyCommand=bad") with self.assertRaises(ValueError): ssh.control_path_for("host", options={"ControlPath": "/tmp/other"}) with self.assertRaises(ValueError): ssh.control_path_for("host", options={"ProxyCommand": "other"}) with self.assertRaises(ValueError): asyncio.run(ssh.ssh("host name", "true")) def test_connection_commands_return_native_bash_handle(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) connection.control_path.touch() with patch.object(ssh, "_bash", return_value=sentinel.handle) as start: handle = connection.run("printf '%s\n' hello") self.assertIs(handle, sentinel.handle) self.assertEqual(handle.command, "ssh user@example.test: printf '%s\n' hello") command = start.call_args.args[0] self.assertIn("ssh", command) self.assertIn(str(connection.control_path), command) self.assertIn("user@example.test", command) self.assertIn("printf", command) def test_run_forwards_stdin_to_the_native_handle(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) connection.control_path.touch() with patch.object(ssh, "_bash", return_value=sentinel.handle) as start: connection.run("cat > out", input="hello") self.assertEqual(start.call_args_list[0].kwargs["input"], "hello") live = ssh._Result(("ssh",), 0, "", "", connection.control_path) with ( patch.object(ssh, "_control", AsyncMock(return_value=live)), patch.object(ssh, "_bash", return_value=sentinel.handle) as start, ): asyncio.run(ssh.ssh("user@example.test", "cat > out", input=b"\x00 binary")) self.assertEqual(start.call_args_list[0].kwargs["input"], b"\x00 binary") def test_copy_and_close_return_native_bash_handles(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) connection.control_path.touch() with patch.object( ssh, "_bash", side_effect=[sentinel.upload, sentinel.download, sentinel.close], ) as start: upload = connection.copy_to("local file", "/tmp/remote file") download = connection.copy_from("/tmp/remote file", "local file") close = connection.close() self.assertIs(upload, sentinel.upload) self.assertIs(download, sentinel.download) self.assertIs(close, sentinel.close) self.assertEqual(upload.command, "scp local file -> user@example.test:/tmp/remote file") self.assertEqual(download.command, "scp user@example.test:/tmp/remote file -> local file") self.assertEqual(close.command, "ssh -O exit user@example.test") self.assertEqual(start.call_count, 3) self.assertIn("scp", start.call_args_list[0].args[0]) self.assertIn("ssh", start.call_args_list[2].args[0]) self.assertIn("-O exit", start.call_args_list[2].args[0]) def test_connection_run_and_copy_refuse_when_control_master_socket_is_missing(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) self.assertFalse(connection.control_path.exists()) with self.assertRaises(ssh.SSHControlError) as ctx: connection.run("hostname") self.assertIn("user@example.test", str(ctx.exception)) self.assertIn("await ssh", str(ctx.exception)) with self.assertRaises(ssh.SSHControlError) as ctx: connection.copy_to("local_file", "/tmp/remote") self.assertIn("user@example.test", str(ctx.exception)) self.assertIn("await ssh", str(ctx.exception)) with self.assertRaises(ssh.SSHControlError) as ctx: connection.copy_from("/tmp/remote", "local_file") self.assertIn("user@example.test", str(ctx.exception)) self.assertIn("await ssh", str(ctx.exception)) def test_socket_identity_consistent_across_options(self): path_default = ssh.control_path_for("user@example.test") path_with_timeout = ssh.control_path_for( "user@example.test", options={"ConnectTimeout": "10"} ) path_with_other = ssh.control_path_for( "user@example.test", options={"Compression": "yes"} ) self.assertEqual(path_default, path_with_timeout) self.assertEqual(path_default, path_with_other) conn_without = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) conn_with = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options={"ConnectTimeout": "10"}, persist="30m", ) self.assertEqual(conn_without.control_path, conn_with.control_path) master_args = ssh._master_start_args(conn_with) self.assertIn("ConnectTimeout=10", master_args) def test_copy_to_and_from_verified_command_structure(self): connection = ssh._connection( "user@example.test", port=None, identity_file=None, proxy_jump=None, options=None, persist="30m", ) connection.control_path.touch() with patch.object(ssh, "_bash", return_value=sentinel.handle) as mock_bash: handle = connection.copy_to("local.bin", "/remote/path.bin", verify=True, sha256=True) self.assertIs(handle, sentinel.handle) cmd = mock_bash.call_args[0][0] self.assertIn("scp", cmd) self.assertIn("wc -c", cmd) self.assertIn("shasum -a 256", cmd) self.assertIn("copy_to verification failed: size mismatch", cmd) self.assertIn("copy_to verification failed: sha256 mismatch", cmd) with patch.object(ssh, "_bash", return_value=sentinel.handle) as mock_bash: handle = connection.copy_from("/remote/path.bin", "local.bin", verify=True, sha256=True) self.assertIs(handle, sentinel.handle) cmd = mock_bash.call_args[0][0] self.assertIn("scp", cmd) self.assertIn("wc -c", cmd) self.assertIn("shasum -a 256", cmd) self.assertIn("copy_from verification failed: size mismatch", cmd) self.assertIn("copy_from verification failed: sha256 mismatch", cmd) def test_copy_and_run_docstrings_state_handle_is_not_completion(self): self.assertIn("NOT command completion", ssh.Connection.run.__doc__) self.assertIn("NOT copy completion", ssh.Connection.copy_to.__doc__) self.assertIn("NOT copy completion", ssh.Connection.copy_from.__doc__) if __name__ == "__main__": unittest.main()