From d6e0fd40256a082cbc2cf2734a9d5cf315056c22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Gremaud?= Date: Wed, 29 May 2024 15:27:19 +0200 Subject: [PATCH] Update tmux-ssh-groups --- tmux_ssh_group.py | 155 ++++++++++++++++++++++++++++++---------------- 1 file changed, 103 insertions(+), 52 deletions(-) diff --git a/tmux_ssh_group.py b/tmux_ssh_group.py index 849c42f..ea6f9db 100644 --- a/tmux_ssh_group.py +++ b/tmux_ssh_group.py @@ -1,5 +1,6 @@ import libtmux import configparser +import argparse import time import sys import os @@ -9,6 +10,16 @@ from pathlib import Path from pyfzf.pyfzf import FzfPrompt +parser = argparse.ArgumentParser() +parser.add_argument("--no-open", action="store_true", help="Do not open tmux window. Mostly for debugging") +parser.add_argument("--separate-sessions", "-s", action="store_true", help="Use a separate sessions for all of the groups") +parser_args = parser.parse_args() + + +def get_window_name(session_name): + name = session_name.replace("-", "").replace(" ", " ").replace(" ", "_") + return f"{'' if parser_args.separate_sessions else 'ssh-multig '}{name}" + BASE_PATH = Path("~/.ssh-tmux/").expanduser() configs = {} @@ -16,76 +27,116 @@ sections = {} for subconffile in os.listdir(BASE_PATH): if not subconffile.endswith(".ini"): continue + + prefix = subconffile.replace(".ini", "").capitalize() config = configparser.ConfigParser() config.read(BASE_PATH / subconffile) configs[subconffile] = config for section in config.sections(): - sections[section] = subconffile - + section_name = f"[{prefix}] {section}" + sections[section_name] = { + "section": section, + "config": subconffile, + } fzf = FzfPrompt() -selections = fzf.prompt(sections.keys(), "--cycle") +selections = fzf.prompt(["Select one or more servers group to open"] + list(sections.keys()), "--cycle --multi --header-lines 1") + if not selections: sys.exit(0) -selected_group = selections[0] - -config = configs[sections[selected_group]] - -servers = config[selected_group]["servers"].split("\n") - -extra_commands = [] -if "commands" in config[selected_group].keys(): - extra_commands = config[selected_group]["commands"].split("\n") - -if not servers: - sys.exit(1) - +# Get active tmux sessions srv = libtmux.Server() -active_sessions = srv.sessions.filter(session_attached='1') +active_sessions = srv.attached_sessions -if active_sessions: - active_session = active_sessions[0] +if parser_args.separate_sessions: + sessions = srv.sessions.filter(name="SSH-MultiG") -else: - raw_try = srv.cmd("display", "-p", "#{session_name}").stdout - if raw_try: - active_session = srv.sessions.filter(name=raw_try[0])[0] + if sessions: + active_session = sessions[0] else: - active_session = srv.sessions[0] + active_session = srv.new_session(session_name="SSH-MultiG") -window_name = f"ssh-multig {selected_group.replace(' ', '_')}" -if windows := active_session.windows.filter(name=window_name): - windows[0].select() - sys.exit(0) +else: + if len(active_sessions) > 1: + session_choice = fzf.prompt(["You have multiple active tmux sessions open. Choose one to create the window on"] + [sess.name for sess in active_sessions], "--cycle --header-lines 1") -window = active_session.new_window(window_name, window_shell=f"ssh {servers[0]}") + if not session_choice: + sys.exit(0) -for server in servers[1:]: - pane = window.split(shell=f"ssh {server}") - # Select tiled layout each time, to ensure enough space + active_session = srv.sessions.filter(name=session_choice[0])[0] + + elif len(active_sessions) == 1: + active_session = active_sessions[0] + + else: + raw_try = srv.cmd("display", "-p", "#{session_name}").stdout + if raw_try: + active_session = srv.sessions.filter(name=raw_try[0])[0] + + else: + active_session = srv.sessions[0] + +for selected_group in selections: + + info = sections[selected_group] + config = configs[info["config"]] + servers = config[info["section"]]["servers"].split("\n") + + extra_commands = [] + if "commands" in config[info["section"]].keys(): + extra_commands = config[info["section"]]["commands"].split("\n") + + if parser_args.no_open: + print(selected_group) + print(f" - Servers ({len(servers)}):") + for server in servers: + print(f" * {server}") + + if extra_commands: + print(f" - Extra commands ({len(extra_commands)}):") + for command in extra_commands: + print(f" * {command}") + + continue + + if not servers: + continue + + window_name = get_window_name(info["section"]) + if windows := active_session.windows.filter(name=window_name): + windows[0].select() + continue + + window = active_session.new_window(window_name, window_shell=f"ssh {servers[0]}") + + for server in servers[1:]: + pane = window.split(shell=f"ssh {server}") + # Select tiled layout each time, to ensure enough space + window.select_layout("tiled") + + # Wait until tmux finished working + time.sleep(0.05) + + # Confirm connection on asking panes + confirmation_needed_text = "Are you sure you want to continue connecting (yes/no/[fingerprint])?" + for pane in window.panes: + pane_content = pane.capture_pane() + if pane_content and confirmation_needed_text == pane_content[-1]: + pane.send_keys("yes") + + window.set_window_option("synchronize-panes", "on") + pane = window.panes[0] + pane.send_keys("sudo su -") + + for command in extra_commands: + pane.send_keys(command) + + window.set_window_option("synchronize-panes", "off") + window.select() window.select_layout("tiled") -# Wait until tmux finished working -time.sleep(0.1) - -# Confirm connection on asking panes -confirmation_needed_text = "Are you sure you want to continue connecting (yes/no/[fingerprint])?" -for pane in window.panes: - pane_content = pane.capture_pane() - if pane_content and confirmation_needed_text == pane_content[-1]: - pane.send_keys("yes") - -window.set_window_option("synchronize-panes", "on") -pane = window.panes[0] -pane.send_keys("sudo su -") - -for command in extra_commands: - pane.send_keys(command) - -window.set_window_option("synchronize-panes", "off") -window.select() -window.select_layout("tiled") +active_session.switch_client()