embedded-lab-control 0.2.0__py3-none-any.whl

labctl/serial/capture.py

@@ -0,0 +1,430 @@
+"""
+Serial port capture and send utilities.
+
+Provides functions to capture serial output and send data via ser2net TCP
+connections. Used by CLI commands and MCP tools.
+"""
+
+import logging
+import re
+import select
+import socket
+import time
+from dataclasses import dataclass
+from typing import Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CaptureResult:
+    """Result of a serial capture operation."""
+
+    output: str = ""
+    lines: int = 0
+    pattern_matched: bool = False
+    elapsed_seconds: float = 0.0
+
+    def to_mcp_string(self, pattern: Optional[str] = None) -> str:
+        """Format as MCP tool return string."""
+        status = f"pattern '{pattern}' matched" if self.pattern_matched else "timeout"
+        header = (
+            f"[Captured {self.lines} lines in {self.elapsed_seconds:.1f}s, "
+            f"{status}]"
+        )
+        return f"{header}\n{self.output}"
+
+
+@dataclass
+class SendResult:
+    """Result of a serial send operation."""
+
+    sent: bool = False
+    bytes_sent: int = 0
+    capture: Optional[CaptureResult] = None
+
+    def to_mcp_string(self) -> str:
+        """Format as MCP tool return string."""
+        if self.capture:
+            return self.capture.to_mcp_string()
+        return f"Sent {self.bytes_sent} bytes"
+
+
+@dataclass
+class SendPacing:
+    """Outbound write pacing to avoid receiver UART RX-FIFO overrun.
+
+    A gapless burst gives the receiver's RX ISR no window to drain its
+    FIFO, so a sustained write can lose a contiguous mid-payload byte
+    window. Writing ``chunk_size`` bytes per slice with ``delay_ms``
+    between slices inserts the gaps the receiver needs. ``chunk_size=1``
+    expresses per-byte pacing.
+
+    The wire still transmits at the configured baud rate; pacing only
+    controls how the payload is fed to ser2net, which is what creates
+    the inter-slice gaps on the UART.
+    """
+
+    chunk_size: int = 0
+    delay_ms: float = 0.0
+
+    @property
+    def enabled(self) -> bool:
+        return self.chunk_size > 0 and self.delay_ms > 0
+
+
+# Pause between the parts of a multi-part send (one per named key). Meant
+# to let consoles that time out a lone ESC quickly (e.g. U-Boot's menus)
+# see it as a key; not checked on hardware, and firmware that waits longer
+# for the rest of a sequence (EDK2 waits up to ~2 s) is not helped by it.
+PART_GAP_SECONDS = 0.05
+
+
+def _write_payload(
+    sock: socket.socket,
+    raw: Union[bytes, list[bytes]],
+    pacing: Optional[SendPacing] = None,
+) -> None:
+    """Write ``raw`` to ``sock``, optionally paced into delayed slices.
+
+    A list is written part by part with PART_GAP_SECONDS between parts.
+    """
+    if isinstance(raw, list):
+        for i, part in enumerate(raw):
+            if i:
+                time.sleep(PART_GAP_SECONDS)
+            _write_payload(sock, part, pacing)
+        return
+    if pacing is None or not pacing.enabled:
+        sock.sendall(raw)
+        return
+    delay = pacing.delay_ms / 1000.0
+    step = pacing.chunk_size
+    for i in range(0, len(raw), step):
+        sock.sendall(raw[i : i + step])
+        if i + step < len(raw):
+            time.sleep(delay)
+
+
+def _payload_len(raw: Union[bytes, list[bytes]]) -> int:
+    return sum(map(len, raw)) if isinstance(raw, list) else len(raw)
+
+
+def resolve_port(manager, port_name: str):
+    """Resolve a port name to a SerialPort object.
+
+    Tries alias first, then falls back to SBC name + console port.
+
+    Args:
+        manager: ResourceManager instance
+        port_name: Port alias or SBC name
+
+    Returns:
+        SerialPort object
+
+    Raises:
+        ValueError: If port cannot be resolved
+    """
+    # Try alias first
+    port = manager.get_serial_port_by_alias(port_name)
+    if port:
+        return port
+
+    # Fall back to SBC name -> console port
+    sbc = manager.get_sbc_by_name(port_name)
+    if sbc:
+        if sbc.console_port:
+            return sbc.console_port
+        raise ValueError(f"SBC '{port_name}' has no console port assigned")
+
+    raise ValueError(f"'{port_name}' is not a known port alias or SBC name")
+
+
+def capture_serial_output(
+    tcp_host: str,
+    tcp_port: int,
+    timeout: float = 15.0,
+    until_pattern: Optional[str] = None,
+    tail: Optional[int] = None,
+) -> CaptureResult:
+    """Capture serial output from a ser2net TCP port.
+
+    Connects to the ser2net TCP port, reads output until timeout or
+    pattern match, then disconnects.
+
+    Args:
+        tcp_host: Host where ser2net is running (usually "localhost")
+        tcp_port: TCP port for the serial connection
+        timeout: Maximum seconds to capture
+        until_pattern: Regex pattern to stop on (matched per line)
+        tail: Return only last N lines (None = all)
+
+    Returns:
+        CaptureResult with captured output and metadata.
+    """
+    compiled_pattern = None
+    if until_pattern:
+        compiled_pattern = re.compile(until_pattern)
+
+    captured_lines: list[str] = []
+    pattern_matched = False
+    start_time = time.monotonic()
+    buffer = b""
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(5.0)  # Connection timeout
+
+    try:
+        sock.connect((tcp_host, tcp_port))
+        sock.setblocking(False)
+
+        while True:
+            elapsed = time.monotonic() - start_time
+            remaining = timeout - elapsed
+            if remaining <= 0:
+                break
+
+            wait = min(remaining, 0.5)
+            readable, _, _ = select.select([sock], [], [], wait)
+
+            if readable:
+                try:
+                    data = sock.recv(4096)
+                    if not data:
+                        break  # Connection closed
+
+                    buffer += data
+
+                    while b"\n" in buffer:
+                        line, buffer = buffer.split(b"\n", 1)
+                        line_str = line.decode("utf-8", errors="replace").rstrip()
+                        captured_lines.append(line_str)
+
+                        if compiled_pattern and compiled_pattern.search(line_str):
+                            pattern_matched = True
+                            break
+
+                    # Check partial buffer for pattern (e.g. prompts without newline)
+                    if not pattern_matched and compiled_pattern and buffer:
+                        partial = buffer.decode("utf-8", errors="replace")
+                        if compiled_pattern.search(partial):
+                            captured_lines.append(partial.rstrip())
+                            buffer = b""
+                            pattern_matched = True
+
+                    if pattern_matched:
+                        break
+
+                except BlockingIOError:
+                    pass
+
+        # Handle any remaining data in buffer (no trailing newline)
+        if buffer:
+            line_str = buffer.decode("utf-8", errors="replace").rstrip()
+            if line_str:
+                captured_lines.append(line_str)
+                if compiled_pattern and compiled_pattern.search(line_str):
+                    pattern_matched = True
+
+    except (ConnectionRefusedError, OSError) as e:
+        raise RuntimeError(
+            f"Cannot connect to serial port at {tcp_host}:{tcp_port}: {e}"
+        )
+    finally:
+        try:
+            sock.close()
+        except Exception:
+            pass
+
+    elapsed = time.monotonic() - start_time
+
+    if tail and len(captured_lines) > tail:
+        output_lines = captured_lines[-tail:]
+    else:
+        output_lines = captured_lines
+
+    return CaptureResult(
+        output="\n".join(output_lines),
+        lines=len(output_lines),
+        pattern_matched=pattern_matched,
+        elapsed_seconds=round(elapsed, 1),
+    )
+
+
+def send_serial_data(
+    tcp_host: str,
+    tcp_port: int,
+    data: Union[str, list[bytes]],
+    newline: bool = True,
+    capture_timeout: Optional[float] = None,
+    capture_until: Optional[str] = None,
+    pacing: Optional[SendPacing] = None,
+) -> SendResult:
+    """Send data to a serial port via ser2net TCP.
+
+    Optionally captures the response after sending.
+
+    Args:
+        tcp_host: Host where ser2net is running
+        tcp_port: TCP port for the serial connection
+        data: Text to send (UTF-8), or a list of byte parts written with
+            a short pause between them
+        newline: Append \\r\\n after data (default True)
+        capture_timeout: If set, capture response for this many seconds
+        capture_until: If set, capture until this regex matches
+        pacing: Optional write pacing to avoid receiver RX-FIFO overrun
+
+    Returns:
+        SendResult with send status and optional capture.
+    """
+    if isinstance(data, list):
+        raw = list(data)
+        if newline:
+            raw.append(b"\r\n")
+    else:
+        raw = data.encode("utf-8")
+        if newline:
+            raw += b"\r\n"
+
+    if capture_timeout is not None:
+        # Send and capture in one connection
+        return _send_and_capture(
+            tcp_host,
+            tcp_port,
+            raw,
+            capture_timeout,
+            capture_until,
+            pacing,
+        )
+
+    # Send only
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(5.0)
+
+    try:
+        sock.connect((tcp_host, tcp_port))
+        _write_payload(sock, raw, pacing)
+        return SendResult(sent=True, bytes_sent=_payload_len(raw))
+    except (ConnectionRefusedError, OSError) as e:
+        raise RuntimeError(
+            f"Cannot connect to serial port at {tcp_host}:{tcp_port}: {e}"
+        )
+    finally:
+        try:
+            sock.close()
+        except Exception:
+            pass
+
+
+def _send_and_capture(
+    tcp_host: str,
+    tcp_port: int,
+    raw_data: Union[bytes, list[bytes]],
+    timeout: float,
+    until_pattern: Optional[str] = None,
+    pacing: Optional[SendPacing] = None,
+) -> SendResult:
+    """Send data then capture response on the same connection."""
+    compiled_pattern = None
+    if until_pattern:
+        compiled_pattern = re.compile(until_pattern)
+
+    captured_lines: list[str] = []
+    pattern_matched = False
+    buffer = b""
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(5.0)
+
+    try:
+        sock.connect((tcp_host, tcp_port))
+
+        # Drain stale data (100ms)
+        sock.setblocking(False)
+        drain_end = time.monotonic() + 0.1
+        while time.monotonic() < drain_end:
+            readable, _, _ = select.select([sock], [], [], 0.05)
+            if readable:
+                try:
+                    sock.recv(4096)
+                except BlockingIOError:
+                    pass
+
+        # Send data
+        sock.setblocking(True)
+        sock.settimeout(5.0)
+        _write_payload(sock, raw_data, pacing)
+        bytes_sent = _payload_len(raw_data)
+
+        # Capture response
+        sock.setblocking(False)
+        start_time = time.monotonic()
+
+        while True:
+            elapsed = time.monotonic() - start_time
+            remaining = timeout - elapsed
+            if remaining <= 0:
+                break
+
+            wait = min(remaining, 0.5)
+            readable, _, _ = select.select([sock], [], [], wait)
+
+            if readable:
+                try:
+                    data = sock.recv(4096)
+                    if not data:
+                        break
+
+                    buffer += data
+
+                    while b"\n" in buffer:
+                        line, buffer = buffer.split(b"\n", 1)
+                        line_str = line.decode("utf-8", errors="replace").rstrip()
+                        captured_lines.append(line_str)
+
+                        if compiled_pattern and compiled_pattern.search(line_str):
+                            pattern_matched = True
+                            break
+
+                    # Check partial buffer for pattern (e.g. prompts without newline)
+                    if not pattern_matched and compiled_pattern and buffer:
+                        partial = buffer.decode("utf-8", errors="replace")
+                        if compiled_pattern.search(partial):
+                            captured_lines.append(partial.rstrip())
+                            buffer = b""
+                            pattern_matched = True
+
+                    if pattern_matched:
+                        break
+
+                except BlockingIOError:
+                    pass
+
+        # Remaining buffer
+        if buffer:
+            line_str = buffer.decode("utf-8", errors="replace").rstrip()
+            if line_str:
+                captured_lines.append(line_str)
+                if compiled_pattern and compiled_pattern.search(line_str):
+                    pattern_matched = True
+
+        elapsed = time.monotonic() - start_time
+
+        capture = CaptureResult(
+            output="\n".join(captured_lines),
+            lines=len(captured_lines),
+            pattern_matched=pattern_matched,
+            elapsed_seconds=round(elapsed, 1),
+        )
+
+        return SendResult(sent=True, bytes_sent=bytes_sent, capture=capture)
+
+    except (ConnectionRefusedError, OSError) as e:
+        raise RuntimeError(
+            f"Cannot connect to serial port at {tcp_host}:{tcp_port}: {e}"
+        )
+    finally:
+        try:
+            sock.close()
+        except Exception:
+            pass

labctl/serial/keys.py

@@ -0,0 +1,89 @@
+"""Named keys for serial consoles, as the bytes a terminal sends.
+
+Tool arguments are text written by a model, which can't reliably carry
+raw control characters (an ESC is dropped before it reaches the server),
+so MCP clients name keys instead (issue #33). Encodings are xterm's,
+except Backspace (0x08, which U-Boot and EDK2 take as backspace; Linux
+ttys usually erase with 0x7F). None has been checked on real firmware.
+Function keys vary by terminal type (EDK2 VT100+ uses ESC 1..ESC 0,
+ESC ! and ESC @ for F1-F12); callers can send those as raw hex.
+"""
+
+from __future__ import annotations
+
+_NAMED = {
+    "esc": b"\x1b",
+    "enter": b"\r",
+    "tab": b"\t",
+    "space": b" ",
+    "backspace": b"\x08",
+    "up": b"\x1b[A",
+    "down": b"\x1b[B",
+    "right": b"\x1b[C",
+    "left": b"\x1b[D",
+    "home": b"\x1b[H",
+    "end": b"\x1b[F",
+    "insert": b"\x1b[2~",
+    "delete": b"\x1b[3~",
+    "pgup": b"\x1b[5~",
+    "pgdn": b"\x1b[6~",
+    "f1": b"\x1bOP",
+    "f2": b"\x1bOQ",
+    "f3": b"\x1bOR",
+    "f4": b"\x1bOS",
+    "f5": b"\x1b[15~",
+    "f6": b"\x1b[17~",
+    "f7": b"\x1b[18~",
+    "f8": b"\x1b[19~",
+    "f9": b"\x1b[20~",
+    "f10": b"\x1b[21~",
+    "f11": b"\x1b[23~",
+    "f12": b"\x1b[24~",
+}
+_ALIASES = {
+    "escape": "esc",
+    "return": "enter",
+    "pageup": "pgup",
+    "pagedown": "pgdn",
+    "page up": "pgup",
+    "page down": "pgdn",
+    "del": "delete",
+    "ins": "insert",
+}
+
+KEY_NAMES = (
+    "ESC, Enter, Tab, Space, Backspace (0x08), Up, Down, Left, Right, "
+    "Home, End, Insert, Delete, PgUp, PgDn, F1-F12, Ctrl-A..Ctrl-Z "
+    "(0x01-0x1A; also written Ctrl+C or ^C)"
+)
+_CTRL_PREFIXES = ("ctrl-", "ctrl+", "control-", "control+", "^")
+
+
+def encode_key(name: str) -> bytes:
+    """The bytes for key ``name`` (case-insensitive); ValueError if unknown."""
+    key = name.strip().lower() if isinstance(name, str) else ""
+    key = _ALIASES.get(key, key)
+    if key in _NAMED:
+        return _NAMED[key]
+    for prefix in _CTRL_PREFIXES:
+        if key.startswith(prefix):
+            letter = key[len(prefix) :]
+            if len(letter) == 1 and "a" <= letter <= "z":
+                return bytes([ord(letter) - ord("a") + 1])
+    raise ValueError(f"unknown key {name!r}; known keys: {KEY_NAMES}")
+
+
+def encode_keys(names: list[str]) -> list[bytes]:
+    """The bytes for each of a sequence of key names, in order."""
+    return [encode_key(n) for n in names]
+
+
+def decode_hex(text: str) -> bytes:
+    """Bytes from hex text, e.g. "1b5b32337e" or "1b 5b 32"; ValueError if
+    malformed."""
+    try:
+        return bytes.fromhex(text)
+    except (TypeError, ValueError):
+        raise ValueError(
+            f"data_hex must be pairs of hex digits (spaces allowed): {text!r}"
+        ) from None