@neoline/hostbridge 2.0.3 → 2.0.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +15 -8
- package/bin/hostbridge.js +80 -9
- package/examples/claude_desktop_config.json +1 -1
- package/examples/codex.config.toml +1 -1
- package/examples/hosts.example.json +1 -0
- package/package.json +1 -1
- package/pyproject.toml +5 -4
- package/src/server_control_mcp/__init__.py +58 -23
- package/src/server_control_mcp/cli.py +1 -1
- package/src/server_control_mcp/client.py +99 -7
- package/src/server_control_mcp/config.py +16 -5
- package/src/server_control_mcp/daemon.py +239 -49
- package/src/server_control_mcp/doctor.py +8 -4
- package/src/server_control_mcp/hosts.py +118 -24
- package/src/server_control_mcp/mock_ssh.py +53 -954
- package/src/server_control_mcp/mock_ssh_exec.py +210 -0
- package/src/server_control_mcp/mock_ssh_session.py +68 -0
- package/src/server_control_mcp/mock_ssh_sftp.py +506 -0
- package/src/server_control_mcp/mock_ssh_tunnel.py +592 -0
- package/src/server_control_mcp/policy.py +86 -14
- package/src/server_control_mcp/protocol.py +18 -17
- package/src/server_control_mcp/remote_task_agent.py +102 -0
- package/src/server_control_mcp/remote_tunnel_agent.py +143 -0
- package/src/server_control_mcp/runtime.py +3 -2
- package/src/server_control_mcp/secrets.py +6 -6
- package/src/server_control_mcp/secure_files.py +121 -0
- package/src/server_control_mcp/server.py +30 -12
- package/src/server_control_mcp/services.py +363 -165
- package/src/server_control_mcp/transports/pty.py +114 -33
- package/src/server_control_mcp/transports/ssh.py +308 -24
|
@@ -20,6 +20,7 @@ from ..hosts import HostConfig
|
|
|
20
20
|
from .base import ExecResult, ShellReadResult, TransferResult, TransportBusy, TransportCapabilities, TransportError
|
|
21
21
|
|
|
22
22
|
_REMOTE_BASE64_DECODE = "{ base64 --decode 2>/dev/null || base64 -d 2>/dev/null || base64 -D; }"
|
|
23
|
+
_REMOTE_BASE64_ENCODE = "{ base64 --wrap=0 2>/dev/null || base64 | tr -d '\\r\\n'; }"
|
|
23
24
|
PTY_BASE64_LINE_LENGTH = 3072
|
|
24
25
|
|
|
25
26
|
|
|
@@ -82,6 +83,28 @@ def _marker_text(value: object) -> str:
|
|
|
82
83
|
return value.group(0) if hasattr(value, "group") else str(value)
|
|
83
84
|
|
|
84
85
|
|
|
86
|
+
def _private_temp_path(kind: str, token: str, name: str | None = None) -> str:
|
|
87
|
+
path = f"$HOME/.hostbridge/tmp/{kind}.{token}"
|
|
88
|
+
return f"{path}/{name}" if name is not None else path
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def _quote_private_temp_path(path: str) -> str:
|
|
92
|
+
prefix = "$HOME/"
|
|
93
|
+
if not path.startswith(prefix):
|
|
94
|
+
raise ValueError("private temporary path must be relative to $HOME")
|
|
95
|
+
return f'"$HOME"/{shlex.quote(path[len(prefix) :])}'
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def _prepare_private_temp_dir(path: str) -> str:
|
|
99
|
+
directory = _quote_private_temp_path(path)
|
|
100
|
+
return (
|
|
101
|
+
'umask 077; tmp_root="$HOME/.hostbridge/tmp"; '
|
|
102
|
+
'[ -n "${HOME:-}" ] && [ ! -L "$HOME/.hostbridge" ] && [ ! -L "$tmp_root" ] && '
|
|
103
|
+
'mkdir -p -- "$tmp_root" && chmod 700 "$HOME/.hostbridge" "$tmp_root" && '
|
|
104
|
+
f"mkdir -m 700 -- {directory}"
|
|
105
|
+
)
|
|
106
|
+
|
|
107
|
+
|
|
85
108
|
class PtyTransport:
|
|
86
109
|
capabilities = TransportCapabilities(
|
|
87
110
|
separate_stderr=False,
|
|
@@ -167,8 +190,9 @@ class PtyTransport:
|
|
|
167
190
|
finally:
|
|
168
191
|
if stdin_path is not None and self.child.isalive():
|
|
169
192
|
with suppress(Exception):
|
|
193
|
+
stdin_dir = stdin_path.rsplit("/", 1)[0]
|
|
170
194
|
self._exec_locked(
|
|
171
|
-
f"trap - EXIT HUP INT TERM; rm -
|
|
195
|
+
f"trap - EXIT HUP INT TERM; rm -rf -- {_quote_private_temp_path(stdin_dir)}",
|
|
172
196
|
timeout=min(timeout, 10),
|
|
173
197
|
output_limit=1024,
|
|
174
198
|
)
|
|
@@ -176,19 +200,33 @@ class PtyTransport:
|
|
|
176
200
|
|
|
177
201
|
def _stage_stdin_locked(self, chunks: Iterable[bytes], *, timeout: float) -> str:
|
|
178
202
|
token = self._token_factory()
|
|
179
|
-
|
|
180
|
-
|
|
203
|
+
temp_dir = _private_temp_path("stdin", token)
|
|
204
|
+
path = _private_temp_path("stdin", token, "stdin")
|
|
205
|
+
base64_path = _private_temp_path("stdin", token, "stdin.b64")
|
|
206
|
+
temp_dir_q = _quote_private_temp_path(temp_dir)
|
|
207
|
+
path_q = _quote_private_temp_path(path)
|
|
208
|
+
base64_path_q = _quote_private_temp_path(base64_path)
|
|
181
209
|
heredoc = f"__{self._marker_prefix}_STDIN_{token}__"
|
|
210
|
+
ready_pattern = re.compile(rf"__{self._marker_prefix}_STDIN_READY_{re.escape(token)}__:(\d+)")
|
|
182
211
|
done_pattern = re.compile(rf"__{self._marker_prefix}_STDIN_DONE_{re.escape(token)}__:(\d+)")
|
|
183
|
-
cleanup = f"rm -
|
|
212
|
+
cleanup = f"rm -rf -- {temp_dir_q}"
|
|
184
213
|
try:
|
|
185
|
-
|
|
214
|
+
setup = _prepare_private_temp_dir(temp_dir)
|
|
215
|
+
self.child.sendline(
|
|
216
|
+
f"if {setup}; then status=0; else status=$?; fi; "
|
|
217
|
+
f"printf '__{self._marker_prefix}_STDIN_READY_{token}__:%s\n' \"$status\""
|
|
218
|
+
)
|
|
219
|
+
self.child.expect(ready_pattern, timeout=timeout)
|
|
220
|
+
ready_match = re.search(r":(\d+)", _marker_text(self.child.after))
|
|
221
|
+
if ready_match is None or int(ready_match.group(1)) != 0:
|
|
222
|
+
raise TransportError("failed to prepare private PTY stdin directory")
|
|
223
|
+
self.child.send(f"trap {shlex.quote(cleanup)} EXIT HUP INT TERM; cat > {path_q} <<'{heredoc}'\n")
|
|
186
224
|
for encoded in _iter_base64(chunks):
|
|
187
225
|
self.child.send(encoded)
|
|
188
226
|
self.child.send(
|
|
189
227
|
f"\n{heredoc}\n"
|
|
190
|
-
f'tmp={
|
|
191
|
-
f'{{ {_REMOTE_BASE64_DECODE} < "$tmp"; }} > {
|
|
228
|
+
f'tmp={base64_path_q}; mv {path_q} "$tmp"; '
|
|
229
|
+
f'{{ {_REMOTE_BASE64_DECODE} < "$tmp"; }} > {path_q}; status=$?; rm -f "$tmp"; '
|
|
192
230
|
f"printf '\n__{self._marker_prefix}_STDIN_DONE_{token}__:%s\n' \"$status\"\n"
|
|
193
231
|
)
|
|
194
232
|
self.child.expect(done_pattern, timeout=timeout)
|
|
@@ -198,6 +236,8 @@ class PtyTransport:
|
|
|
198
236
|
self.child.send("\x03")
|
|
199
237
|
with suppress(Exception):
|
|
200
238
|
self.child.sendline(f"trap - EXIT HUP INT TERM; {cleanup}")
|
|
239
|
+
with suppress(Exception):
|
|
240
|
+
self.child.close(force=True)
|
|
201
241
|
raise
|
|
202
242
|
match = re.search(r":(\d+)", _marker_text(self.child.after))
|
|
203
243
|
if match is None or int(match.group(1)) != 0:
|
|
@@ -216,22 +256,54 @@ class PtyTransport:
|
|
|
216
256
|
start_marker = f"__{self._marker_prefix}_START_{token}__"
|
|
217
257
|
exit_pattern = re.compile(rf"__{self._marker_prefix}_EXIT_{re.escape(token)}__:(\d+)")
|
|
218
258
|
encoded_command = base64.b64encode(command.encode("utf-8")).decode("ascii")
|
|
219
|
-
stdin_redirect = f" < {
|
|
259
|
+
stdin_redirect = f" < {_quote_private_temp_path(stdin_path)}" if stdin_path is not None else ""
|
|
260
|
+
status_dir = _private_temp_path("exec", token)
|
|
261
|
+
status_path = _private_temp_path("exec", token, "status")
|
|
262
|
+
cleanup_dirs = [status_dir]
|
|
263
|
+
if stdin_path is not None:
|
|
264
|
+
cleanup_dirs.append(stdin_path.rsplit("/", 1)[0])
|
|
265
|
+
signal_cleanup = "rm -rf -- " + " ".join(_quote_private_temp_path(path) for path in cleanup_dirs)
|
|
266
|
+
remaining_trap = (
|
|
267
|
+
f"trap {shlex.quote('rm -rf -- ' + _quote_private_temp_path(cleanup_dirs[1]))} EXIT HUP INT TERM"
|
|
268
|
+
if len(cleanup_dirs) > 1
|
|
269
|
+
else "trap - EXIT HUP INT TERM"
|
|
270
|
+
)
|
|
271
|
+
status_path_q = _quote_private_temp_path(status_path)
|
|
272
|
+
status_dir_q = _quote_private_temp_path(status_dir)
|
|
273
|
+
setup = _prepare_private_temp_dir(status_dir)
|
|
220
274
|
remote_line = (
|
|
275
|
+
f"if {setup}; then trap {shlex.quote(signal_cleanup)} EXIT HUP INT TERM; "
|
|
221
276
|
f"__scm_cmd=$(printf %s {shlex.quote(encoded_command)} | {_REMOTE_BASE64_DECODE}); "
|
|
222
|
-
f"printf '\n{start_marker}
|
|
223
|
-
f
|
|
224
|
-
"
|
|
225
|
-
f"
|
|
277
|
+
f"printf '\n{start_marker}'; "
|
|
278
|
+
f'{{ {self.remote_shell} -lc "$__scm_cmd"{stdin_redirect}; '
|
|
279
|
+
f'printf %s "$?" > {status_path_q}; }} 2>&1 | '
|
|
280
|
+
f"tail -c {output_limit} | {_REMOTE_BASE64_ENCODE}; "
|
|
281
|
+
f"__scm_status=$(cat {status_path_q} 2>/dev/null || printf 255); "
|
|
282
|
+
f"rm -rf -- {status_dir_q}; {remaining_trap}; "
|
|
283
|
+
f"printf '__{self._marker_prefix}_EXIT_{token}__:%s\n' \"$__scm_status\"; "
|
|
284
|
+
f"else printf '\n{start_marker}'; printf '__{self._marker_prefix}_EXIT_{token}__:255\n'; fi"
|
|
226
285
|
)
|
|
227
286
|
try:
|
|
228
287
|
self.child.sendline(remote_line)
|
|
229
288
|
self.child.expect(exit_pattern, timeout=timeout)
|
|
230
289
|
except pexpect.TIMEOUT:
|
|
231
|
-
|
|
290
|
+
output = b""
|
|
291
|
+
with suppress(TransportError):
|
|
292
|
+
output = self._output_after_marker(str(self.child.before), start_marker, output_limit)
|
|
293
|
+
with suppress(Exception):
|
|
294
|
+
self.child.send("\x03")
|
|
295
|
+
with suppress(Exception):
|
|
296
|
+
self.child.close(force=True)
|
|
297
|
+
return ExecResult(None, output, timed_out=True)
|
|
232
298
|
match = re.search(r":(\d+)", _marker_text(self.child.after))
|
|
233
299
|
exit_code = int(match.group(1)) if match else None
|
|
234
|
-
|
|
300
|
+
try:
|
|
301
|
+
output = self._output_after_marker(str(self.child.before), start_marker, output_limit)
|
|
302
|
+
except TransportError:
|
|
303
|
+
with suppress(Exception):
|
|
304
|
+
self.child.close(force=True)
|
|
305
|
+
raise
|
|
306
|
+
return ExecResult(exit_code, output)
|
|
235
307
|
|
|
236
308
|
@staticmethod
|
|
237
309
|
def _output_after_marker(raw: str, marker: str, output_limit: int) -> bytes:
|
|
@@ -239,13 +311,12 @@ class PtyTransport:
|
|
|
239
311
|
normalized = normalized.replace("\r\n", "\n").replace("\r", "\n")
|
|
240
312
|
marker_index = normalized.rfind(marker)
|
|
241
313
|
output = normalized[marker_index + len(marker) :] if marker_index >= 0 else normalized
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
return encoded[-output_limit:]
|
|
314
|
+
try:
|
|
315
|
+
encoded = "".join(output.split()).encode("ascii")
|
|
316
|
+
decoded = base64.b64decode(encoded, validate=True)
|
|
317
|
+
except (UnicodeEncodeError, binascii.Error, ValueError) as exc:
|
|
318
|
+
raise TransportError("PTY command output contains invalid base64 framing") from exc
|
|
319
|
+
return decoded[-output_limit:]
|
|
249
320
|
|
|
250
321
|
def upload(
|
|
251
322
|
self,
|
|
@@ -263,9 +334,14 @@ class PtyTransport:
|
|
|
263
334
|
if not self.operation_lock.acquire(blocking=False):
|
|
264
335
|
raise TransportBusy("PTY transport already has an active operation")
|
|
265
336
|
token = self._token_factory()
|
|
266
|
-
|
|
267
|
-
|
|
337
|
+
temp_dir = _private_temp_path("upload", token)
|
|
338
|
+
tmp_b64 = _private_temp_path("upload", token, "payload.b64")
|
|
339
|
+
tmp_out = _private_temp_path("upload", token, "payload.out")
|
|
340
|
+
temp_dir_q = _quote_private_temp_path(temp_dir)
|
|
341
|
+
tmp_b64_q = _quote_private_temp_path(tmp_b64)
|
|
342
|
+
tmp_out_q = _quote_private_temp_path(tmp_out)
|
|
268
343
|
heredoc = f"__{self._marker_prefix}_UPLOAD_{token}__"
|
|
344
|
+
ready_pattern = re.compile(rf"__{self._marker_prefix}_UPLOAD_READY_{re.escape(token)}__:(\d+)")
|
|
269
345
|
done_pattern = re.compile(rf"__{self._marker_prefix}_UPLOAD_DONE_{re.escape(token)}__:(\d+)")
|
|
270
346
|
digest = hashlib.sha256()
|
|
271
347
|
transferred = 0
|
|
@@ -282,27 +358,32 @@ class PtyTransport:
|
|
|
282
358
|
try:
|
|
283
359
|
if not self.child.isalive():
|
|
284
360
|
raise TransportError("PTY transport is not alive")
|
|
285
|
-
|
|
361
|
+
setup = _prepare_private_temp_dir(temp_dir)
|
|
362
|
+
self.child.sendline(
|
|
363
|
+
f"if {setup}; then status=0; else status=$?; fi; "
|
|
364
|
+
f"printf '__{self._marker_prefix}_UPLOAD_READY_{token}__:%s\n' \"$status\""
|
|
365
|
+
)
|
|
366
|
+
self.child.expect(ready_pattern, timeout=60)
|
|
367
|
+
ready_match = re.search(r":(\d+)", _marker_text(self.child.after))
|
|
368
|
+
if ready_match is None or int(ready_match.group(1)) != 0:
|
|
369
|
+
raise TransportError("failed to prepare private PTY upload directory")
|
|
370
|
+
self.child.send(f"cat > {tmp_b64_q} <<'{heredoc}'\n")
|
|
286
371
|
for encoded in _iter_base64(counted_chunks()):
|
|
287
372
|
self.child.send(encoded)
|
|
288
|
-
self.child.send(
|
|
289
|
-
f"\n{heredoc}\nprintf '\n__{self._marker_prefix}_UPLOAD_DONE_{token}__:%s\n' \"$?\"\n"
|
|
290
|
-
)
|
|
373
|
+
self.child.send(f"\n{heredoc}\nprintf '\n__{self._marker_prefix}_UPLOAD_DONE_{token}__:%s\n' \"$?\"\n")
|
|
291
374
|
self.child.expect(done_pattern, timeout=60)
|
|
292
375
|
match = re.search(r":(\d+)", _marker_text(self.child.after))
|
|
293
376
|
if match is None or int(match.group(1)) != 0:
|
|
294
377
|
raise TransportError("failed to stream PTY upload")
|
|
295
378
|
if expected_size is not None and transferred != expected_size:
|
|
296
|
-
raise TransportError(
|
|
297
|
-
f"upload size mismatch: expected {expected_size}, received {transferred}"
|
|
298
|
-
)
|
|
379
|
+
raise TransportError(f"upload size mismatch: expected {expected_size}, received {transferred}")
|
|
299
380
|
target_q = shlex.quote(remote_path)
|
|
300
381
|
finalize = (
|
|
301
382
|
"set -e; "
|
|
302
|
-
f"target={target_q}; tmp_b64={
|
|
383
|
+
f"target={target_q}; tmp_b64={tmp_b64_q}; tmp_out={tmp_out_q}; "
|
|
303
384
|
'mkdir -p "$(dirname "$target")"; '
|
|
304
385
|
f'{{ {_REMOTE_BASE64_DECODE} < "$tmp_b64"; }} > "$tmp_out"; '
|
|
305
|
-
f'chmod {format(mode, "04o")} "$tmp_out"; mv "$tmp_out" "$target"; rm -
|
|
386
|
+
f'chmod {format(mode, "04o")} "$tmp_out"; mv "$tmp_out" "$target"; rm -rf -- {temp_dir_q}; '
|
|
306
387
|
"if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
|
|
307
388
|
"else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
|
|
308
389
|
'bytes=$(wc -c < "$target" | tr -d " "); printf "BYTES=%s\\nSHA256=%s\\n" "$bytes" "$sha"'
|
|
@@ -320,7 +401,7 @@ class PtyTransport:
|
|
|
320
401
|
if self.child.isalive():
|
|
321
402
|
with suppress(Exception):
|
|
322
403
|
self._exec_locked(
|
|
323
|
-
f"rm -
|
|
404
|
+
f"rm -rf -- {temp_dir_q}",
|
|
324
405
|
timeout=10,
|
|
325
406
|
output_limit=1024,
|
|
326
407
|
)
|
|
@@ -4,15 +4,51 @@ from __future__ import annotations
|
|
|
4
4
|
|
|
5
5
|
import asyncio
|
|
6
6
|
import hashlib
|
|
7
|
+
import threading
|
|
7
8
|
import uuid
|
|
8
|
-
from collections
|
|
9
|
+
from collections import deque
|
|
10
|
+
from collections.abc import AsyncGenerator, Awaitable, Iterable, Iterator, Sequence
|
|
9
11
|
from contextlib import suppress
|
|
10
12
|
from dataclasses import dataclass
|
|
11
13
|
from pathlib import Path
|
|
14
|
+
from typing import TypeVar
|
|
12
15
|
|
|
13
16
|
import asyncssh
|
|
14
17
|
|
|
15
|
-
from .base import ExecResult, TransferResult, TransportCapabilities
|
|
18
|
+
from .base import ExecResult, ShellReadResult, TransferResult, TransportCapabilities, TransportError
|
|
19
|
+
|
|
20
|
+
T = TypeVar("T")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class _CombinedOutputBuffer:
|
|
24
|
+
def __init__(self, limit: int) -> None:
|
|
25
|
+
self._limit = limit
|
|
26
|
+
self._segments: deque[tuple[bool, bytes]] = deque()
|
|
27
|
+
self._size = 0
|
|
28
|
+
|
|
29
|
+
def append(self, chunk: bytes, *, stderr: bool) -> None:
|
|
30
|
+
if not chunk:
|
|
31
|
+
return
|
|
32
|
+
self._segments.append((stderr, bytes(chunk)))
|
|
33
|
+
self._size += len(chunk)
|
|
34
|
+
excess = self._size - self._limit
|
|
35
|
+
while excess > 0:
|
|
36
|
+
segment_stderr, segment = self._segments[0]
|
|
37
|
+
if len(segment) <= excess:
|
|
38
|
+
self._segments.popleft()
|
|
39
|
+
self._size -= len(segment)
|
|
40
|
+
excess -= len(segment)
|
|
41
|
+
else:
|
|
42
|
+
self._segments[0] = (segment_stderr, segment[excess:])
|
|
43
|
+
self._size -= excess
|
|
44
|
+
excess = 0
|
|
45
|
+
|
|
46
|
+
def outputs(self) -> tuple[bytes, bytes]:
|
|
47
|
+
stdout = bytearray()
|
|
48
|
+
stderr = bytearray()
|
|
49
|
+
for is_stderr, segment in self._segments:
|
|
50
|
+
(stderr if is_stderr else stdout).extend(segment)
|
|
51
|
+
return bytes(stdout), bytes(stderr)
|
|
16
52
|
|
|
17
53
|
|
|
18
54
|
@dataclass(frozen=True, slots=True)
|
|
@@ -21,10 +57,12 @@ class SshConnectionConfig:
|
|
|
21
57
|
port: int = 22
|
|
22
58
|
username: str | None = None
|
|
23
59
|
client_keys: Sequence[str] = ()
|
|
24
|
-
known_hosts: str | None =
|
|
60
|
+
known_hosts: str | None | tuple[()] = ()
|
|
61
|
+
tunnel: str | None = None
|
|
62
|
+
password: str | None = None
|
|
25
63
|
|
|
26
64
|
|
|
27
|
-
class
|
|
65
|
+
class _AsyncSshTransport:
|
|
28
66
|
capabilities = TransportCapabilities(
|
|
29
67
|
separate_stderr=True,
|
|
30
68
|
native_binary_transfer=True,
|
|
@@ -35,7 +73,7 @@ class SshTransport:
|
|
|
35
73
|
self._connection = connection
|
|
36
74
|
|
|
37
75
|
@classmethod
|
|
38
|
-
async def connect(cls, config: SshConnectionConfig) ->
|
|
76
|
+
async def connect(cls, config: SshConnectionConfig, *, connect_timeout: float | None = None) -> _AsyncSshTransport:
|
|
39
77
|
if not config.host.strip():
|
|
40
78
|
raise ValueError("SSH host must not be empty")
|
|
41
79
|
if not 1 <= config.port <= 65535:
|
|
@@ -46,6 +84,9 @@ class SshTransport:
|
|
|
46
84
|
username=config.username,
|
|
47
85
|
client_keys=list(config.client_keys),
|
|
48
86
|
known_hosts=config.known_hosts,
|
|
87
|
+
tunnel=config.tunnel,
|
|
88
|
+
password=config.password,
|
|
89
|
+
connect_timeout=connect_timeout,
|
|
49
90
|
encoding=None,
|
|
50
91
|
)
|
|
51
92
|
return cls(connection)
|
|
@@ -64,28 +105,76 @@ class SshTransport:
|
|
|
64
105
|
raise ValueError("timeout must be positive")
|
|
65
106
|
if output_limit <= 0:
|
|
66
107
|
raise ValueError("output_limit must be positive")
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
108
|
+
open_task: asyncio.Future[asyncssh.SSHClientProcess[bytes]] = asyncio.ensure_future(
|
|
109
|
+
self._connection.create_process(command, encoding=None)
|
|
110
|
+
)
|
|
111
|
+
try:
|
|
112
|
+
process = await asyncio.shield(open_task)
|
|
113
|
+
except asyncio.CancelledError:
|
|
114
|
+
self._connection.abort()
|
|
115
|
+
close_task = asyncio.create_task(self._connection.wait_closed())
|
|
116
|
+
while not close_task.done():
|
|
117
|
+
try:
|
|
118
|
+
await asyncio.shield(close_task)
|
|
119
|
+
except asyncio.CancelledError:
|
|
120
|
+
continue
|
|
121
|
+
if not open_task.done():
|
|
122
|
+
open_task.cancel()
|
|
123
|
+
await asyncio.gather(open_task, return_exceptions=True)
|
|
124
|
+
raise
|
|
125
|
+
output = _CombinedOutputBuffer(output_limit)
|
|
126
|
+
stdout_task = asyncio.create_task(self._read_stream(process.stdout, output, stderr=False))
|
|
127
|
+
stderr_task = asyncio.create_task(self._read_stream(process.stderr, output, stderr=True))
|
|
70
128
|
input_task = asyncio.create_task(self._write_stdin(process, stdin))
|
|
129
|
+
wait_task = asyncio.create_task(process.wait_closed())
|
|
130
|
+
tasks = (stdout_task, stderr_task, input_task, wait_task)
|
|
71
131
|
try:
|
|
72
|
-
await asyncio.wait_for(
|
|
73
|
-
await
|
|
74
|
-
stdout, stderr =
|
|
132
|
+
await asyncio.wait_for(asyncio.gather(asyncio.shield(wait_task), input_task), timeout=timeout)
|
|
133
|
+
await asyncio.gather(stdout_task, stderr_task)
|
|
134
|
+
stdout, stderr = output.outputs()
|
|
75
135
|
return ExecResult(process.exit_status, stdout, stderr)
|
|
76
136
|
except TimeoutError:
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
137
|
+
await self._abort_process(process, tasks)
|
|
138
|
+
stdout, stderr = output.outputs()
|
|
139
|
+
return ExecResult(None, stdout, stderr, timed_out=True)
|
|
140
|
+
except BaseException:
|
|
141
|
+
cleanup = asyncio.create_task(self._abort_process(process, tasks))
|
|
142
|
+
while not cleanup.done():
|
|
143
|
+
try:
|
|
144
|
+
await asyncio.shield(cleanup)
|
|
145
|
+
except asyncio.CancelledError:
|
|
146
|
+
continue
|
|
147
|
+
await cleanup
|
|
148
|
+
raise
|
|
149
|
+
|
|
150
|
+
@staticmethod
|
|
151
|
+
async def _abort_process(
|
|
152
|
+
process: asyncssh.SSHClientProcess[bytes],
|
|
153
|
+
tasks: tuple[asyncio.Task[None], asyncio.Task[None], asyncio.Task[None], asyncio.Task[None]],
|
|
154
|
+
) -> None:
|
|
155
|
+
process.terminate()
|
|
156
|
+
process.close()
|
|
157
|
+
wait_task = tasks[-1]
|
|
158
|
+
try:
|
|
159
|
+
await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
|
|
160
|
+
except BaseException:
|
|
161
|
+
process.kill()
|
|
162
|
+
process.close()
|
|
163
|
+
with suppress(BaseException):
|
|
164
|
+
await asyncio.wait_for(asyncio.shield(wait_task), timeout=1)
|
|
165
|
+
for task in tasks:
|
|
166
|
+
if not task.done():
|
|
81
167
|
task.cancel()
|
|
82
|
-
|
|
83
|
-
return ExecResult(None, b"", b"", timed_out=True)
|
|
168
|
+
await asyncio.gather(*tasks, return_exceptions=True)
|
|
84
169
|
|
|
85
170
|
@staticmethod
|
|
86
171
|
async def _write_stdin(process: asyncssh.SSHClientProcess[bytes], stdin: Iterable[bytes] | None) -> None:
|
|
87
172
|
if stdin is not None:
|
|
88
|
-
|
|
173
|
+
iterator = iter(stdin)
|
|
174
|
+
while True:
|
|
175
|
+
present, chunk = await asyncio.to_thread(_next_chunk, iterator)
|
|
176
|
+
if not present:
|
|
177
|
+
break
|
|
89
178
|
if not isinstance(chunk, bytes):
|
|
90
179
|
raise TypeError("stdin chunks must be bytes")
|
|
91
180
|
process.stdin.write(chunk)
|
|
@@ -93,16 +182,12 @@ class SshTransport:
|
|
|
93
182
|
process.stdin.write_eof()
|
|
94
183
|
|
|
95
184
|
@staticmethod
|
|
96
|
-
async def _read_stream(stream,
|
|
97
|
-
output = bytearray()
|
|
185
|
+
async def _read_stream(stream, output: _CombinedOutputBuffer, *, stderr: bool) -> None:
|
|
98
186
|
while True:
|
|
99
187
|
chunk = await stream.read(64 * 1024)
|
|
100
188
|
if not chunk:
|
|
101
189
|
break
|
|
102
|
-
output.
|
|
103
|
-
if len(output) > output_limit:
|
|
104
|
-
del output[: len(output) - output_limit]
|
|
105
|
-
return bytes(output)
|
|
190
|
+
output.append(bytes(chunk), stderr=stderr)
|
|
106
191
|
|
|
107
192
|
async def upload_file(
|
|
108
193
|
self,
|
|
@@ -142,6 +227,57 @@ class SshTransport:
|
|
|
142
227
|
sftp.exit()
|
|
143
228
|
return TransferResult(transferred, digest.hexdigest())
|
|
144
229
|
|
|
230
|
+
async def upload(
|
|
231
|
+
self,
|
|
232
|
+
chunks: Iterable[bytes],
|
|
233
|
+
remote_path: str,
|
|
234
|
+
*,
|
|
235
|
+
mode: int,
|
|
236
|
+
expected_size: int | None,
|
|
237
|
+
) -> TransferResult:
|
|
238
|
+
self._validate_remote_path(remote_path)
|
|
239
|
+
temporary = f"{remote_path}.hostbridge-{uuid.uuid4().hex}.tmp"
|
|
240
|
+
digest = hashlib.sha256()
|
|
241
|
+
transferred = 0
|
|
242
|
+
iterator = iter(chunks)
|
|
243
|
+
sftp = await self._connection.start_sftp_client()
|
|
244
|
+
try:
|
|
245
|
+
async with sftp.open(temporary, "wb", encoding=None) as remote_file:
|
|
246
|
+
while True:
|
|
247
|
+
present, chunk = await asyncio.to_thread(_next_chunk, iterator)
|
|
248
|
+
if not present:
|
|
249
|
+
break
|
|
250
|
+
if not isinstance(chunk, bytes):
|
|
251
|
+
raise TypeError("upload chunks must be bytes")
|
|
252
|
+
await remote_file.write(chunk)
|
|
253
|
+
digest.update(chunk)
|
|
254
|
+
transferred += len(chunk)
|
|
255
|
+
if expected_size is not None and transferred != expected_size:
|
|
256
|
+
raise TransportError(f"upload size mismatch: expected {expected_size}, received {transferred}")
|
|
257
|
+
await sftp.chmod(temporary, mode)
|
|
258
|
+
await sftp.rename(temporary, remote_path)
|
|
259
|
+
except BaseException:
|
|
260
|
+
with suppress(Exception):
|
|
261
|
+
await sftp.remove(temporary)
|
|
262
|
+
raise
|
|
263
|
+
finally:
|
|
264
|
+
sftp.exit()
|
|
265
|
+
return TransferResult(transferred, digest.hexdigest())
|
|
266
|
+
|
|
267
|
+
async def download(self, remote_path: str, *, chunk_size: int) -> AsyncGenerator[bytes, None]:
|
|
268
|
+
self._validate_chunk_size(chunk_size)
|
|
269
|
+
self._validate_remote_path(remote_path)
|
|
270
|
+
sftp = await self._connection.start_sftp_client()
|
|
271
|
+
try:
|
|
272
|
+
async with sftp.open(remote_path, "rb", encoding=None) as remote_file:
|
|
273
|
+
while True:
|
|
274
|
+
raw_chunk = await remote_file.read(chunk_size)
|
|
275
|
+
if not raw_chunk:
|
|
276
|
+
break
|
|
277
|
+
yield raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
|
|
278
|
+
finally:
|
|
279
|
+
sftp.exit()
|
|
280
|
+
|
|
145
281
|
async def download_file(
|
|
146
282
|
self,
|
|
147
283
|
remote_path: str,
|
|
@@ -190,6 +326,154 @@ class SshTransport:
|
|
|
190
326
|
if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
|
|
191
327
|
raise ValueError("remote_path must be a non-empty path without NUL bytes")
|
|
192
328
|
|
|
329
|
+
async def shell_write(self, data: bytes) -> None:
|
|
330
|
+
if not isinstance(data, bytes):
|
|
331
|
+
raise TypeError("shell data must be bytes")
|
|
332
|
+
process = await self._ensure_shell()
|
|
333
|
+
process.stdin.write(data)
|
|
334
|
+
await process.stdin.drain()
|
|
335
|
+
|
|
336
|
+
async def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
|
|
337
|
+
if timeout < 0:
|
|
338
|
+
raise ValueError("timeout must be non-negative")
|
|
339
|
+
if max_bytes <= 0:
|
|
340
|
+
raise ValueError("max_bytes must be positive")
|
|
341
|
+
process = await self._ensure_shell()
|
|
342
|
+
try:
|
|
343
|
+
data = await asyncio.wait_for(process.stdout.read(max_bytes), timeout=timeout)
|
|
344
|
+
except TimeoutError:
|
|
345
|
+
data = b""
|
|
346
|
+
return ShellReadResult(bytes(data), process.exit_status is None and not self._connection.is_closed())
|
|
347
|
+
|
|
348
|
+
async def _ensure_shell(self) -> asyncssh.SSHClientProcess[bytes]:
|
|
349
|
+
shell = getattr(self, "_shell", None)
|
|
350
|
+
if shell is None or shell.exit_status is not None:
|
|
351
|
+
shell = await self._connection.create_process(term_type="xterm", encoding=None)
|
|
352
|
+
self._shell = shell
|
|
353
|
+
return shell
|
|
354
|
+
|
|
193
355
|
async def close(self) -> None:
|
|
356
|
+
shell = getattr(self, "_shell", None)
|
|
357
|
+
if shell is not None:
|
|
358
|
+
shell.terminate()
|
|
359
|
+
shell.close()
|
|
360
|
+
with suppress(Exception):
|
|
361
|
+
await asyncio.wait_for(shell.wait_closed(), timeout=1)
|
|
194
362
|
self._connection.close()
|
|
195
363
|
await self._connection.wait_closed()
|
|
364
|
+
|
|
365
|
+
|
|
366
|
+
def _next_chunk(iterator: Iterator[bytes]) -> tuple[bool, bytes]:
|
|
367
|
+
try:
|
|
368
|
+
return True, next(iterator)
|
|
369
|
+
except StopIteration:
|
|
370
|
+
return False, b""
|
|
371
|
+
|
|
372
|
+
|
|
373
|
+
class _EventLoopThread:
|
|
374
|
+
def __init__(self) -> None:
|
|
375
|
+
self.loop = asyncio.new_event_loop()
|
|
376
|
+
self.thread = threading.Thread(target=self._run, name="hostbridge-ssh", daemon=True)
|
|
377
|
+
self.started = threading.Event()
|
|
378
|
+
self.thread.start()
|
|
379
|
+
self.started.wait()
|
|
380
|
+
|
|
381
|
+
def _run(self) -> None:
|
|
382
|
+
asyncio.set_event_loop(self.loop)
|
|
383
|
+
self.started.set()
|
|
384
|
+
self.loop.run_forever()
|
|
385
|
+
self.loop.run_until_complete(self.loop.shutdown_asyncgens())
|
|
386
|
+
self.loop.close()
|
|
387
|
+
|
|
388
|
+
def run(self, awaitable: Awaitable[T]) -> T:
|
|
389
|
+
if not self.thread.is_alive():
|
|
390
|
+
raise TransportError("SSH transport is closed")
|
|
391
|
+
return asyncio.run_coroutine_threadsafe(_await_value(awaitable), self.loop).result()
|
|
392
|
+
|
|
393
|
+
def stop(self) -> None:
|
|
394
|
+
if self.thread.is_alive():
|
|
395
|
+
self.loop.call_soon_threadsafe(self.loop.stop)
|
|
396
|
+
self.thread.join(timeout=2)
|
|
397
|
+
|
|
398
|
+
|
|
399
|
+
class SshTransport:
|
|
400
|
+
capabilities = _AsyncSshTransport.capabilities
|
|
401
|
+
|
|
402
|
+
def __init__(self, core: _AsyncSshTransport, runner: _EventLoopThread) -> None:
|
|
403
|
+
self._core = core
|
|
404
|
+
self._runner = runner
|
|
405
|
+
self._close_lock = threading.Lock()
|
|
406
|
+
self._closed = False
|
|
407
|
+
|
|
408
|
+
@classmethod
|
|
409
|
+
def connect(cls, config: SshConnectionConfig, *, connect_timeout: float | None = None) -> SshTransport:
|
|
410
|
+
runner = _EventLoopThread()
|
|
411
|
+
try:
|
|
412
|
+
core = runner.run(_AsyncSshTransport.connect(config, connect_timeout=connect_timeout))
|
|
413
|
+
except BaseException:
|
|
414
|
+
runner.stop()
|
|
415
|
+
raise
|
|
416
|
+
return cls(core, runner)
|
|
417
|
+
|
|
418
|
+
def exec(
|
|
419
|
+
self,
|
|
420
|
+
command: str,
|
|
421
|
+
*,
|
|
422
|
+
stdin: Iterable[bytes] | None,
|
|
423
|
+
timeout: float,
|
|
424
|
+
output_limit: int,
|
|
425
|
+
) -> ExecResult:
|
|
426
|
+
return self._run(self._core.exec(command, stdin=stdin, timeout=timeout, output_limit=output_limit))
|
|
427
|
+
|
|
428
|
+
def upload(
|
|
429
|
+
self,
|
|
430
|
+
chunks: Iterable[bytes],
|
|
431
|
+
remote_path: str,
|
|
432
|
+
*,
|
|
433
|
+
mode: int,
|
|
434
|
+
expected_size: int | None,
|
|
435
|
+
) -> TransferResult:
|
|
436
|
+
return self._run(self._core.upload(chunks, remote_path, mode=mode, expected_size=expected_size))
|
|
437
|
+
|
|
438
|
+
def download(self, remote_path: str, *, chunk_size: int) -> Iterator[bytes]:
|
|
439
|
+
stream = self._core.download(remote_path, chunk_size=chunk_size)
|
|
440
|
+
|
|
441
|
+
def chunks() -> Iterator[bytes]:
|
|
442
|
+
try:
|
|
443
|
+
while True:
|
|
444
|
+
try:
|
|
445
|
+
yield self._run(stream.__anext__())
|
|
446
|
+
except StopAsyncIteration:
|
|
447
|
+
return
|
|
448
|
+
finally:
|
|
449
|
+
self._run(stream.aclose())
|
|
450
|
+
|
|
451
|
+
return chunks()
|
|
452
|
+
|
|
453
|
+
def shell_write(self, data: bytes) -> None:
|
|
454
|
+
self._run(self._core.shell_write(data))
|
|
455
|
+
|
|
456
|
+
def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
|
|
457
|
+
return self._run(self._core.shell_read(timeout=timeout, max_bytes=max_bytes))
|
|
458
|
+
|
|
459
|
+
def close(self) -> None:
|
|
460
|
+
with self._close_lock:
|
|
461
|
+
if self._closed:
|
|
462
|
+
return
|
|
463
|
+
self._closed = True
|
|
464
|
+
try:
|
|
465
|
+
self._runner.run(self._core.close())
|
|
466
|
+
finally:
|
|
467
|
+
self._runner.stop()
|
|
468
|
+
|
|
469
|
+
def _run(self, awaitable: Awaitable[T]) -> T:
|
|
470
|
+
try:
|
|
471
|
+
return self._runner.run(awaitable)
|
|
472
|
+
except (StopAsyncIteration, TypeError, ValueError, TransportError):
|
|
473
|
+
raise
|
|
474
|
+
except Exception as exc:
|
|
475
|
+
raise TransportError(str(exc) or exc.__class__.__name__) from exc
|
|
476
|
+
|
|
477
|
+
|
|
478
|
+
async def _await_value(awaitable: Awaitable[T]) -> T:
|
|
479
|
+
return await awaitable
|