300 lines
11 KiB
Python
300 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
import stat
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
|
|
from herdr_hub.remotes import (
|
|
REMOTE_DIRECT_COMMAND,
|
|
REMOTE_LIST_SCRIPT,
|
|
RegistryError,
|
|
RemoteEntry,
|
|
RemoteQueryPool,
|
|
RemoteRegistry,
|
|
RemoteTargetError,
|
|
delete_remote_session,
|
|
decode_remote_output,
|
|
parse_remote_session_result,
|
|
remote_delete_session_command,
|
|
remote_session_command,
|
|
remote_stop_session_command,
|
|
ssh_destination_arguments,
|
|
stop_remote_session,
|
|
)
|
|
|
|
|
|
class RemoteRegistryTests(unittest.TestCase):
|
|
def test_round_trip_update_remove_and_permissions(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "config" / "remotes.json"
|
|
registry = RemoteRegistry(path)
|
|
self.assertEqual(registry.load(), [])
|
|
self.assertTrue(registry.upsert(RemoteEntry("orangepi")))
|
|
self.assertFalse(registry.upsert(RemoteEntry("orangepi")))
|
|
self.assertTrue(registry.upsert(RemoteEntry("prod", "server", "ssh")))
|
|
self.assertEqual(
|
|
registry.load(),
|
|
[RemoteEntry("orangepi"), RemoteEntry("prod", "server", "ssh")],
|
|
)
|
|
self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600)
|
|
self.assertTrue(registry.remove("orangepi"))
|
|
self.assertFalse(registry.remove("missing"))
|
|
self.assertEqual(
|
|
registry.load(),
|
|
[RemoteEntry("prod", "server", "ssh")],
|
|
)
|
|
|
|
def test_discovery_does_not_overwrite_preferences_and_ssh_is_sticky(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
|
registry.upsert(RemoteEntry("host", "server"))
|
|
registry.remember_many([RemoteEntry("host", "local")])
|
|
self.assertEqual(registry.load(), [RemoteEntry("host", "server")])
|
|
|
|
registry.remember_many([RemoteEntry("host", "local", "ssh")])
|
|
self.assertEqual(
|
|
registry.load(),
|
|
[RemoteEntry("host", "server", "ssh")],
|
|
)
|
|
|
|
def test_corrupt_registry_is_not_overwritten(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "remotes.json"
|
|
path.write_text("not json", encoding="utf-8")
|
|
registry = RemoteRegistry(path)
|
|
with self.assertRaises(RegistryError):
|
|
registry.upsert(RemoteEntry("orangepi"))
|
|
self.assertEqual(path.read_text(encoding="utf-8"), "not json")
|
|
|
|
def test_legacy_registry_defaults_to_native_transport(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "remotes.json"
|
|
path.write_text(
|
|
json.dumps(
|
|
{
|
|
"version": 1,
|
|
"remotes": [
|
|
{"target": "legacy", "remote_keybindings": "server"}
|
|
],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
self.assertEqual(
|
|
RemoteRegistry(path).load(),
|
|
[RemoteEntry("legacy", "server", "native")],
|
|
)
|
|
|
|
|
|
class RemoteSshTests(unittest.TestCase):
|
|
def test_alias_and_url_targets_become_safe_ssh_arguments(self) -> None:
|
|
self.assertEqual(ssh_destination_arguments("orangepi"), ["orangepi"])
|
|
self.assertEqual(
|
|
ssh_destination_arguments("ssh://me@example.com:2222"),
|
|
["-p", "2222", "-l", "me", "example.com"],
|
|
)
|
|
command = remote_session_command("/usr/bin/ssh", RemoteEntry("host name"))
|
|
self.assertEqual(command[0], "/usr/bin/ssh")
|
|
self.assertIn("host name", command)
|
|
self.assertEqual(command[-1], REMOTE_DIRECT_COMMAND)
|
|
fallback = remote_session_command(
|
|
"/usr/bin/ssh",
|
|
RemoteEntry("orangepi"),
|
|
posix_fallback=True,
|
|
)
|
|
self.assertEqual(fallback[-1], REMOTE_LIST_SCRIPT)
|
|
self.assertNotIn("sh", command[:1])
|
|
with self.assertRaises(RemoteTargetError):
|
|
ssh_destination_arguments("ssh://me:secret@example.com")
|
|
with self.assertRaises(RemoteTargetError):
|
|
ssh_destination_arguments("ssh://example.com/a/path")
|
|
|
|
def test_delete_stopped_remote_session_uses_safe_command_and_posix_fallback(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
fake_ssh = Path(directory) / "ssh"
|
|
fake_ssh.write_text(
|
|
"#!/bin/sh\n"
|
|
"last=''\n"
|
|
"for argument do last=$argument; done\n"
|
|
"case \"$last\" in\n"
|
|
" *\"if [ -x\"*)\n"
|
|
" printf '%s\\n' '{\"deleted\":true}'\n"
|
|
" exit 0\n"
|
|
" ;;\n"
|
|
"esac\n"
|
|
"exit 127\n",
|
|
encoding="utf-8",
|
|
)
|
|
fake_ssh.chmod(0o755)
|
|
delete_remote_session(
|
|
str(fake_ssh),
|
|
RemoteEntry("orangepi"),
|
|
"old-session",
|
|
timeout=1,
|
|
)
|
|
|
|
command = remote_delete_session_command(
|
|
"/usr/bin/ssh",
|
|
RemoteEntry("orangepi"),
|
|
"old-session",
|
|
)
|
|
self.assertEqual(
|
|
command[-1],
|
|
"herdr session delete --json -- old-session",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
remote_delete_session_command(
|
|
"/usr/bin/ssh",
|
|
RemoteEntry("orangepi"),
|
|
"bad session",
|
|
)
|
|
|
|
def test_stop_running_remote_session_uses_safe_command_and_posix_fallback(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
fake_ssh = Path(directory) / "ssh"
|
|
fake_ssh.write_text(
|
|
"#!/bin/sh\n"
|
|
"last=''\n"
|
|
"for argument do last=$argument; done\n"
|
|
"case \"$last\" in\n"
|
|
" *\"if [ -x\"*) exit 0 ;;\n"
|
|
"esac\n"
|
|
"exit 127\n",
|
|
encoding="utf-8",
|
|
)
|
|
fake_ssh.chmod(0o755)
|
|
stop_remote_session(
|
|
str(fake_ssh),
|
|
RemoteEntry("orangepi"),
|
|
"work",
|
|
timeout=1,
|
|
)
|
|
|
|
command = remote_stop_session_command(
|
|
"/usr/bin/ssh",
|
|
RemoteEntry("orangepi"),
|
|
"work",
|
|
)
|
|
self.assertEqual(command[-1], "herdr session stop -- work")
|
|
with self.assertRaises(ValueError):
|
|
remote_stop_session_command(
|
|
"/usr/bin/ssh",
|
|
RemoteEntry("orangepi"),
|
|
"bad session",
|
|
)
|
|
|
|
def test_remote_session_json_keeps_running_and_stopped(self) -> None:
|
|
result = parse_remote_session_result(
|
|
RemoteEntry("orangepi"),
|
|
0,
|
|
'{"sessions":['
|
|
'{"name":"old","running":false},'
|
|
'{"name":"default","running":true}]}',
|
|
"",
|
|
)
|
|
self.assertIsNone(result.error)
|
|
self.assertEqual(
|
|
[(session.name, session.running) for session in result.sessions],
|
|
[("default", True), ("old", False)],
|
|
)
|
|
|
|
windows = parse_remote_session_result(
|
|
RemoteEntry("windows"),
|
|
0,
|
|
'{"sessions":[{"name":"default","running":true,'
|
|
'"session_dir":"C:\\\\Users\\\\me\\\\AppData\\\\Roaming\\\\herdr"}]}',
|
|
"",
|
|
)
|
|
self.assertEqual(windows.platform, "windows")
|
|
|
|
def test_query_pool_reports_results_and_spawn_failures(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
fake_ssh = Path(directory) / "ssh"
|
|
fake_ssh.write_text(
|
|
"#!/bin/sh\n"
|
|
"printf '%s\\n' '{\"sessions\":[{\"name\":\"work\",\"running\":true}]}'\n",
|
|
encoding="utf-8",
|
|
)
|
|
fake_ssh.chmod(0o755)
|
|
pool = RemoteQueryPool(str(fake_ssh), query_timeout=1)
|
|
self.addCleanup(pool.close)
|
|
pool.enqueue(RemoteEntry("orangepi"))
|
|
results = []
|
|
deadline = time.monotonic() + 2
|
|
while not results and time.monotonic() < deadline:
|
|
results = pool.poll()
|
|
time.sleep(0.01)
|
|
self.assertEqual(results[0].sessions[0].name, "work")
|
|
|
|
missing = RemoteQueryPool(str(Path(directory) / "missing"))
|
|
self.addCleanup(missing.close)
|
|
missing.enqueue(RemoteEntry("orangepi"))
|
|
failure = missing.poll()
|
|
self.assertEqual(len(failure), 1)
|
|
self.assertIn("cannot start", failure[0].error or "")
|
|
|
|
def test_query_pool_falls_back_for_posix_path_and_decodes_windows_errors(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
fallback_ssh = root / "fallback-ssh"
|
|
fallback_ssh.write_text(
|
|
"#!/bin/sh\n"
|
|
"last=''\n"
|
|
"for argument do last=$argument; done\n"
|
|
"case \"$last\" in\n"
|
|
" *\"if [ -x\"*)\n"
|
|
" printf '%s\\n' '{\"sessions\":[{\"name\":\"work\",\"running\":true}]}'\n"
|
|
" exit 0\n"
|
|
" ;;\n"
|
|
"esac\n"
|
|
"exit 127\n",
|
|
encoding="utf-8",
|
|
)
|
|
fallback_ssh.chmod(0o755)
|
|
pool = RemoteQueryPool(str(fallback_ssh), query_timeout=1)
|
|
self.addCleanup(pool.close)
|
|
pool.enqueue(RemoteEntry("linux-host"))
|
|
results = []
|
|
deadline = time.monotonic() + 2
|
|
while not results and time.monotonic() < deadline:
|
|
results = pool.poll()
|
|
time.sleep(0.01)
|
|
self.assertEqual(results[0].sessions[0].name, "work")
|
|
|
|
windows_ssh = root / "windows-ssh"
|
|
windows_ssh.write_bytes(
|
|
b"#!/bin/sh\nprintf '\\264\\355\\316\\363' >&2\nexit 1\n"
|
|
)
|
|
windows_ssh.chmod(0o755)
|
|
windows_pool = RemoteQueryPool(str(windows_ssh), query_timeout=1)
|
|
self.addCleanup(windows_pool.close)
|
|
windows_pool.enqueue(RemoteEntry("windows-host"))
|
|
windows_results = []
|
|
deadline = time.monotonic() + 2
|
|
while not windows_results and time.monotonic() < deadline:
|
|
windows_results = windows_pool.poll()
|
|
time.sleep(0.01)
|
|
self.assertIn("错误", windows_results[0].error or "")
|
|
|
|
self.assertEqual(decode_remote_output("错误".encode("gbk")), "错误")
|
|
|
|
def test_query_pool_times_out_without_blocking_other_results(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
slow_ssh = Path(directory) / "ssh"
|
|
slow_ssh.write_text("#!/bin/sh\nexec sleep 5\n", encoding="utf-8")
|
|
slow_ssh.chmod(0o755)
|
|
pool = RemoteQueryPool(str(slow_ssh), query_timeout=0.05)
|
|
self.addCleanup(pool.close)
|
|
pool.enqueue(RemoteEntry("slow"))
|
|
time.sleep(0.08)
|
|
results = pool.poll()
|
|
self.assertEqual(len(results), 1)
|
|
self.assertIn("timed out", results[0].error or "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|