from __future__ import annotations import fcntl import os import pty import select import signal import sys import termios import tty from collections.abc import Callable, Sequence from .broker import InputBroker class PtyTap: """Full-duplex PTY proxy. Input is observed, never rewritten.""" def __init__(self, command: Sequence[str], broker: InputBroker, *, on_line: Callable[[str], None] | None = None): self.command = list(command) self.broker = broker self.on_line = on_line def _observe_input(self, buffer: bytearray, data: bytes) -> None: buffer.extend(data) while True: delimiters = [index for marker in (b"\r", b"\n") if (index := buffer.find(marker)) >= 0] if not delimiters: return index = min(delimiters) raw = bytes(buffer[:index]) del buffer[: index + 1] line = raw.decode("utf-8", errors="replace") routed = self.broker.inspect_line(line) if routed and self.on_line: self.on_line(line) def run(self) -> int: stdin_fd = sys.stdin.fileno() stdout_fd = sys.stdout.fileno() master, slave = pty.openpty() copy_winsize(stdin_fd, slave) pid = os.fork() if pid == 0: try: os.close(master) os.setsid() fcntl.ioctl(slave, termios.TIOCSCTTY, 0) for fd in (0, 1, 2): os.dup2(slave, fd) if slave > 2: os.close(slave) os.execvp(self.command[0], self.command) except BaseException as exc: os.write(2, f"tamq tap: cannot start {self.command[0]}: {exc}\n".encode()) os._exit(127) os.close(slave) input_buffer = bytearray() saved_terminal = termios.tcgetattr(stdin_fd) if os.isatty(stdin_fd) else None saved_handlers: dict[int, signal.Handlers] = {} input_open = True def resize(_signum=None, _frame=None) -> None: copy_winsize(stdin_fd, master) def forward(signum, _frame) -> None: try: os.killpg(pid, signum) except ProcessLookupError: pass status = 0 try: if saved_terminal is not None: tty.setraw(stdin_fd) saved_handlers[signal.SIGWINCH] = signal.signal(signal.SIGWINCH, resize) for signum in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP): saved_handlers[signum] = signal.signal(signum, forward) resize() while True: sources = [master] if input_open: sources.append(stdin_fd) try: readable, _, _ = select.select(sources, [], []) except InterruptedError: continue if master in readable: try: data = os.read(master, 65536) except OSError: break if not data: break write_all(stdout_fd, data) if input_open and stdin_fd in readable: data = os.read(stdin_fd, 65536) if not data: input_open = False try: os.killpg(pid, signal.SIGHUP) except ProcessLookupError: pass continue write_all(master, data) self._observe_input(input_buffer, data) except BaseException: try: os.killpg(pid, signal.SIGHUP) except ProcessLookupError: pass raise finally: for signum, handler in saved_handlers.items(): signal.signal(signum, handler) if saved_terminal is not None: termios.tcsetattr(stdin_fd, termios.TCSAFLUSH, saved_terminal) os.close(master) try: _, status = os.waitpid(pid, 0) except ChildProcessError: status = 0 return os.waitstatus_to_exitcode(status) def copy_winsize(source_fd: int, target_fd: int) -> bool: """Copy terminal rows/columns and pixel dimensions when a source tty exists.""" try: size = fcntl.ioctl(source_fd, termios.TIOCGWINSZ, b"\0" * 8) fcntl.ioctl(target_fd, termios.TIOCSWINSZ, size) return True except OSError: return False def write_all(fd: int, data: bytes) -> None: view = memoryview(data) while view: written = os.write(fd, view) if written == 0: raise OSError("terminal write returned zero bytes") view = view[written:]