Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
15 kB · 356 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357import asyncioimport osimport unittestfrom pathlib import Pathfrom tempfile import TemporaryDirectoryfrom unittest.mock import AsyncMock, patch, sentinel
import sshimport 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()