Update tmux-ssh-groups

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,40 +27,52 @@ 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]
+# Get active tmux sessions
-
-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)
-
 srv = libtmux.Server()
-active_sessions = srv.sessions.filter(session_attached='1')
+active_sessions = srv.attached_sessions
 
-if active_sessions:
+if parser_args.separate_sessions:
-    active_session = active_sessions[0]
+    sessions = srv.sessions.filter(name="SSH-MultiG")
+
+    if sessions:
+        active_session = sessions[0]
+
+    else:
+        active_session = srv.new_session(session_name="SSH-MultiG")
 
 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")
+
+        if not session_choice:
+            sys.exit(0)
+
+        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]
@@ -57,35 +80,63 @@ else:
         else:
             active_session = srv.sessions[0]
 
-window_name = f"ssh-multig {selected_group.replace(' ', '_')}"
+for selected_group in selections:
-if windows := active_session.windows.filter(name=window_name):
+
+    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()
-    sys.exit(0)
+        continue
 
-window = active_session.new_window(window_name, window_shell=f"ssh {servers[0]}")
+    window = active_session.new_window(window_name, window_shell=f"ssh {servers[0]}")
 
-for server in servers[1:]:
+    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
+    # Wait until tmux finished working
-time.sleep(0.1)
+    time.sleep(0.05)
 
-# Confirm connection on asking panes
+    # Confirm connection on asking panes
-confirmation_needed_text = "Are you sure you want to continue connecting (yes/no/[fingerprint])?"
+    confirmation_needed_text = "Are you sure you want to continue connecting (yes/no/[fingerprint])?"
-for pane in window.panes:
+    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")
+    window.set_window_option("synchronize-panes", "on")
-pane = window.panes[0]
+    pane = window.panes[0]
-pane.send_keys("sudo su -")
+    pane.send_keys("sudo su -")
 
-for command in extra_commands:
+    for command in extra_commands:
         pane.send_keys(command)
 
-window.set_window_option("synchronize-panes", "off")
+    window.set_window_option("synchronize-panes", "off")
-window.select()
+    window.select()
-window.select_layout("tiled")
+    window.select_layout("tiled")
+
+active_session.switch_client()