@neoline/hostbridge 1.4.4 → 2.0.1

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.
@@ -0,0 +1,427 @@
1
+ """Ordered interactive PTY transport."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import base64
6
+ import binascii
7
+ import hashlib
8
+ import re
9
+ import shlex
10
+ import threading
11
+ import time
12
+ import uuid
13
+ from collections.abc import Callable, Iterable, Iterator
14
+ from contextlib import suppress
15
+ from typing import Protocol
16
+
17
+ import pexpect
18
+
19
+ from ..hosts import HostConfig
20
+ from .base import ExecResult, ShellReadResult, TransferResult, TransportBusy, TransportCapabilities, TransportError
21
+
22
+ _REMOTE_BASE64_DECODE = "{ base64 --decode 2>/dev/null || base64 -d 2>/dev/null || base64 -D; }"
23
+ PTY_BASE64_LINE_LENGTH = 3072
24
+
25
+
26
+ class ChildProcess(Protocol):
27
+ before: str
28
+ after: object
29
+
30
+ def sendline(self, line: str) -> None:
31
+ raise NotImplementedError
32
+
33
+ def send(self, text: str) -> None:
34
+ raise NotImplementedError
35
+
36
+ def expect(self, patterns: object, timeout: int | float = -1) -> int:
37
+ raise NotImplementedError
38
+
39
+ def isalive(self) -> bool:
40
+ raise NotImplementedError
41
+
42
+ def close(self, force: bool = False) -> None:
43
+ raise NotImplementedError
44
+
45
+ def read_nonblocking(self, size: int = 1, timeout: int | float = -1) -> str:
46
+ raise NotImplementedError
47
+
48
+
49
+ ChildFactory = Callable[[str, list[str], int | float], ChildProcess]
50
+
51
+
52
+ def _default_child_factory(command: str, args: list[str], timeout: int | float) -> ChildProcess:
53
+ return pexpect.spawn(
54
+ command,
55
+ args,
56
+ encoding="utf-8",
57
+ codec_errors="replace",
58
+ timeout=timeout,
59
+ echo=False,
60
+ )
61
+
62
+
63
+ def _iter_base64(chunks: Iterable[bytes]) -> Iterator[str]:
64
+ def split_line(encoded: str) -> Iterator[str]:
65
+ for offset in range(0, len(encoded), PTY_BASE64_LINE_LENGTH):
66
+ yield encoded[offset : offset + PTY_BASE64_LINE_LENGTH]
67
+
68
+ remainder = b""
69
+ for chunk in chunks:
70
+ if not isinstance(chunk, bytes):
71
+ raise TypeError("stdin chunks must be bytes")
72
+ data = remainder + chunk
73
+ complete = len(data) - (len(data) % 3)
74
+ if complete:
75
+ yield from split_line(base64.b64encode(data[:complete]).decode("ascii"))
76
+ remainder = data[complete:]
77
+ if remainder:
78
+ yield from split_line(base64.b64encode(remainder).decode("ascii"))
79
+
80
+
81
+ def _marker_text(value: object) -> str:
82
+ return value.group(0) if hasattr(value, "group") else str(value)
83
+
84
+
85
+ class PtyTransport:
86
+ capabilities = TransportCapabilities(
87
+ separate_stderr=False,
88
+ native_binary_transfer=False,
89
+ parallel_channels=1,
90
+ )
91
+
92
+ @classmethod
93
+ def connect(
94
+ cls,
95
+ host: HostConfig,
96
+ *,
97
+ child_factory: ChildFactory = _default_child_factory,
98
+ connect_timeout: int | float = 60,
99
+ remote_shell: str = "bash",
100
+ token_factory: Callable[[], str] | None = None,
101
+ ) -> PtyTransport:
102
+ if connect_timeout <= 0:
103
+ raise ValueError("connect_timeout must be positive")
104
+ tokens = token_factory or (lambda: uuid.uuid4().hex)
105
+ child = child_factory(host.command, list(host.args), connect_timeout)
106
+ try:
107
+ for step in host.login_steps:
108
+ child.expect(step.expect, timeout=connect_timeout)
109
+ if step.send_secret is not None:
110
+ secret = host.secrets.get(step.send_secret)
111
+ if secret is None:
112
+ raise TransportError(f"login step references missing secret {step.send_secret!r}")
113
+ child.send(f"{secret}\r")
114
+ else:
115
+ child.send(f"{step.send or ''}\r")
116
+ if host.login_steps:
117
+ child.expect(host.ready_expect or r"(?m)[^\r\n]*[#$]\s*$", timeout=connect_timeout)
118
+ token = tokens()
119
+ marker = f"__HB_READY_{token}__"
120
+ child.send(f"printf '\n{marker}\n'\r")
121
+ child.expect(re.escape(marker), timeout=connect_timeout)
122
+ except Exception as exc:
123
+ child.close(force=True)
124
+ if isinstance(exc, TransportError):
125
+ raise
126
+ raise TransportError(f"failed to establish PTY transport for {host.id}") from exc
127
+ return cls(child, remote_shell=remote_shell, token_factory=tokens)
128
+
129
+ def __init__(
130
+ self,
131
+ child: ChildProcess,
132
+ *,
133
+ remote_shell: str = "bash",
134
+ token_factory: Callable[[], str] | None = None,
135
+ marker_prefix: str = "HB",
136
+ ) -> None:
137
+ self.child = child
138
+ self.remote_shell = remote_shell
139
+ self._token_factory = token_factory or (lambda: uuid.uuid4().hex)
140
+ self._marker_prefix = marker_prefix
141
+ self.operation_lock = threading.Lock()
142
+ self._io_lock = threading.Lock()
143
+
144
+ def exec(
145
+ self,
146
+ command: str,
147
+ *,
148
+ stdin: Iterable[bytes] | None,
149
+ timeout: float,
150
+ output_limit: int,
151
+ ) -> ExecResult:
152
+ if not command.strip():
153
+ raise ValueError("command must not be empty")
154
+ if timeout <= 0:
155
+ raise ValueError("timeout must be positive")
156
+ if output_limit <= 0:
157
+ raise ValueError("output_limit must be positive")
158
+ if not self.operation_lock.acquire(blocking=False):
159
+ raise TransportBusy("PTY transport already has an active operation")
160
+ stdin_path: str | None = None
161
+ try:
162
+ if not self.child.isalive():
163
+ raise TransportError("PTY transport is not alive")
164
+ if stdin is not None:
165
+ stdin_path = self._stage_stdin_locked(stdin, timeout=timeout)
166
+ command = f"{command} < {shlex.quote(stdin_path)}"
167
+ return self._exec_locked(command, timeout=timeout, output_limit=output_limit)
168
+ finally:
169
+ if stdin_path is not None and self.child.isalive():
170
+ with suppress(Exception):
171
+ self._exec_locked(
172
+ f"trap - EXIT HUP INT TERM; rm -f -- {shlex.quote(stdin_path)} {shlex.quote(f'{stdin_path}.b64')}",
173
+ timeout=min(timeout, 10),
174
+ output_limit=1024,
175
+ )
176
+ self.operation_lock.release()
177
+
178
+ def _stage_stdin_locked(self, chunks: Iterable[bytes], *, timeout: float) -> str:
179
+ token = self._token_factory()
180
+ path = f"/tmp/hostbridge.stdin.{token}"
181
+ base64_path = f"{path}.b64"
182
+ heredoc = f"__{self._marker_prefix}_STDIN_{token}__"
183
+ done_pattern = re.compile(rf"__{self._marker_prefix}_STDIN_DONE_{re.escape(token)}__:(\d+)")
184
+ cleanup = f"rm -f -- {shlex.quote(path)} {shlex.quote(base64_path)}"
185
+ try:
186
+ self.child.send(f"trap {shlex.quote(cleanup)} EXIT HUP INT TERM; cat > {shlex.quote(path)} <<'{heredoc}'\n")
187
+ for encoded in _iter_base64(chunks):
188
+ self.child.send(encoded)
189
+ self.child.send(
190
+ f"\n{heredoc}\n"
191
+ f'tmp={shlex.quote(base64_path)}; mv {shlex.quote(path)} "$tmp"; '
192
+ f'{{ {_REMOTE_BASE64_DECODE} < "$tmp"; }} > {shlex.quote(path)}; status=$?; rm -f "$tmp"; '
193
+ f"printf '\n__{self._marker_prefix}_STDIN_DONE_{token}__:%s\n' \"$status\"\n"
194
+ )
195
+ self.child.expect(done_pattern, timeout=timeout)
196
+ except BaseException:
197
+ if self.child.isalive():
198
+ with suppress(Exception):
199
+ self.child.send("\x03")
200
+ with suppress(Exception):
201
+ self.child.sendline(f"trap - EXIT HUP INT TERM; {cleanup}")
202
+ raise
203
+ match = re.search(r":(\d+)", _marker_text(self.child.after))
204
+ if match is None or int(match.group(1)) != 0:
205
+ raise TransportError("failed to stage PTY stdin")
206
+ return path
207
+
208
+ def _exec_locked(self, command: str, *, timeout: float, output_limit: int) -> ExecResult:
209
+ token = self._token_factory()
210
+ start_marker = f"__{self._marker_prefix}_START_{token}__"
211
+ exit_pattern = re.compile(rf"__{self._marker_prefix}_EXIT_{re.escape(token)}__:(\d+)")
212
+ encoded_command = base64.b64encode(command.encode("utf-8")).decode("ascii")
213
+ remote_line = (
214
+ f"__scm_cmd=$(printf %s {shlex.quote(encoded_command)} | {_REMOTE_BASE64_DECODE}); "
215
+ f"printf '\n{start_marker}\n'; "
216
+ f"{self.remote_shell} -lc \"$__scm_cmd\"; "
217
+ "__scm_status=$?; "
218
+ f"printf '\n__{self._marker_prefix}_EXIT_{token}__:%s\n' \"$__scm_status\""
219
+ )
220
+ try:
221
+ self.child.sendline(remote_line)
222
+ self.child.expect(exit_pattern, timeout=timeout)
223
+ except pexpect.TIMEOUT:
224
+ return ExecResult(None, self._output_after_marker(str(self.child.before), start_marker, output_limit), timed_out=True)
225
+ match = re.search(r":(\d+)", _marker_text(self.child.after))
226
+ exit_code = int(match.group(1)) if match else None
227
+ return ExecResult(exit_code, self._output_after_marker(str(self.child.before), start_marker, output_limit))
228
+
229
+ @staticmethod
230
+ def _output_after_marker(raw: str, marker: str, output_limit: int) -> bytes:
231
+ normalized = raw.replace("\x1b[?2004h", "").replace("\x1b[?2004l", "")
232
+ normalized = normalized.replace("\r\n", "\n").replace("\r", "\n")
233
+ marker_index = normalized.rfind(marker)
234
+ output = normalized[marker_index + len(marker) :] if marker_index >= 0 else normalized
235
+ cleaned = [
236
+ line
237
+ for line in output.split("\n")
238
+ if not line.startswith("printf '") and not re.search(r"__[A-Z]+_(?:START|EXIT)_", line)
239
+ ]
240
+ encoded = "\n".join(cleaned).strip("\n").encode("utf-8", errors="replace")
241
+ return encoded[-output_limit:]
242
+
243
+ def upload(
244
+ self,
245
+ chunks: Iterable[bytes],
246
+ remote_path: str,
247
+ *,
248
+ mode: int,
249
+ expected_size: int | None,
250
+ ) -> TransferResult:
251
+ self._validate_remote_path(remote_path)
252
+ if not isinstance(mode, int) or isinstance(mode, bool) or not 0 <= mode <= 0o7777:
253
+ raise ValueError("mode must be an integer between 0 and 07777")
254
+ if expected_size is not None and expected_size < 0:
255
+ raise ValueError("expected_size must be non-negative")
256
+ if not self.operation_lock.acquire(blocking=False):
257
+ raise TransportBusy("PTY transport already has an active operation")
258
+ token = self._token_factory()
259
+ tmp_b64 = f"/tmp/hostbridge.upload.{token}.b64"
260
+ tmp_out = f"/tmp/hostbridge.upload.{token}.out"
261
+ heredoc = f"__{self._marker_prefix}_UPLOAD_{token}__"
262
+ done_pattern = re.compile(rf"__{self._marker_prefix}_UPLOAD_DONE_{re.escape(token)}__:(\d+)")
263
+ digest = hashlib.sha256()
264
+ transferred = 0
265
+
266
+ def counted_chunks() -> Iterator[bytes]:
267
+ nonlocal transferred
268
+ for chunk in chunks:
269
+ if not isinstance(chunk, bytes):
270
+ raise TypeError("upload chunks must be bytes")
271
+ transferred += len(chunk)
272
+ digest.update(chunk)
273
+ yield chunk
274
+
275
+ try:
276
+ if not self.child.isalive():
277
+ raise TransportError("PTY transport is not alive")
278
+ self.child.send(f"cat > {shlex.quote(tmp_b64)} <<'{heredoc}'\n")
279
+ for encoded in _iter_base64(counted_chunks()):
280
+ self.child.send(encoded)
281
+ self.child.send(
282
+ f"\n{heredoc}\nprintf '\n__{self._marker_prefix}_UPLOAD_DONE_{token}__:%s\n' \"$?\"\n"
283
+ )
284
+ self.child.expect(done_pattern, timeout=60)
285
+ match = re.search(r":(\d+)", _marker_text(self.child.after))
286
+ if match is None or int(match.group(1)) != 0:
287
+ raise TransportError("failed to stream PTY upload")
288
+ if expected_size is not None and transferred != expected_size:
289
+ raise TransportError(
290
+ f"upload size mismatch: expected {expected_size}, received {transferred}"
291
+ )
292
+ target_q = shlex.quote(remote_path)
293
+ finalize = (
294
+ "set -e; "
295
+ f"target={target_q}; tmp_b64={shlex.quote(tmp_b64)}; tmp_out={shlex.quote(tmp_out)}; "
296
+ 'mkdir -p "$(dirname "$target")"; '
297
+ f'{{ {_REMOTE_BASE64_DECODE} < "$tmp_b64"; }} > "$tmp_out"; '
298
+ f'chmod {format(mode, "04o")} "$tmp_out"; mv "$tmp_out" "$target"; rm -f "$tmp_b64"; '
299
+ "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
300
+ "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
301
+ 'bytes=$(wc -c < "$target" | tr -d " "); printf "BYTES=%s\\nSHA256=%s\\n" "$bytes" "$sha"'
302
+ )
303
+ result = self._exec_locked(finalize, timeout=60, output_limit=4096)
304
+ if result.timed_out or result.exit_code != 0:
305
+ raise TransportError("failed to finalize PTY upload")
306
+ fields = self._metadata_fields(result.stdout)
307
+ remote_size = int(fields.get("BYTES", transferred))
308
+ remote_sha = fields.get("SHA256", digest.hexdigest())
309
+ if remote_size != transferred or remote_sha != digest.hexdigest():
310
+ raise TransportError("PTY upload integrity check failed")
311
+ return TransferResult(transferred, remote_sha)
312
+ except Exception:
313
+ if self.child.isalive():
314
+ with suppress(Exception):
315
+ self._exec_locked(
316
+ f"rm -f -- {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}",
317
+ timeout=10,
318
+ output_limit=1024,
319
+ )
320
+ raise
321
+ finally:
322
+ self.operation_lock.release()
323
+
324
+ def download(self, remote_path: str, *, chunk_size: int) -> Iterator[bytes]:
325
+ self._validate_remote_path(remote_path)
326
+ if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
327
+ raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
328
+
329
+ def stream() -> Iterator[bytes]:
330
+ if not self.operation_lock.acquire(blocking=False):
331
+ raise TransportBusy("PTY transport already has an active operation")
332
+ digest = hashlib.sha256()
333
+ transferred = 0
334
+ try:
335
+ if not self.child.isalive():
336
+ raise TransportError("PTY transport is not alive")
337
+ remote_q = shlex.quote(remote_path)
338
+ metadata = (
339
+ f"target={remote_q}; "
340
+ 'if [ ! -e "$target" ] && [ ! -L "$target" ]; then printf "__HB_FILE_NOT_FOUND__\\n"; exit 44; fi; '
341
+ 'if [ ! -r "$target" ]; then printf "__HB_PERMISSION_DENIED__\\n"; exit 45; fi; '
342
+ "bytes=$(wc -c < \"$target\" | tr -d ' '); "
343
+ "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
344
+ "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
345
+ 'printf "BYTES=%s\\nSHA256=%s\\n" "$bytes" "$sha"'
346
+ )
347
+ result = self._exec_locked(metadata, timeout=60, output_limit=4096)
348
+ if result.timed_out or result.exit_code != 0:
349
+ if b"__HB_FILE_NOT_FOUND__" in result.stdout:
350
+ raise TransportError(f"remote file not found: {remote_path}")
351
+ if b"__HB_PERMISSION_DENIED__" in result.stdout:
352
+ raise TransportError(f"permission denied reading remote file: {remote_path}")
353
+ raise TransportError(f"failed to inspect remote file {remote_path}")
354
+ fields = self._metadata_fields(result.stdout)
355
+ byte_count = int(fields.get("BYTES", "0"))
356
+ expected_sha = fields.get("SHA256", "")
357
+ chunk_count = (byte_count + chunk_size - 1) // chunk_size
358
+ for index in range(chunk_count):
359
+ command = (
360
+ f"dd if={remote_q} bs={chunk_size} skip={index} count=1 2>/dev/null | "
361
+ "{ base64 --wrap=0 2>/dev/null || base64 | tr -d '\\n'; }"
362
+ )
363
+ result = self._exec_locked(command, timeout=60, output_limit=chunk_size * 2 + 4096)
364
+ if result.timed_out or result.exit_code != 0:
365
+ raise TransportError(f"failed to download PTY chunk {index}")
366
+ try:
367
+ chunk = base64.b64decode(b"".join(result.stdout.split()), validate=True)
368
+ except (binascii.Error, ValueError) as exc:
369
+ raise TransportError(f"remote chunk {index} contains invalid base64") from exc
370
+ transferred += len(chunk)
371
+ digest.update(chunk)
372
+ yield chunk
373
+ if transferred != byte_count or (expected_sha and digest.hexdigest() != expected_sha):
374
+ raise TransportError("PTY download integrity check failed")
375
+ finally:
376
+ self.operation_lock.release()
377
+
378
+ return stream()
379
+
380
+ @staticmethod
381
+ def _validate_remote_path(remote_path: str) -> None:
382
+ if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
383
+ raise ValueError("remote_path must be a non-empty path without NUL bytes")
384
+
385
+ @staticmethod
386
+ def _metadata_fields(output: bytes) -> dict[str, str]:
387
+ fields: dict[str, str] = {}
388
+ for line in output.decode("utf-8", errors="replace").splitlines():
389
+ key, separator, value = line.partition("=")
390
+ if separator:
391
+ fields[key] = value
392
+ return fields
393
+
394
+ def close(self) -> None:
395
+ self.child.close(force=True)
396
+
397
+ def shell_write(self, data: bytes) -> None:
398
+ if not isinstance(data, bytes):
399
+ raise TypeError("shell data must be bytes")
400
+ if not self.child.isalive():
401
+ raise TransportError("PTY transport is not alive")
402
+ with self._io_lock:
403
+ self.child.send(data.decode("utf-8", errors="replace"))
404
+
405
+ def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
406
+ if timeout < 0:
407
+ raise ValueError("timeout must be non-negative")
408
+ if max_bytes <= 0:
409
+ raise ValueError("max_bytes must be positive")
410
+ output = bytearray()
411
+ deadline = time.monotonic() + timeout
412
+ with self._io_lock:
413
+ while len(output) < max_bytes:
414
+ remaining = max(0.0, deadline - time.monotonic())
415
+ try:
416
+ chunk = self.child.read_nonblocking(
417
+ size=min(4096, max_bytes - len(output)),
418
+ timeout=remaining,
419
+ )
420
+ except (pexpect.TIMEOUT, pexpect.EOF):
421
+ break
422
+ if not chunk:
423
+ break
424
+ output.extend(chunk.encode("utf-8", errors="replace"))
425
+ if remaining <= 0:
426
+ break
427
+ return ShellReadResult(bytes(output), self.child.isalive())
@@ -0,0 +1,195 @@
1
+ """Native AsyncSSH transport for direct SSH hosts."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import hashlib
7
+ import uuid
8
+ from collections.abc import Iterable, Sequence
9
+ from contextlib import suppress
10
+ from dataclasses import dataclass
11
+ from pathlib import Path
12
+
13
+ import asyncssh
14
+
15
+ from .base import ExecResult, TransferResult, TransportCapabilities
16
+
17
+
18
+ @dataclass(frozen=True, slots=True)
19
+ class SshConnectionConfig:
20
+ host: str
21
+ port: int = 22
22
+ username: str | None = None
23
+ client_keys: Sequence[str] = ()
24
+ known_hosts: str | None = None
25
+
26
+
27
+ class SshTransport:
28
+ capabilities = TransportCapabilities(
29
+ separate_stderr=True,
30
+ native_binary_transfer=True,
31
+ parallel_channels=8,
32
+ )
33
+
34
+ def __init__(self, connection: asyncssh.SSHClientConnection) -> None:
35
+ self._connection = connection
36
+
37
+ @classmethod
38
+ async def connect(cls, config: SshConnectionConfig) -> SshTransport:
39
+ if not config.host.strip():
40
+ raise ValueError("SSH host must not be empty")
41
+ if not 1 <= config.port <= 65535:
42
+ raise ValueError("SSH port must be between 1 and 65535")
43
+ connection = await asyncssh.connect(
44
+ config.host,
45
+ port=config.port,
46
+ username=config.username,
47
+ client_keys=list(config.client_keys),
48
+ known_hosts=config.known_hosts,
49
+ encoding=None,
50
+ )
51
+ return cls(connection)
52
+
53
+ async def exec(
54
+ self,
55
+ command: str,
56
+ *,
57
+ stdin: Iterable[bytes] | None,
58
+ timeout: float,
59
+ output_limit: int,
60
+ ) -> ExecResult:
61
+ if not command.strip():
62
+ raise ValueError("command must not be empty")
63
+ if timeout <= 0:
64
+ raise ValueError("timeout must be positive")
65
+ if output_limit <= 0:
66
+ raise ValueError("output_limit must be positive")
67
+ process = await self._connection.create_process(command, encoding=None)
68
+ stdout_task = asyncio.create_task(self._read_stream(process.stdout, output_limit))
69
+ stderr_task = asyncio.create_task(self._read_stream(process.stderr, output_limit))
70
+ input_task = asyncio.create_task(self._write_stdin(process, stdin))
71
+ try:
72
+ await asyncio.wait_for(process.wait_closed(), timeout=timeout)
73
+ await input_task
74
+ stdout, stderr = await asyncio.gather(stdout_task, stderr_task)
75
+ return ExecResult(process.exit_status, stdout, stderr)
76
+ except TimeoutError:
77
+ process.terminate()
78
+ with suppress(Exception):
79
+ await process.wait_closed()
80
+ for task in (stdout_task, stderr_task, input_task):
81
+ task.cancel()
82
+ await asyncio.gather(stdout_task, stderr_task, input_task, return_exceptions=True)
83
+ return ExecResult(None, b"", b"", timed_out=True)
84
+
85
+ @staticmethod
86
+ async def _write_stdin(process: asyncssh.SSHClientProcess[bytes], stdin: Iterable[bytes] | None) -> None:
87
+ if stdin is not None:
88
+ for chunk in stdin:
89
+ if not isinstance(chunk, bytes):
90
+ raise TypeError("stdin chunks must be bytes")
91
+ process.stdin.write(chunk)
92
+ await process.stdin.drain()
93
+ process.stdin.write_eof()
94
+
95
+ @staticmethod
96
+ async def _read_stream(stream, output_limit: int) -> bytes:
97
+ output = bytearray()
98
+ while True:
99
+ chunk = await stream.read(64 * 1024)
100
+ if not chunk:
101
+ break
102
+ output.extend(chunk)
103
+ if len(output) > output_limit:
104
+ del output[: len(output) - output_limit]
105
+ return bytes(output)
106
+
107
+ async def upload_file(
108
+ self,
109
+ local_path: Path | str,
110
+ remote_path: str,
111
+ *,
112
+ chunk_size: int = 256 * 1024,
113
+ mode: int = 0o644,
114
+ ) -> TransferResult:
115
+ source = Path(local_path)
116
+ self._validate_chunk_size(chunk_size)
117
+ self._validate_remote_path(remote_path)
118
+ temporary = f"{remote_path}.hostbridge-{uuid.uuid4().hex}.tmp"
119
+ digest = hashlib.sha256()
120
+ transferred = 0
121
+ sftp = await self._connection.start_sftp_client()
122
+ try:
123
+ local_file = source.open("rb")
124
+ try:
125
+ async with sftp.open(temporary, "wb") as remote_file:
126
+ while True:
127
+ chunk = await asyncio.to_thread(local_file.read, chunk_size)
128
+ if not chunk:
129
+ break
130
+ await remote_file.write(chunk)
131
+ digest.update(chunk)
132
+ transferred += len(chunk)
133
+ finally:
134
+ local_file.close()
135
+ await sftp.chmod(temporary, mode)
136
+ await sftp.rename(temporary, remote_path)
137
+ except Exception:
138
+ with suppress(Exception):
139
+ await sftp.remove(temporary)
140
+ raise
141
+ finally:
142
+ sftp.exit()
143
+ return TransferResult(transferred, digest.hexdigest())
144
+
145
+ async def download_file(
146
+ self,
147
+ remote_path: str,
148
+ local_path: Path | str,
149
+ *,
150
+ chunk_size: int = 256 * 1024,
151
+ ) -> TransferResult:
152
+ destination = Path(local_path)
153
+ self._validate_chunk_size(chunk_size)
154
+ self._validate_remote_path(remote_path)
155
+ temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
156
+ digest = hashlib.sha256()
157
+ transferred = 0
158
+ sftp = await self._connection.start_sftp_client()
159
+ try:
160
+ destination.parent.mkdir(parents=True, exist_ok=True)
161
+ local_file = temporary.open("wb")
162
+ try:
163
+ async with sftp.open(remote_path, "rb") as remote_file:
164
+ while True:
165
+ raw_chunk = await remote_file.read(chunk_size)
166
+ if not raw_chunk:
167
+ break
168
+ chunk = raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
169
+ await asyncio.to_thread(local_file.write, chunk)
170
+ digest.update(chunk)
171
+ transferred += len(chunk)
172
+ finally:
173
+ local_file.close()
174
+ temporary.chmod(0o600)
175
+ temporary.replace(destination)
176
+ except Exception:
177
+ temporary.unlink(missing_ok=True)
178
+ raise
179
+ finally:
180
+ sftp.exit()
181
+ return TransferResult(transferred, digest.hexdigest())
182
+
183
+ @staticmethod
184
+ def _validate_chunk_size(chunk_size: int) -> None:
185
+ if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
186
+ raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
187
+
188
+ @staticmethod
189
+ def _validate_remote_path(remote_path: str) -> None:
190
+ if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
191
+ raise ValueError("remote_path must be a non-empty path without NUL bytes")
192
+
193
+ async def close(self) -> None:
194
+ self._connection.close()
195
+ await self._connection.wait_closed()