tmux-amq/src/tamq/ptytap.py

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:]