#!/usr/bin/env python3
"""Run a command on a pty and strip terminal control *strings* from its output.

Why this exists
---------------
Modern Linux emits OSC strings (ESC ] ...) into the terminal stream:
  * systemd pam_systemd writes "OSC 3008" session context at login
  * shells/prompts emit "OSC 133" (semantic prompt) and "OSC 7" (cwd)

The DEC VT1200's terminal emulator (1990 firmware) does not understand
control strings and locks up on them.  This filter sits between the process
logging in and latd's pty, dropping OSC/DCS/SOS/PM/APC strings so the
terminal only ever sees sequences it can handle.

Usage:  osc-filter.py [command [args...]]      (default: /bin/login)

It is intended to be used as a latd service command, e.g.:
  latcp -A -a VT1200 -C /path/to/osc-filter.py
"""
import errno
import fcntl
import os
import pty
import select
import signal
import struct
import sys
import syslog
import termios
import tty

# String introducers (byte after ESC) that begin a control string terminated
# by BEL (0x07) or ST (ESC \):
#   ] OSC   P DCS   X SOS   ^ PM   _ APC
STRING_INTROS = b"]PX^_"
MAX_STRING = 8192

(ST_DATA, ST_ESC, ST_STRING, ST_STRING_ESC) = range(4)


class Filter:
    """Stateful control-string stripper (handles sequences split across reads)."""

    def __init__(self):
        self.state = ST_DATA
        self.n = 0

    def feed(self, data):
        out = bytearray()
        for b in data:
            if self.state == ST_DATA:
                if b == 0x1B:
                    self.state = ST_ESC
                else:
                    out.append(b)
            elif self.state == ST_ESC:
                if b in STRING_INTROS:
                    self.state = ST_STRING
                    self.n = 0
                elif b == 0x1B:
                    out.append(0x1B)          # keep, still expect next byte
                else:
                    out.append(0x1B)
                    out.append(b)
                    self.state = ST_DATA
            elif self.state == ST_STRING:
                self.n += 1
                if b == 0x07:                 # BEL ends the string
                    self.state = ST_DATA
                elif b == 0x1B:
                    self.state = ST_STRING_ESC
                elif self.n > MAX_STRING:     # runaway: give up, resync
                    self.state = ST_DATA
            else:                             # ST_STRING_ESC: saw ESC in string
                if b == 0x5C:                 # '\' -> ST ends the string
                    self.state = ST_DATA
                elif b == 0x1B:
                    pass                      # another ESC, keep waiting
                else:
                    self.state = ST_STRING
        return bytes(out)


def set_winsize(dst_fd, src_fd):
    try:
        packed = fcntl.ioctl(src_fd, termios.TIOCGWINSZ, b"\0" * 8)
        rows, cols = struct.unpack("HHHH", packed)[:2]
        if rows and cols:
            fcntl.ioctl(dst_fd, termios.TIOCSWINSZ,
                        struct.pack("HHHH", rows, cols, 0, 0))
    except OSError:
        pass


def main():
    command = sys.argv[1:] or ["/bin/login"]

    syslog.openlog("osc-filter")
    syslog.syslog(syslog.LOG_INFO, "starting: %s" % " ".join(command))

    # Make the outer terminal (latd's pty) a transparent byte pipe: no local
    # echo, no signal generation (so Ctrl-C reaches login), no line editing.
    # The inner pty owned by login/child does all of that.
    outer_termios = None
    try:
        outer_termios = termios.tcgetattr(0)
        tty.setraw(0)
    except termios.error:
        outer_termios = None

    child, master = pty.fork()
    if child == 0:
        os.execvp(command[0], command)
        os._exit(127)

    filt = Filter()
    set_winsize(master, 0)
    signal.signal(signal.SIGWINCH, lambda *_: set_winsize(master, 0))

    try:
        relay(child, master, filt)
    finally:
        if outer_termios is not None:
            try:
                termios.tcsetattr(0, termios.TCSANOW, outer_termios)
            except termios.error:
                pass


def relay(child, master, filt):
    stdin_open = True
    hungup = False
    while True:
        watch = [master] + ([0] if stdin_open else [])
        try:
            ready, _, _ = select.select(watch, [], [])
        except InterruptedError:
            continue
        except OSError:
            break

        if stdin_open and 0 in ready:         # terminal -> child
            try:
                data = os.read(0, 4096)
            except OSError:
                data = b""
            if not data:
                # Terminal side closed: hang up the child, keep draining
                # its output until it actually exits.
                stdin_open = False
                if not hungup:
                    hungup = True
                    try:
                        os.kill(child, signal.SIGHUP)
                    except OSError:
                        pass
            else:
                try:
                    os.write(master, data)
                except OSError:
                    break

        if master in ready:                   # child -> terminal (filtered)
            try:
                data = os.read(master, 4096)
            except OSError as e:
                if e.errno == errno.EIO:      # child exited
                    break
                raise
            if not data:
                break
            out = filt.feed(data)
            if out:
                try:
                    os.write(1, out)
                except OSError:
                    break

    try:
        os.kill(child, signal.SIGHUP)
    except OSError:
        pass
    try:
        os.waitpid(child, 0)
    except OSError:
        pass


if __name__ == "__main__":
    main()