#!/usr/bin/env python3
"""Dokki Terminal's attach client for Linux hosts (the twin of `DokkiHelper pty-attach`;
docs/plans/2026-10-22-terminal-attach.md C1-C3).

    dokki-pty-attach --ticket-file <abs path>        dokki-pty-attach --probe

The computer's own terminal window as one more client of a PTY session in the host. The host
opens the terminal with this script and the PATH of a ticket file; the single-use token is
inside the file, never in argv, the environment or a log.

Trust boundary (C1). Any failure exits 2 with one line on stderr and touches nothing:
- stdin is a tty (else "not a terminal", the ticket untouched);
- the basename is exactly `a-<32 lowercase hex>.json`;
- the folder resolves to one of the host's two runtime folders: `$XDG_RUNTIME_DIR/dokki-term`
  (when XDG_RUNTIME_DIR is absolute, a directory of ours, 0700) or `/tmp/dokki-term-<uid>`;
  it is opened with O_NOFOLLOW | O_DIRECTORY and judged on that descriptor: a directory,
  ours, no group/other bits (and, with /proc, really that folder);
- the file is opened once, relative to the folder's descriptor, with O_NOFOLLOW (a symlink
  fails with ELOOP); every check uses fstat on the descriptor: regular, ours, exactly 0600,
  1..512 bytes; it is read from the same descriptor. No path-based stat of the file exists;
- the JSON has exactly v (1), endpoint (<that folder>/attach.sock), token (43 base64url) and
  session (s_<16 hex>);
- only then is the file unlinked (relative to the folder), before connecting.

Wire (C2): src/term/frame.js's `type u8 | length u32be | payload`, only 1 DATA (<= 64 KiB)
and 6 CONTROL (a JSON object, <= 4096 bytes). `hello` first, with the window size; keys go
through the detach key (Ctrl-] d, C3) then out as DATA; SIGWINCH is coalesced over 50 ms into
`resize`. The host's DATA goes to stdout (blocking writes); `ready` sets the window title;
`size` is ignored; `exit` / `error` / EOF end the client.

Leaving on its own (detach, exit, error, host gone): RESET and the host's label are written
while still raw (byte for byte), then the tty is restored; exit 0 / the program's code / 1 / 1.
The window closing (SIGHUP, SIGTERM, or the tty read failing): the `hangup` frame goes FIRST,
after any frame already half sent, within 200 ms; then a best-effort tty restore (EIO/EBADF
are expected); nothing is written to stdout; exit 129 / 143. Signals only wake the loop (the
wakeup fd); all work happens in the loop, and the socket is non-blocking, so a busy host can
never hold up a hang-up.

python3 stdlib only. Unit and end-to-end tests: native/linux/tests/test_dokki_pty_attach.py.
"""

import errno
import fcntl
import json
import os
import re
import select
import selectors
import signal
import socket
import stat
import struct
import sys
import termios
import time
import tty

T_DATA, T_CONTROL = 1, 6
HEADER = 5
MAX_DATA = 64 * 1024
MAX_CONTROL = 4096
MAX_TICKET = 512
MAX_COLS, MAX_ROWS = 1000, 500
DETACH = 0x1D
DETACH_TIMEOUT_MS = 1000
RESIZE_COALESCE_MS = 50
HANGUP_SEND_MS = 200
OUT_HIGH_WATER = 256 * 1024
PROBE = '{"ok":true,"v":1}'
SOCKET_NAME = "attach.sock"

# Identical in the three clients (DokkiCore/PtyAttach.swift, native/win/PtyAttach.Rules.cs).
# `ESC 7` first: leaving the alternate screen restores the saved cursor, and when the program never entered it (most agents) Terminal.app jumps to a stale saved position and the label overwrites old lines; saving the cursor first makes that a no-op (real E2E, 2026-10-22).
RESET = (b"\x1b7\x1b[?1049l\x1b[?25h\x1b[?2004l\x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l"
         b"\x1b[?1004l\x1b[<u\x1b[?1l\x1b>\x1b[0m\r\n")

ENGLISH = {
    "exit": "The program ended (code {code}).",
    "detached": "Detached. The session keeps running in Dokki: continue it on your phone or with Open on the computer.",
    "ended": "This session had already ended.",
    "hostGone": "Dokki closed, and the session ended with it.",
    "error": "Dokki couldn't connect this window ({code}).",
}

TICKET_NAME = re.compile(r"a-[0-9a-f]{32}\.json")
TOKEN = re.compile(r"[A-Za-z0-9_-]{43}")
SESSION = re.compile(r"s_[0-9a-f]{16}")
ERROR_CODE = re.compile(r"[a-z-]{1,32}")


class TicketError(ValueError):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code


class FrameError(ValueError):
    pass


# ---- arguments ---------------------------------------------------------------

def is_clean_abs_path(path):
    """Absolute, <= 1024 bytes, no control characters, no empty, `.` or `..` segment."""
    if not isinstance(path, str) or not path.startswith("/") or len(path) < 2 or len(path.encode("utf-8", "surrogateescape")) > 1024:
        return False
    if any(ord(c) < 0x20 or 0x7F <= ord(c) <= 0x9F or c in "  " for c in path):
        return False
    parts = path.split("/")[1:]
    if parts and parts[-1] == "":
        parts = parts[:-1]
    return bool(parts) and not any(p in ("", ".", "..") for p in parts)


def parse_args(argv):
    """("probe",) | ("ticket", path) | None."""
    if argv == ["--probe"]:
        return ("probe",)
    if len(argv) == 2 and argv[0] == "--ticket-file" and is_clean_abs_path(argv[1]) and not argv[1].endswith("/"):
        return ("ticket", argv[1])
    return None


# ---- the ticket file (C1) ----------------------------------------------------

class HostFS:
    """What reading a ticket needs. Files are only looked at through a descriptor."""

    def uid(self):
        return os.getuid()

    def realpath(self, path):
        return os.path.realpath(path)

    def lstat_dir(self, path):
        """Only for XDG_RUNTIME_DIR itself (which folder to expect), never for the ticket."""
        try:
            return os.lstat(path)
        except OSError:
            return None

    def open_dir(self, path):
        return os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC)

    def open_file(self, dir_fd, name):
        # O_NONBLOCK: a FIFO planted under the name must not hang the open (fstat refuses it next).
        return os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC | os.O_NONBLOCK, dir_fd=dir_fd)

    def fstat(self, fd):
        return os.fstat(fd)

    def fd_path(self, fd):
        """Where a descriptor really is (/proc), or None where there is no /proc."""
        try:
            return os.readlink("/proc/self/fd/%d" % fd)
        except OSError:
            return None

    def read(self, fd, limit):
        out = b""
        while len(out) < limit:
            chunk = os.read(fd, limit - len(out))
            if not chunk:
                break
            out += chunk
        return out

    def close(self, fd):
        os.close(fd)

    def unlink(self, dir_fd, name):
        os.unlink(name, dir_fd=dir_fd)


def expected_folders(environ, fs):
    """The host's runtime folders (Decision 2), resolved: XDG first when it qualifies, then /tmp."""
    uid = fs.uid()
    out = []
    xdg = environ.get("XDG_RUNTIME_DIR", "")
    if xdg.startswith("/"):
        st = fs.lstat_dir(xdg)
        if st is not None and stat.S_ISDIR(st.st_mode) and st.st_uid == uid and (st.st_mode & 0o777) == 0o700:
            out.append(fs.realpath(xdg).rstrip("/") + "/dokki-term")
    out.append(fs.realpath("/tmp").rstrip("/") + "/dokki-term-%d" % uid)
    return out


def is_integer(value, want=None):
    if isinstance(value, bool):
        return False
    if isinstance(value, float):
        if not value.is_integer():
            return False
        value = int(value)
    return isinstance(value, int) and (want is None or value == want)


def parse_ticket(data):
    """The JSON object with exactly v, endpoint, token, session, or None."""
    if not data.startswith(b"{"):
        return None
    try:
        obj = json.loads(data.decode("utf-8"))
    except (UnicodeDecodeError, ValueError):
        return None
    if not isinstance(obj, dict) or set(obj) != {"v", "endpoint", "token", "session"} or not is_integer(obj["v"], 1):
        return None
    endpoint, token, session = obj["endpoint"], obj["token"], obj["session"]
    if not all(isinstance(v, str) for v in (endpoint, token, session)):
        return None
    if not TOKEN.fullmatch(token) or not SESSION.fullmatch(session):
        return None
    return {"endpoint": endpoint, "token": token, "session": session}


def read_ticket(path, environ=None, fs=None, folders=None):
    """C1 in order; the file is unlinked only when every check passed. Raises TicketError."""
    fs = fs or HostFS()
    if not is_clean_abs_path(path) or path.endswith("/"):
        raise TicketError("arguments", "bad ticket path")
    parent, name = os.path.split(path)
    if not TICKET_NAME.fullmatch(name):
        raise TicketError("basename", "bad ticket name")
    allowed = folders if folders is not None else expected_folders(os.environ if environ is None else environ, fs)
    real_parent = fs.realpath(parent)
    if real_parent not in allowed:
        raise TicketError("folder", "the ticket is not in Dokki's private folder")
    try:
        dir_fd = fs.open_dir(parent)
    except OSError:
        raise TicketError("folder", "the ticket folder could not be opened (a symlink?)")
    try:
        st = fs.fstat(dir_fd)
        where = fs.fd_path(dir_fd)
        if (not stat.S_ISDIR(st.st_mode) or st.st_uid != fs.uid() or st.st_mode & 0o077
                or (where is not None and where != real_parent)):
            raise TicketError("folder", "the ticket folder is not private")
        try:
            fd = fs.open_file(dir_fd, name)
        except OSError:
            raise TicketError("open", "the ticket file could not be opened")
        try:
            st = fs.fstat(fd)
            if (not stat.S_ISREG(st.st_mode) or st.st_uid != fs.uid() or (st.st_mode & 0o7777) != 0o600
                    or not 0 < st.st_size <= MAX_TICKET):
                raise TicketError("file", "the ticket file is not a private regular file")
            try:
                data = fs.read(fd, MAX_TICKET + 1)
            except OSError:
                raise TicketError("file", "the ticket file could not be read")
            if not 0 < len(data) <= MAX_TICKET:
                raise TicketError("file", "the ticket file is not a private regular file")
        finally:
            fs.close(fd)
        ticket = parse_ticket(data)
        if ticket is None:
            raise TicketError("json", "bad ticket content")
        if ticket["endpoint"] not in (parent + "/" + SOCKET_NAME, real_parent + "/" + SOCKET_NAME):
            raise TicketError("endpoint", "bad ticket endpoint")
        try:
            fs.unlink(dir_fd, name)
        except OSError:
            raise TicketError("unlink", "the ticket file could not be removed")
        return ticket
    finally:
        fs.close(dir_fd)


# ---- frames (C2) -------------------------------------------------------------

def frame(ftype, payload):
    return struct.pack(">BI", ftype, len(payload)) + bytes(payload)


def encode_control(obj):
    """A CONTROL frame, or None past 4096 bytes of JSON."""
    body = json.dumps(obj, separators=(",", ":"), sort_keys=True, ensure_ascii=False).encode("utf-8")
    return frame(T_CONTROL, body) if len(body) <= MAX_CONTROL else None


def encode_data(data):
    if not data:
        return frame(T_DATA, b"")
    return b"".join(frame(T_DATA, data[i:i + MAX_DATA]) for i in range(0, len(data), MAX_DATA))


DETACH_FRAME = encode_control({"t": "detach"})
HANGUP_FRAME = encode_control({"t": "hangup"})


class Decoder:
    """The host's frames: ("data", bytes) or ("control", dict). Fails for good on the first bad one."""

    def __init__(self):
        self.buf = bytearray()
        self.failed = None

    def feed(self, chunk):
        if self.failed:
            raise self.failed
        self.buf += chunk
        out = []
        try:
            while len(self.buf) >= HEADER:
                ftype, length = struct.unpack_from(">BI", self.buf, 0)
                if ftype not in (T_DATA, T_CONTROL):
                    raise FrameError("unknown frame type %d" % ftype)
                if length > (MAX_DATA if ftype == T_DATA else MAX_CONTROL):
                    raise FrameError("frame too large")
                if len(self.buf) < HEADER + length:
                    break
                payload = bytes(self.buf[HEADER:HEADER + length])
                del self.buf[:HEADER + length]
                if ftype == T_DATA:
                    out.append(("data", payload))
                    continue
                try:
                    obj = json.loads(payload.decode("utf-8"))
                except (UnicodeDecodeError, ValueError):
                    raise FrameError("bad control json")
                if not isinstance(obj, dict):
                    raise FrameError("control is not an object")
                out.append(("control", obj))
        except FrameError as error:
            self.failed = error
            self.buf = bytearray()
            raise
        return out


# ---- host messages, labels -----------------------------------------------------

def _stripped(ch):
    v = ord(ch)
    return (v < 0x20 or 0x7F <= v <= 0x9F or v in (0x061C, 0x200E, 0x200F, 0x2028, 0x2029)
            or 0x202A <= v <= 0x202E or 0x2066 <= v <= 0x2069)


def sanitize(value, limit):
    out = []
    for ch in value:
        if not _stripped(ch):
            out.append(ch)
            if len(out) == limit:
                break
    return "".join(out)


def sanitize_label(value):
    return sanitize(value, 200)


def sanitize_title(value):
    return sanitize(value, 40)


def title_sequence(name):
    return ("\x1b]0;Dokki · %s\x07" % sanitize_title(name)).encode("utf-8")


def labels_from(raw):
    out = dict(ENGLISH)
    if isinstance(raw, dict):
        for key in ENGLISH:
            value = raw.get(key)
            if isinstance(value, str):
                clean = sanitize_label(value)
                if clean:
                    out[key] = clean
    return out


def fill(template, code):
    return template.replace("{code}", str(code))


def message(obj):
    t = obj.get("t")
    if t == "ready":
        name = obj.get("name")
        return ("ready", name if isinstance(name, str) else "", labels_from(obj.get("labels")))
    if t == "size":
        return ("size",)
    if t == "exit":
        code = obj.get("code")
        if not is_integer(code):
            return ("exit", 1)
        return ("exit", max(0, min(255, int(code))))
    if t == "error":
        code = obj.get("code")
        return ("error", code if isinstance(code, str) and ERROR_CODE.fullmatch(code) else "bad-message")
    return ("other",)


def clamp_size(cols, rows):
    if cols < 1 or rows < 1:
        return (80, 24)
    return (min(cols, MAX_COLS), min(rows, MAX_ROWS))


# ---- the detach key (C3) -------------------------------------------------------

class DetachKey:
    """Ctrl-] then d/D detaches; Ctrl-] Ctrl-] sends one; Ctrl-] and any other byte send both;
    a lone Ctrl-] with nothing after it for 1000 ms is sent as is."""

    def __init__(self):
        self.since = None

    @property
    def pending(self):
        return self.since is not None

    @property
    def deadline(self):
        return None if self.since is None else self.since + DETACH_TIMEOUT_MS

    def feed(self, data, now_ms):
        out = bytearray(self.tick(now_ms))
        for byte in data:
            if self.since is not None:
                self.since = None
                if byte in (0x64, 0x44):
                    return bytes(out), True
                if byte == DETACH:
                    out.append(DETACH)
                else:
                    out += bytes((DETACH, byte))
            elif byte == DETACH:
                self.since = now_ms
            else:
                out.append(byte)
        return bytes(out), False

    def tick(self, now_ms):
        if self.since is None or now_ms - self.since < DETACH_TIMEOUT_MS:
            return b""
        self.since = None
        return bytes((DETACH,))


# ---- the client ----------------------------------------------------------------

def write_all(fd, data):
    view = memoryview(data)
    while view:
        try:
            n = os.write(fd, view)
        except InterruptedError:
            continue
        except OSError:
            return False
        view = view[n:]
    return True


def window_size(fd):
    try:
        rows, cols, _, _ = struct.unpack("HHHH", fcntl.ioctl(fd, termios.TIOCGWINSZ, b"\0" * 8))
    except OSError:
        return clamp_size(0, 0)
    return clamp_size(cols, rows)


class Client:
    def __init__(self, sock, saved, stdin=0, stdout=1):
        self.sock = sock
        self.saved = saved
        self.stdin = stdin
        self.stdout = stdout
        self.out = bytearray()          # frames not yet on the socket (whole frames, or the rest of one)
        self.labels = dict(ENGLISH)
        self.key = DetachKey()
        self.decoder = Decoder()
        self.size = (80, 24)
        self.resize_at = None
        self.sel = selectors.DefaultSelector()
        self.stdin_on = False

    @staticmethod
    def now_ms():
        return int(time.monotonic() * 1000)

    # -- socket output: non-blocking, so nothing in the loop ever waits on the host --
    def queue(self, data):
        self.out += data
        self.flush()

    def flush(self):
        while self.out:
            try:
                n = self.sock.send(self.out)
            except (BlockingIOError, InterruptedError):
                break
            except OSError:
                self.out.clear()  # the host is gone: the read side reports it
                break
            del self.out[:n]
        self.update()

    def update(self):
        self.sel.modify(self.sock, selectors.EVENT_READ | (selectors.EVENT_WRITE if self.out else 0), "sock")
        want = len(self.out) < OUT_HIGH_WATER
        if want != self.stdin_on:
            if want:
                self.sel.register(self.stdin, selectors.EVENT_READ, "stdin")
            else:
                self.sel.unregister(self.stdin)
            self.stdin_on = want

    def drain(self, ms):
        """Blocking flush with a deadline (a full buffer can't hold up the exit)."""
        deadline = self.now_ms() + ms
        while self.out:
            left = deadline - self.now_ms()
            if left <= 0:
                return False
            try:
                n = self.sock.send(self.out)
            except (BlockingIOError, InterruptedError):
                select.select([], [self.sock], [], left / 1000.0)
                continue
            except OSError:
                return False
            del self.out[:n]
        return True

    def run(self, token):
        wake_r, wake_w = os.pipe()
        for fd in (wake_r, wake_w):
            fcntl.fcntl(fd, fcntl.F_SETFL, fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK)
            fcntl.fcntl(fd, fcntl.F_SETFD, fcntl.FD_CLOEXEC)
        signal.set_wakeup_fd(wake_w)
        for sig in (signal.SIGWINCH, signal.SIGHUP, signal.SIGTERM):
            signal.signal(sig, lambda *_: None)  # only wakes the loop (the wakeup fd says which)
        self.sock.setblocking(False)
        self.sel.register(self.sock, selectors.EVENT_READ, "sock")
        self.sel.register(wake_r, selectors.EVENT_READ, "wake")
        self.size = window_size(self.stdin)
        self.queue(encode_control({"t": "hello", "v": 1, "token": token, "cols": self.size[0], "rows": self.size[1]}))
        tty.setraw(self.stdin, termios.TCSANOW)
        self.resize_at = self.now_ms()  # a resize that came before the handler existed
        while True:
            now = self.now_ms()
            deadlines = [d for d in (self.key.deadline, self.resize_at) if d is not None]
            timeout = max(0, min(deadlines) - now) / 1000.0 if deadlines else None
            events = self.sel.select(timeout)
            events.sort(key=lambda e: e[0].data != "wake")  # signals first
            for key, mask in events:
                if key.data == "wake":
                    try:
                        signals = os.read(wake_r, 256)
                    except BlockingIOError:
                        signals = b""
                    for sig in signals:
                        if sig in (signal.SIGHUP, signal.SIGTERM):
                            self.hangup(sig)
                        if sig == signal.SIGWINCH and self.resize_at is None:
                            self.resize_at = self.now_ms() + RESIZE_COALESCE_MS
                elif key.data == "stdin":
                    try:
                        data = os.read(self.stdin, 4096)
                    except (BlockingIOError, InterruptedError):
                        continue
                    except OSError:
                        data = b""
                    if not data:
                        self.hangup(signal.SIGHUP)  # the terminal was revoked: the window is closing
                    forward, detach = self.key.feed(data, self.now_ms())
                    if forward:
                        self.queue(encode_data(forward))
                    if detach:
                        self.finish("detach")
                else:
                    if mask & selectors.EVENT_WRITE:
                        self.flush()
                    if mask & selectors.EVENT_READ:
                        self.read_socket()
            now = self.now_ms()
            late = self.key.tick(now)
            if late:
                self.queue(encode_data(late))
            if self.resize_at is not None and now >= self.resize_at:
                self.resize_at = None
                size = window_size(self.stdin)
                if size != self.size:
                    self.size = size
                    self.queue(encode_control({"t": "resize", "cols": size[0], "rows": size[1]}))

    def read_socket(self):
        try:
            data = self.sock.recv(65536)
        except (BlockingIOError, InterruptedError):
            return
        except OSError:
            data = b""
        if not data:
            self.finish("hostGone")
        try:
            frames = self.decoder.feed(data)
        except FrameError:
            self.finish("error", "bad-message")
        for kind, value in frames:
            if kind == "data":
                write_all(self.stdout, value)
                continue
            msg = message(value)
            if msg[0] == "ready":
                self.labels = msg[2]
                write_all(self.stdout, title_sequence(msg[1]))
            elif msg[0] == "exit":
                self.finish("exit", msg[1])
            elif msg[0] == "error":
                self.finish("error", msg[1])

    def finish(self, end, value=None):
        if end == "detach":
            self.queue(DETACH_FRAME)
            self.drain(1000)
            text, code = self.labels["detached"], 0
        elif end == "exit":
            text, code = fill(self.labels["exit"], value), value
        elif end == "error":
            text = self.labels["ended"] if value in ("unknown-session", "gone") else fill(self.labels["error"], value)
            code = 1
        else:
            text, code = self.labels["hostGone"], 1
        # Still raw: the bytes land exactly as written; then the modes come back.
        write_all(self.stdout, RESET + text.encode("utf-8") + b"\r\n")
        try:
            termios.tcsetattr(self.stdin, termios.TCSADRAIN, self.saved)
        except termios.error:
            pass
        try:
            self.sock.close()
        except OSError:
            pass
        os._exit(code)

    def hangup(self, sig):
        """The window is closing: the frame first, then a best-effort restore, nothing on stdout."""
        self.out += HANGUP_FRAME
        self.drain(HANGUP_SEND_MS)
        try:
            termios.tcsetattr(self.stdin, termios.TCSANOW, self.saved)
        except (termios.error, OSError):
            pass
        os._exit(128 + sig)


def attach(ticket, saved, stdin=0, stdout=1):
    # Until the loop's handlers exist a signal is ignored rather than fatal; a window that
    # closed meanwhile is still noticed (the tty read fails) and hung up properly.
    for sig in (signal.SIGPIPE, signal.SIGHUP, signal.SIGTERM, signal.SIGWINCH):
        signal.signal(sig, signal.SIG_IGN)
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(ticket["endpoint"])
    except OSError:
        sock.close()
        sys.stderr.write(fill(ENGLISH["error"], "unavailable") + "\n")
        return 1
    Client(sock, saved, stdin, stdout).run(ticket["token"])
    return 1


def main(argv, folders=None, environ=None):
    parsed = parse_args(argv)
    if parsed is None:
        sys.stderr.write("usage: dokki-pty-attach --ticket-file <absolute path> | --probe\n")
        return 2
    if parsed[0] == "probe":
        sys.stdout.write(PROBE + "\n")
        sys.stdout.flush()
        return 0
    if not os.isatty(0):
        sys.stderr.write("dokki-pty-attach: not a terminal\n")
        return 2
    try:
        saved = termios.tcgetattr(0)
        ticket = read_ticket(parsed[1], environ=environ, folders=folders)
    except termios.error:
        sys.stderr.write("dokki-pty-attach: not a terminal\n")
        return 2
    except TicketError as error:
        sys.stderr.write("dokki-pty-attach: %s\n" % error)
        return 2
    return attach(ticket, saved)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
