"""Copy a report file from a running VM through its serial console or ssh.
A VM reached through ssh gets a separate ssh session that runs one command:
it prints the ``sha256sum`` of the file and then the file itself. The
student's terminal is not touched, so this works while ``pysnap connect`` is
attached. Every other VM uses its serial console.
The serial transfer uses only the guest shell and coreutils:
1. The VM must be running and expose ``UART1`` as a TCP server. VirtualBox
serves one client, so when ``pysnap connect`` is attached the transfer is
delegated to that process through its control channel
(:mod:`pysnap.report.control`), which runs the same protocol.
2. PySnap wakes the console with Enter and refuses to continue when the
prompt installed by ``report`` is visible: a recording is running and every
further command would end up in the student's report.
3. ``echo PYSNAP_$((20+22))`` must answer ``PYSNAP_42``. The echoed command
line cannot contain that text, so the answer proves a working shell and
not a login screen or a full-screen program.
4. One command prints a begin marker, the file as one ``base64`` line, its
``sha256sum`` and an end marker. The markers are split in the command so
that its echo never looks like a marker.
5. PySnap decodes the data, compares the checksum, writes the file
atomically and clears the guest screen.
Service commands start with a space, so shells with ``HISTCONTROL`` set to
``ignorespace`` or ``ignoreboth`` keep them out of the history.
"""
from __future__ import annotations
import base64
import binascii
from dataclasses import dataclass
import hashlib
import os
from pathlib import Path
import re
import secrets
import shlex
import socket
import tempfile
import threading
from time import monotonic
import subprocess
from typing import Callable, Protocol
from pysnap.errors import PySnapError
from pysnap.ssh import client as ssh_client
from pysnap.terminal.transport import open_serial_socket
DEFAULT_CONNECT_TIMEOUT = 10.0
DEFAULT_SHELL_TIMEOUT = 5.0
DEFAULT_IDLE_TIMEOUT = 10.0
DEFAULT_MAX_BYTES = 16 * 1024 * 1024
DEFAULT_SSH_TIMEOUT = 120.0
# Exit status of the ssh transfer command when the file cannot be read.
_SSH_MISSING_STATUS = 3
_QUIET_PERIOD = 0.3
PROBE_COMMAND = b" echo PYSNAP_$((20+22))\r"
PROBE_ANSWER = b"PYSNAP_42"
_ANSI_SEQUENCE = re.compile(rb"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b[()][0-9A-Za-z]|\x1b.")
# The prompt installed by ``report`` (``[\u@NN-HOST \W]# ``).
_REPORT_PROMPT = re.compile(rb"\[[^@\]\s]+@\d+-[^\s\]]+ [^\]]*\][#$] ?")
Connector = Callable[[str, int, float], socket.socket]
[docs]
def run_transfer(
console: "ConsoleIO",
vm_name: str,
remote_path: str,
*,
shell_timeout: float = DEFAULT_SHELL_TIMEOUT,
idle_timeout: float = DEFAULT_IDLE_TIMEOUT,
max_bytes: int = DEFAULT_MAX_BYTES,
) -> tuple[bytes, str]:
"""Run the transfer protocol on a serial console.
:param console: Console transport.
:param vm_name: VM name used in messages.
:param remote_path: Quoted shell path of the file.
:param shell_timeout: Seconds to wait for the shell check answer.
:param idle_timeout: Seconds without data after which the transfer fails.
:param max_bytes: Largest accepted file size.
:returns: File content and its SHA-256 checksum.
:raises ExtractError: When the console is not usable or the transfer fails.
"""
# Ctrl-U first: a line the student typed but did not finish is erased
# instead of being run by the Enter (it can be restored with Ctrl-Y).
greeting = console.exchange(b"\x15\r", quiet=_QUIET_PERIOD, limit=2.0)
if _REPORT_PROMPT.search(_plain(greeting)):
raise ExtractError(
f'A report recording is running on "{vm_name}". Finish it with exit or '
"Ctrl-D in pysnap connect, then extract the report."
)
console.send(PROBE_COMMAND)
if not console.read_until(lambda data: PROBE_ANSWER in data, shell_timeout):
raise ExtractError(
f'The console of "{vm_name}" is not at a shell prompt. Open it with '
"pysnap connect, log in or leave the running program and try again."
)
console.exchange(b"", quiet=_QUIET_PERIOD, limit=1.0)
data, checksum = _transfer(console, remote_path, idle_timeout, max_bytes)
console.send(b" clear\r")
console.exchange(b"", quiet=_QUIET_PERIOD, limit=1.0)
return data, checksum
[docs]
def ssh_transfer(
service,
vm_name: str,
remote_path: str,
*,
max_bytes: int = DEFAULT_MAX_BYTES,
timeout: float = DEFAULT_SSH_TIMEOUT,
) -> tuple[bytes, str]:
"""Copy a file through a separate ssh session with the PySnap key.
:param service: PySnap service; gives the ssh target and the key store.
:param vm_name: VM to copy from.
:param remote_path: Quoted shell path of the file.
:param max_bytes: Largest accepted file size.
:param timeout: Seconds the whole transfer may take.
:returns: File content and its SHA-256 checksum.
:raises ExtractError: When ssh fails or the file cannot be read.
"""
try:
target = service.ssh_target(vm_name)
except PySnapError as error:
raise ExtractError(str(error)) from None
# One more byte than allowed tells a file that is too large.
script = (
f"f={remote_path}; if [ -f \"$f\" ] && [ -r \"$f\" ]; then "
f"sha256sum < \"$f\" && head -c {max_bytes + 1} < \"$f\"; "
f"else exit {_SSH_MISSING_STATUS}; fi"
)
try:
result = ssh_client.run_remote(
target, service.ssh_key_store, f"sh -c {shlex.quote(script)}", timeout=timeout
)
except subprocess.TimeoutExpired:
raise ExtractError(f"The transfer did not finish within {timeout:.0f} seconds.") from None
if result.returncode == _SSH_MISSING_STATUS:
raise ExtractError(f"The file {remote_path} was not found or is not readable in the VM.")
if result.returncode != 0:
message = result.stderr.decode(errors="replace").strip()
if "Permission denied" in message:
raise ExtractError(
f'The ssh server of "{vm_name}" refused the PySnap key for {target.user}; '
f"run pysnap ssh-key {vm_name}."
)
raise ExtractError(
f'The ssh transfer from "{vm_name}" failed: {message or f"exit status {result.returncode}"}.'
)
header, separator, data = result.stdout.partition(b"\n")
if not separator or not header.split():
raise ExtractError("The VM sent an incomplete transfer.")
if len(data) > max_bytes:
raise ExtractError(f"The file is larger than {max_bytes} bytes.")
expected = header.split()[0].decode("ascii", errors="replace")
checksum = hashlib.sha256(data).hexdigest()
if checksum != expected:
raise ExtractError("The checksum of the received file does not match; try again.")
return data, checksum
[docs]
def remote_shell_path(name: str) -> str:
"""Return the quoted shell path of a file in the VM.
:param name: File name, looked up in the home directory, or a path.
:returns: Shell expression for the path.
:raises ExtractError: For empty names or names with control characters.
"""
if not name or any(ord(char) < 32 for char in name) or name.endswith("/"):
raise ExtractError(f"Invalid file name {name!r}.")
if "/" in name:
return shlex.quote(name)
return f'"$HOME"/{shlex.quote(name)}'
def _connectable_port(service, vm_name: str) -> int:
"""Check that the VM can serve a transfer and return its serial port."""
vm_info = service.show_vm(vm_name)
if vm_info.serial_port is None:
raise ExtractError(
f'Virtual machine "{vm_name}" has no TCP serial port; run pysnap plug first.'
)
if (vm_info.vm_state or "").lower() != "running":
raise ExtractError(
f'Virtual machine "{vm_name}" is not running; start it with pysnap connect.'
)
return vm_info.serial_port
def _transfer(console: "ConsoleIO", remote_path: str, idle_timeout: float, max_bytes: int) -> tuple[bytes, str]:
"""Run the transfer command and return the verified file content."""
token = secrets.token_hex(8)
begin, end, missing = (f"PYSNAP-{kind}-{token}".encode() for kind in ("BEGIN", "END", "MISSING"))
command = (
f" f={remote_path}; if [ -f \"$f\" ] && [ -r \"$f\" ]; then "
f"printf '%s-%s\\n' PYSNAP-BEGIN {token}; base64 -w0 \"$f\"; printf '\\n'; "
f"sha256sum < \"$f\"; printf '%s-%s\\n' PYSNAP-END {token}; "
f"else printf '%s-%s\\n' PYSNAP-MISSING {token}; fi\r"
)
console.send(command.encode())
limit = max_bytes * 4 // 3 + 4096
def complete(data: bytes) -> bool:
return end in data or missing in data or len(data) > limit
if not console.read_until(complete, idle_timeout, idle=True):
raise ExtractError("The transfer stopped: no data arrived from the VM in time.")
output = console.buffer.replace(b"\r", b"")
if missing in output:
raise ExtractError(f"The file {remote_path} was not found or is not readable in the VM.")
if end not in output:
raise ExtractError(f"The file is larger than {max_bytes} bytes.")
body = output[output.index(begin + b"\n") + len(begin) + 1:output.index(end)]
lines = body.split(b"\n")
if len(lines) < 2:
raise ExtractError("The VM sent an incomplete transfer.")
try:
data = base64.b64decode(lines[0], validate=True)
except (binascii.Error, ValueError):
raise ExtractError("The VM sent damaged data; try again.") from None
expected = lines[1].split()[0].decode("ascii", errors="replace") if lines[1].split() else ""
checksum = hashlib.sha256(data).hexdigest()
if checksum != expected:
raise ExtractError("The checksum of the received file does not match; try again.")
if len(data) > max_bytes:
raise ExtractError(f"The file is larger than {max_bytes} bytes.")
return data, checksum
def _plain(data: bytes) -> bytes:
"""Remove terminal control sequences from console output."""
return _ANSI_SEQUENCE.sub(b"", data)
def _write_atomically(target: Path, data: bytes) -> None:
"""Write a file through a temporary file in the same directory."""
target.parent.mkdir(parents=True, exist_ok=True)
handle, temporary = tempfile.mkstemp(prefix=f".{target.name}.", dir=target.parent)
try:
with os.fdopen(handle, "wb") as stream:
stream.write(data)
os.replace(temporary, target)
except BaseException:
Path(temporary).unlink(missing_ok=True)
raise
[docs]
class ConsoleIO(Protocol):
"""Console transport used by :func:`run_transfer`."""
buffer: bytes
[docs]
def send(self, data: bytes) -> None:
"""Send bytes and start a new response buffer."""
[docs]
def exchange(self, data: bytes, *, quiet: float, limit: float) -> bytes:
"""Send bytes and collect output until it stays quiet."""
[docs]
def read_until(self, done: Callable[[bytes], bool], timeout: float, *, idle: bool = False) -> bool:
"""Collect output until ``done`` holds."""
class _BufferedConsole:
"""Shared waiting logic; subclasses provide ``_write`` and ``_receive``."""
def __init__(self) -> None:
self.buffer = b""
def send(self, data: bytes) -> None:
"""Send bytes and start a new response buffer."""
self.buffer = b""
if data:
self._write(data)
def exchange(self, data: bytes, *, quiet: float, limit: float) -> bytes:
"""Send bytes and collect output until it stays quiet."""
self.send(data)
deadline = monotonic() + limit
last_data = monotonic()
while monotonic() < deadline and monotonic() - last_data < quiet:
if self._receive():
last_data = monotonic()
return self.buffer
def read_until(self, done: Callable[[bytes], bool], timeout: float, *, idle: bool = False) -> bool:
"""Collect output until ``done`` holds; ``idle`` restarts the timer on data."""
deadline = monotonic() + timeout
while not done(self.buffer):
if monotonic() >= deadline:
return False
if self._receive() and idle:
deadline = monotonic() + timeout
return True
def _write(self, data: bytes) -> None:
raise NotImplementedError
def _receive(self) -> bool:
raise NotImplementedError
[docs]
class SocketConsole(_BufferedConsole):
"""Console on a blocking serial TCP socket."""
def __init__(self, connection: socket.socket) -> None:
"""Wrap a connected socket.
:param connection: Serial TCP socket.
"""
super().__init__()
self.connection = connection
self.connection.settimeout(0.05)
def _write(self, data: bytes) -> None:
self.connection.sendall(data)
def _receive(self) -> bool:
"""Read one chunk; return whether data arrived."""
try:
chunk = self.connection.recv(65536)
except socket.timeout:
return False
except OSError as error:
raise ExtractError(f"The serial connection failed: {error}") from None
if not chunk:
raise ExtractError("The VM closed the serial connection.")
self.buffer += chunk
return True
[docs]
class QueueConsole(_BufferedConsole):
"""Console fed by another thread, such as the ``pysnap connect`` loop.
The owner calls :meth:`feed` with guest output and :meth:`close` when the
serial connection ends; ``write`` must be safe to call from the transfer
thread.
"""
def __init__(self, write: Callable[[bytes], None]) -> None:
"""Initialize the console.
:param write: Sends bytes to the guest.
"""
super().__init__()
self._writer = write
self._pending = b""
self._closed = False
self._condition = threading.Condition()
[docs]
def feed(self, data: bytes) -> None:
"""Hand guest output to the transfer.
:param data: Received bytes.
"""
with self._condition:
self._pending += data
self._condition.notify_all()
[docs]
def close(self) -> None:
"""Signal that the serial connection has ended."""
with self._condition:
self._closed = True
self._condition.notify_all()
def _write(self, data: bytes) -> None:
self._writer(data)
def _receive(self) -> bool:
"""Wait briefly for fed data; return whether data arrived."""
with self._condition:
if not self._pending and not self._closed:
self._condition.wait(0.05)
if self._pending:
self.buffer += self._pending
self._pending = b""
return True
if self._closed:
raise ExtractError("The serial connection of pysnap connect was closed.")
return False