tiny mux
This commit is contained in:
parent
9a187a8b70
commit
88bf0e2357
4
.gitignore
vendored
Normal file
4
.gitignore
vendored
Normal file
@ -0,0 +1,4 @@
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
.env
|
||||
3
test_exec.py
Normal file
3
test_exec.py
Normal file
@ -0,0 +1,3 @@
|
||||
import os
|
||||
os.execvp("ssh", ["ssh", "localhost"])
|
||||
|
||||
393
tiny_mux.py
Executable file
393
tiny_mux.py
Executable file
@ -0,0 +1,393 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tiny_mux: minimal SSH workspace switcher that multiplexes *remote* tmux sessions.
|
||||
|
||||
- Reads sessions from ~/.tiny_mux.conf (one per line).
|
||||
- Runs exactly one child at a time: ssh + tmux session.
|
||||
- Hotkeys (prefix = Ctrl-a):
|
||||
Ctrl-a <digit> : switch to session # (1-based from config file)
|
||||
Ctrl-a q : quit tiny_mux (kills current ssh)
|
||||
Ctrl-a a : send literal Ctrl-a to child
|
||||
Ctrl-a w : show popup menu, choose session
|
||||
"""
|
||||
|
||||
import errno
|
||||
import fcntl
|
||||
import os
|
||||
import selectors
|
||||
import signal
|
||||
import struct
|
||||
import sys
|
||||
import termios
|
||||
import tty
|
||||
import pty
|
||||
|
||||
CONFIG_PATH = os.path.expanduser("~/.tiny_mux.conf")
|
||||
|
||||
# ioctl constants (macOS + Linux)
|
||||
TIOCSWINSZ = termios.TIOCSWINSZ
|
||||
TIOCGWINSZ = termios.TIOCGWINSZ
|
||||
|
||||
|
||||
# ---------- VM wrapper logic ----------
|
||||
|
||||
def build_vm_command(hostname: str, skip_tmux: bool = False, extra_args=None):
|
||||
"""Reimplementation of ~/bin/vm wrapper."""
|
||||
user = os.environ.get("USER", "unknown")
|
||||
host = hostname
|
||||
ssh_identity = None
|
||||
|
||||
if "@" in host:
|
||||
user, host = host.split("@", 1)
|
||||
|
||||
tmux_session = host
|
||||
|
||||
if host.endswith("-orb"):
|
||||
ssh_host = f"{host[:-4]}@orb"
|
||||
elif host.endswith("-utm"):
|
||||
ssh_host = f"{host[:-4]}.utm.local"
|
||||
ssh_identity = os.path.expanduser("~/.ssh/id_ed25519_internal")
|
||||
elif host.endswith("-workstation"):
|
||||
ssh_host = f"{host[:-12]}.workstation.lan"
|
||||
ssh_identity = os.path.expanduser("~/.ssh/id_ed25519_internal")
|
||||
else:
|
||||
print(f"[tiny_mux] Error: unknown host pattern '{host}'", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
ssh_cmd = ["ssh", "-tt"]
|
||||
if ssh_identity:
|
||||
ssh_cmd += ["-i", ssh_identity, "-o", "IdentitiesOnly=yes"]
|
||||
ssh_cmd.append(f"{user}@{ssh_host}")
|
||||
|
||||
if not skip_tmux:
|
||||
ssh_cmd += ["tmux", "new", "-A", "-s", tmux_session]
|
||||
elif extra_args:
|
||||
ssh_cmd += extra_args
|
||||
|
||||
return ssh_cmd
|
||||
|
||||
|
||||
# ---------- Terminal helpers ----------
|
||||
|
||||
def load_sessions(path: str):
|
||||
sessions = []
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
s = line.strip()
|
||||
if not s or s.startswith("#"):
|
||||
continue
|
||||
sessions.append(s)
|
||||
except FileNotFoundError:
|
||||
print(f"[tiny_mux] Config not found: {path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
if not sessions:
|
||||
print(f"[tiny_mux] No sessions in {path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
return sessions
|
||||
|
||||
|
||||
def get_winsize(fd: int):
|
||||
try:
|
||||
data = fcntl.ioctl(fd, TIOCGWINSZ, b"\x00" * 8)
|
||||
rows, cols, xp, yp = struct.unpack("HHHH", data)
|
||||
if rows == 0 or cols == 0:
|
||||
rows = int(os.environ.get("LINES", 24)) or 24
|
||||
cols = int(os.environ.get("COLUMNS", 80)) or 80
|
||||
xp = yp = 0
|
||||
return rows, cols, xp, yp
|
||||
except OSError:
|
||||
return 24, 80, 0, 0
|
||||
|
||||
|
||||
def set_winsize(fd: int, rows: int, cols: int, xp: int = 0, yp: int = 0):
|
||||
data = struct.pack("HHHH", rows, cols, xp, yp)
|
||||
fcntl.ioctl(fd, TIOCSWINSZ, data)
|
||||
|
||||
|
||||
class RawTTY:
|
||||
def __init__(self, fd: int):
|
||||
self.fd = fd
|
||||
self._saved = None
|
||||
|
||||
def __enter__(self):
|
||||
self._saved = termios.tcgetattr(self.fd)
|
||||
tty.setraw(self.fd, when=termios.TCSANOW)
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
if self._saved is not None:
|
||||
termios.tcsetattr(self.fd, termios.TCSADRAIN, self._saved)
|
||||
|
||||
|
||||
# ---------- TinyMux core ----------
|
||||
|
||||
class TinyMux:
|
||||
def __init__(self, sessions):
|
||||
self.sessions = sessions
|
||||
self.current_idx = None
|
||||
|
||||
self.master_fd = None
|
||||
self.child_pid = None
|
||||
self._selector = selectors.DefaultSelector()
|
||||
|
||||
self.stdin_fd = sys.stdin.fileno()
|
||||
self.stdout_fd = sys.stdout.fileno()
|
||||
|
||||
self.prefix = b"\x01" # Ctrl-a
|
||||
self._saw_prefix = False
|
||||
|
||||
# Save original cooked tty mode
|
||||
self._saved_tty = termios.tcgetattr(self.stdin_fd)
|
||||
|
||||
signal.signal(signal.SIGWINCH, self._on_sigwinch)
|
||||
signal.signal(signal.SIGCHLD, self._on_sigchld)
|
||||
|
||||
# ---------- Process / PTY management ----------
|
||||
|
||||
def _close_master(self):
|
||||
if self.master_fd is not None:
|
||||
try:
|
||||
self._selector.unregister(self.master_fd)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
os.close(self.master_fd)
|
||||
except OSError:
|
||||
pass
|
||||
self.master_fd = None
|
||||
|
||||
def _kill_child_group(self):
|
||||
if self.child_pid:
|
||||
try:
|
||||
os.killpg(self.child_pid, signal.SIGHUP)
|
||||
except Exception:
|
||||
try:
|
||||
os.kill(self.child_pid, signal.SIGHUP)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _reap_child(self):
|
||||
if self.child_pid:
|
||||
try:
|
||||
while True:
|
||||
pid, _ = os.waitpid(self.child_pid, os.WNOHANG)
|
||||
if pid == 0:
|
||||
break
|
||||
if pid == self.child_pid:
|
||||
self.child_pid = None
|
||||
break
|
||||
except ChildProcessError:
|
||||
self.child_pid = None
|
||||
|
||||
def _spawn(self, idx: int):
|
||||
if self.child_pid:
|
||||
self._kill_child_group()
|
||||
self._reap_child()
|
||||
self._close_master()
|
||||
|
||||
master_fd, slave_fd = pty.openpty()
|
||||
rows, cols, xp, yp = get_winsize(self.stdout_fd)
|
||||
set_winsize(slave_fd, rows, cols, xp, yp)
|
||||
|
||||
pid = os.fork()
|
||||
if pid == 0:
|
||||
# --- Child ---
|
||||
try:
|
||||
try:
|
||||
os.setsid()
|
||||
except OSError as e:
|
||||
if e.errno != errno.EPERM:
|
||||
raise
|
||||
|
||||
os.close(master_fd)
|
||||
os.dup2(slave_fd, 0)
|
||||
os.dup2(slave_fd, 1)
|
||||
os.dup2(slave_fd, 2)
|
||||
if slave_fd > 2:
|
||||
os.close(slave_fd)
|
||||
|
||||
session_name = self.sessions[idx]
|
||||
argv = build_vm_command(session_name, skip_tmux=False)
|
||||
os.execvp(argv[0], argv)
|
||||
except Exception as e:
|
||||
try:
|
||||
os.write(2, f"tiny_mux child exec failed: {e}\n".encode())
|
||||
except Exception:
|
||||
pass
|
||||
os._exit(1)
|
||||
|
||||
# --- Parent ---
|
||||
os.close(slave_fd)
|
||||
self.child_pid = pid
|
||||
self.master_fd = master_fd
|
||||
self.current_idx = idx
|
||||
try:
|
||||
os.setpgid(self.child_pid, self.child_pid)
|
||||
except Exception:
|
||||
pass
|
||||
self._selector.register(self.master_fd, selectors.EVENT_READ)
|
||||
self._propagate_winsize()
|
||||
self._signal_winch()
|
||||
|
||||
def _propagate_winsize(self):
|
||||
if self.master_fd is None:
|
||||
return
|
||||
rows, cols, xp, yp = get_winsize(self.stdout_fd)
|
||||
try:
|
||||
set_winsize(self.master_fd, rows, cols, xp, yp)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def _signal_winch(self):
|
||||
if self.child_pid:
|
||||
try:
|
||||
os.killpg(self.child_pid, signal.SIGWINCH)
|
||||
except Exception:
|
||||
try:
|
||||
os.kill(self.child_pid, signal.SIGWINCH)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _on_sigwinch(self, _signum, _frame):
|
||||
self._propagate_winsize()
|
||||
self._signal_winch()
|
||||
|
||||
def _on_sigchld(self, _signum, _frame):
|
||||
self._reap_child()
|
||||
|
||||
# ---------- Input handling ----------
|
||||
|
||||
def _choose_session(self):
|
||||
"""Close ssh, clear screen, show session list, wait for choice, reconnect."""
|
||||
self._kill_child_group()
|
||||
self._close_master()
|
||||
self._reap_child()
|
||||
|
||||
# Clear terminal
|
||||
sys.stdout.write("\033[2J\033[H")
|
||||
sys.stdout.flush()
|
||||
|
||||
# Temporarily cooked mode
|
||||
termios.tcsetattr(self.stdin_fd, termios.TCSADRAIN, self._saved_tty)
|
||||
try:
|
||||
print("[tiny_mux] Available sessions:")
|
||||
for i, s in enumerate(self.sessions, 1):
|
||||
mark = " *" if self.current_idx == i - 1 else ""
|
||||
print(f" {i}. {s}{mark}")
|
||||
sys.stdout.write("\nType number (or q to cancel): ")
|
||||
sys.stdout.flush()
|
||||
|
||||
ch = sys.stdin.read(1)
|
||||
if ch.isdigit():
|
||||
idx = int(ch) - 1
|
||||
if 0 <= idx < len(self.sessions):
|
||||
self._spawn(idx)
|
||||
else:
|
||||
# If canceled, just respawn previous session
|
||||
if self.current_idx is not None:
|
||||
self._spawn(self.current_idx)
|
||||
finally:
|
||||
tty.setraw(self.stdin_fd, when=termios.TCSANOW)
|
||||
|
||||
def _handle_stdin(self):
|
||||
try:
|
||||
data = os.read(self.stdin_fd, 1024)
|
||||
except OSError as e:
|
||||
if e.errno in (errno.EAGAIN, errno.EINTR):
|
||||
return True
|
||||
return False
|
||||
if not data:
|
||||
return False
|
||||
|
||||
out = bytearray()
|
||||
i = 0
|
||||
while i < len(data):
|
||||
b = data[i:i+1]
|
||||
i += 1
|
||||
if self._saw_prefix:
|
||||
self._saw_prefix = False
|
||||
if b == b"q":
|
||||
return False
|
||||
elif b == b"a":
|
||||
out.extend(self.prefix)
|
||||
elif b.isdigit():
|
||||
d = b[0] - ord('0')
|
||||
idx = (9 if d == 0 else d - 1)
|
||||
if 0 <= idx < len(self.sessions):
|
||||
self._spawn(idx)
|
||||
elif b == b"w":
|
||||
self._choose_session()
|
||||
elif b == self.prefix:
|
||||
self._saw_prefix = True
|
||||
else:
|
||||
out.extend(b)
|
||||
|
||||
if out and self.master_fd is not None:
|
||||
try:
|
||||
os.write(self.master_fd, out)
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _handle_master(self):
|
||||
try:
|
||||
chunk = os.read(self.master_fd, 4096)
|
||||
except OSError as e:
|
||||
if e.errno in (errno.EIO, errno.EBADF):
|
||||
return False
|
||||
if e.errno in (errno.EAGAIN, errno.EINTR):
|
||||
return True
|
||||
return False
|
||||
if not chunk:
|
||||
return False
|
||||
try:
|
||||
os.write(self.stdout_fd, chunk)
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def run(self):
|
||||
with RawTTY(self.stdin_fd):
|
||||
self._selector.register(self.stdin_fd, selectors.EVENT_READ)
|
||||
self._spawn(0)
|
||||
while True:
|
||||
events = self._selector.select(timeout=None)
|
||||
for key, _mask in events:
|
||||
if key.fileobj == self.stdin_fd:
|
||||
if not self._handle_stdin():
|
||||
self.shutdown()
|
||||
return
|
||||
elif key.fileobj == self.master_fd:
|
||||
if not self._handle_master():
|
||||
self.shutdown()
|
||||
return
|
||||
|
||||
def shutdown(self):
|
||||
try:
|
||||
self._selector.unregister(self.stdin_fd)
|
||||
except Exception:
|
||||
pass
|
||||
self._kill_child_group()
|
||||
self._close_master()
|
||||
self._reap_child()
|
||||
|
||||
|
||||
def main():
|
||||
sessions = load_sessions(CONFIG_PATH)
|
||||
mux = TinyMux(sessions)
|
||||
try:
|
||||
mux.run()
|
||||
except KeyboardInterrupt:
|
||||
mux.shutdown()
|
||||
except Exception as e:
|
||||
try:
|
||||
mux.shutdown()
|
||||
finally:
|
||||
print(f"[tiny_mux] fatal error: {e}", file=sys.stderr)
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user