csshi

#!/usr/bin/env python3
""" Launch multiple ssh sessions in a single iTerm2 window and broadcast keyboard input

Copyright (C) 2021-2024 David Stark

This program is free software; you can redistribute it and/or modify it under the terms of
the GNU General Public License as published by the Free Software Foundation; version 2 only.

This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with this program;
if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""

import argparse
import asyncio
import os
import subprocess
import sys
from math import ceil, floor, sqrt
from os.path import basename
from typing import NoReturn, Optional

import iterm2
from AppKit import NSWorkspace
from iterm2.profile import TitleComponents

# We prepend U+2064 INVISIBLE PLUS to session names
# so we can find sessions we have created later on
session_name_prefix = "\u2064"


def error(message: str) -> NoReturn:
    """Squeak and keel over"""
    sys.stderr.write("ERROR: " + message + "\n")
    sys.exit(1)


def debug(message: str) -> None:
    """Tell me more"""
    if debugging:
        sys.stderr.write("DEBUG: " + message + "\n")


def get_arg_parser() -> argparse.ArgumentParser:
    """Parse it. Parse it real good."""
    parser = argparse.ArgumentParser(
        description="Run multiple ssh connections concurrently in an iTerm2 terminal."
    )
    parser.add_argument(
        "-d", "--debug", dest="debug", action="store_true", help="Turn on debugging"
    )
    parser.add_argument(
        "-k",
        "--kill-inactive",
        dest="kill",
        action="store_true",
        help="Don't create spacer panes and kill inactive panes",
    )
    parser.add_argument(
        "-c",
        "--caffeinate",
        dest="caffeine",
        action="store_true",
        help="Try to keep the Mac awake",
    )
    parser.add_argument(
        "-n",
        "--no-broadcast",
        dest="nobroadcast",
        action="store_true",
        help="Don't enable input broadcast",
    )
    parser.add_argument(
        "-l",
        "--username",
        dest="username",
        help="Login username, overridden per-host by specifying 'user@host'",
    )
    parser.add_argument(
        "-p",
        "--port",
        dest="port",
        type=int,
        help="ssh server port, overridden per-host by specifying 'host:port",
    )
    parser.add_argument(
        "-J",
        "--jump",
        dest="jumphost",
        help="ProxyJump host specification, like [user@]jumphost[:port]",
    )
    parser.add_argument(
        "-s",
        "--sleep",
        dest="sleep",
        default=0,
        type=float,
        help="Sleep time in seconds between hosts. Defaults to 0",
    )
    parser.add_argument(
        "-o",
        "--options",
        action="append",
        dest="options",
        help="Raw ssh options string, e.g. '-i ~/.ssh/id_ansible'. Can be specified multiple times.",
    )
    parser.add_argument(
        "-b",
        "--binary",
        dest="binary",
        default="ssh",
        help="SSH binary to use. Defaults to 'ssh'",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "-C",
        "--columns",
        dest="columns",
        type=int,
        help="Maximum number of columns to use.",
    )
    group.add_argument(
        "-R", "--rows", dest="rows", type=int, help="Maximum number of rows to use."
    )
    parser.add_argument(
        "destinations",
        metavar="destination",
        nargs="*",
        help="[user@]host[:port] ssh server specification. Use square brackets around IPv6 IP adresses",
    )
    parser.add_argument(
        "-L",
        "--list-sessions",
        dest="list_sessions",
        action="store_true",
        help="Print space-separated list of csshi sessions, one line per window",
    )
    return parser


def args_to_hostnames(args: argparse.Namespace) -> list[str]:
    """Create a list of hostnames from the user args."""
    hostnames = []
    for host_spec in args.destinations:
        spec = host_spec
        hostname = None

        # Strip down spec until it's just a hostname
        if "@" in spec:
            spec = spec.split("@")[1]
        if "]" in spec:
            # IPv6 [2100:1111::1111] square bracket syntax
            hostname = spec.split("]")[0].lstrip("[")
            spec = spec.split("]")[1]
        if ":" in spec:
            spec = spec.split(":")[0]
        if spec:
            hostname = spec

        if not hostname:
            error(f"ssh host '{host_spec}' does not appear to include a hostname.")

        hostnames.append(hostname)

    return hostnames


def args_to_ssh_strings(args: argparse.Namespace) -> list[str]:
    """Create a list of ssh command strings from the user args.
    ssh destination specs like user@ and :port override the '-l' and '-p' args.
    """
    ssh_strings = []
    for host_spec in args.destinations:
        spec = host_spec
        ssh_string = ""
        user = None
        hostname = None
        port = None
        if args.username:
            debug(f"Setting default user {args.username}")
            user = args.username
        if args.port:
            debug(f"Setting default port {args.port}")
            port = args.port

        # Strip down spec until it's just a hostname
        if "@" in spec:
            user = spec.split("@")[0]
            spec = spec.split("@")[1]
        if "]" in spec:
            # IPv6 [2100:1111::1111] square bracket syntax
            hostname = spec.split("]")[0].lstrip("[")
            spec = spec.split("]")[1]
        if ":" in spec:
            port = spec.split(":")[1]
            spec = spec.split(":")[0]
        if spec:
            hostname = spec

        if not hostname:
            error(f"ssh host '{host_spec}' does not appear to include a hostname.")

        ssh_string = hostname
        if user:
            debug(f"User for host {hostname} is {user}")
            ssh_string = " ".join(["-l", user, ssh_string])
        if port:
            debug(f"Port for host {hostname} is {port}")
            ssh_string = " ".join(["-p", str(port), ssh_string])
        if args.jumphost:
            debug(f"Setting jumphost for host {hostname} to {args.jumphost}")
            ssh_string = " ".join(["-J", args.jumphost, ssh_string])
        if args.options:
            for option_count in range(len(args.options)):
                debug(
                    f'Setting option for host {hostname}: "{args.options[option_count]}"'
                )
                ssh_string = " ".join(["-o", args.options[option_count], ssh_string])
        ssh_string = " ".join([args.binary, ssh_string])
        if args.caffeine:
            ssh_string = " ".join(["caffeinate -dimsu", ssh_string])

        debug(f"ssh_string for host {hostname} is: {ssh_string}")
        ssh_strings.append(ssh_string)

    return ssh_strings


def get_user_shell() -> Optional[str]:
    """Find the user's shell, or whichever bash is first in PATH"""
    if os.environ.get("SHELL"):
        shell = os.environ.get("SHELL")
    else:
        try:
            pcomplete = subprocess.run(
                "command -v bash", shell=True, check=True, capture_output=True
            )
        except subprocess.CalledProcessError:
            error("Could not find a local shell to run. Set $SHELL or install bash.")
        shell = str(pcomplete.stdout.decode("utf-8").strip())
    debug(f"Using shell {shell}")
    return shell


def args_to_columns(args: argparse.Namespace) -> int:
    """Set a calculated number of columns, or user specified
    whichever is less
    """
    floor_square = floor(sqrt(len(args.destinations)))
    if args.columns is not None and args.columns < 1:
        error("Columns cannot be < 1")
    if args.rows is not None and args.rows < 1:
        error("Rows cannot be < 1")
    if args.columns:
        if args.columns <= floor_square:
            debug(f"Using columns argument: {args.columns}")
            return args.columns
        debug("Not using columns argument - sqrt is lower")
    if args.rows:
        columns = ceil(len(args.destinations) / args.rows)
        debug(
            f"Using {columns} columns from {len(args.destinations)} destinations, {args.rows} rows"
        )
        return columns
    return floor_square


def gen_ssh_profile() -> iterm2.LocalWriteOnlyProfile:
    """iTerm profile for ssh panes
    The ssh command is set in the main loop with set_command()
    The session title is set in the main loop with set_name()
    """
    profile = iterm2.LocalWriteOnlyProfile()
    profile.set_title_components([TitleComponents.SESSION_NAME])
    profile.set_allow_title_setting(False)
    profile.set_use_custom_command("Yes")
    profile.set_close_sessions_on_end(False)
    return profile


def gen_spacer_profile() -> iterm2.LocalWriteOnlyProfile:
    """iTerm profile for spacer panes
    Just spit out a message a wait
    """
    profile = iterm2.LocalWriteOnlyProfile()
    profile.set_use_custom_command("Yes")
    profile.set_command(f'{shell} -c "/bin/echo PANE NOT IN USE"')
    profile.set_close_sessions_on_end(False)
    return profile


async def sessions_list(connection: iterm2.Connection) -> None:
    debug("Connecting to iTerm")
    app = await iterm2.async_get_app(connection)
    debug("Connected")
    windows = {}  # type: dict[str, list[iterm2.session.Session]]
    for window in app.windows:
        windows[window.window_id] = []
        for tab in window.tabs:
            for session in tab.sessions:
                if session.name.startswith(session_name_prefix):
                    windows[window.window_id].append(
                        session.name[len(session_name_prefix) :]
                    )
    for session_list in windows.values():
        if session_list:
            print(" ".join(session_list))


async def sessions_run(connection: iterm2.Connection) -> None:
    """Do the thing.
    This is getting a bit long, but it's all procedural
    stuff, so just sit back and enjoy the ride.
    """
    # Fire up iTerm if it's not already running
    debug("Launching iTerm...")
    NSWorkspace.sharedWorkspace().launchApplication_("iTerm")
    debug("iTerm launched")
    ssh_exec_string = f"{shell} -l -c '{{ssh_string}}'"
    debug(f"Command string is {ssh_exec_string}")
    ssh_profile = gen_ssh_profile()
    spacer_profile = gen_spacer_profile()

    debug("Connecting to iTerm")
    app = await iterm2.async_get_app(connection)
    debug("Connected")

    host_ind = 0

    async with iterm2.SessionTerminationMonitor(connection) as mon:
        # Launch the first pane straight into ssh
        debug("Launching 1st pane")
        ssh_profile.set_command(
            ssh_exec_string.format(ssh_string=ssh_strings[host_ind])
        )
        ssh_profile.set_name(f"{session_name_prefix}{hostnames[host_ind]}")
        myterm = await iterm2.Window.async_create(
            connection, profile_customizations=ssh_profile
        )

        tab = myterm.current_tab

        # Launch the rest of the first row into ssh
        # First row will always be all ssh panes, no spacers
        column_heads = [tab.current_session]
        for _ in range(cols - 1):
            debug(f"Sleeping for {sleep} seconds")
            await asyncio.sleep(sleep)
            host_ind += 1
            ssh_profile.set_command(
                ssh_exec_string.format(ssh_string=ssh_strings[host_ind])
            )
            ssh_profile.set_name(f"{session_name_prefix}{hostnames[host_ind]}")
            debug("Splitting pane for ssh")
            column_heads.append(
                await tab.current_session.async_split_pane(
                    vertical=True, profile_customizations=ssh_profile
                )
            )

        # Create the rest of the pane rows, either as ssh or spacer
        total_sessions = cols
        column_tails = column_heads.copy()  # new pane rows are split from the bottom
        for _ in range(rows - 1):
            for j in range(cols):
                total_sessions += 1
                if total_sessions <= len(ssh_strings):
                    # Still launching ssh panes
                    host_ind += 1
                    ssh_profile.set_command(
                        ssh_exec_string.format(ssh_string=ssh_strings[host_ind])
                    )
                    ssh_profile.set_name(f"{session_name_prefix}{hostnames[host_ind]}")
                    debug(f"Sleeping for {sleep} seconds")
                    await asyncio.sleep(sleep)
                    debug("Splitting pane for ssh")
                    last = await column_tails[j].async_split_pane(
                        vertical=False, profile_customizations=ssh_profile
                    )
                elif not kill_defunct:
                    # We've launched all of the ssh panes
                    # If '-k' flag is _not_ specified, launch spacer panes to
                    # fill the last row and keep the pane sizes consistent
                    debug("Splitting pane for spacer")
                    last = await column_tails[j].async_split_pane(
                        vertical=False, profile_customizations=spacer_profile
                    )
                column_tails[j] = last

        if broadcast:
            # Attach user input to the ssh panes
            domain = iterm2.broadcast.BroadcastDomain()
            for session in tab.sessions:
                domain.add_session(session)
            bc_domains = app.broadcast_domains
            bc_domains.append(domain)
            debug("Enabling input broadcast")
            await iterm2.async_set_broadcast_domains(connection, bc_domains)

        # Tidy up the window
        debug("Arranging panes")
        await iterm2.MainMenu.async_select_menu_item(
            connection, "Arrange Split Panes Evenly"
        )
        debug("Activating 1st pane")
        await tab.sessions[0].async_activate()

        # Wait for session end if '-k' flag supplied
        while kill_defunct and tab.sessions:
            debug(f"Active sessions: {tab.sessions}")
            debug("Waiting for next session to close...")
            session_id = await mon.async_get()
            debug(f"Got termination for {session_id}")
            session = app.get_session_by_id(session_id)
            if session:
                if session in tab.sessions:
                    debug(f"Closing session {session_id}")
                    await session.async_close()
                    if tab.sessions == [session]:
                        debug(f"Session for {session_id} was last session")
                        return
                else:
                    debug(f"Ignoring session {session_id} - not in our tab")
            else:
                debug(f"Session for {session_id} gone")


if __name__ == "__main__":
    parser = get_arg_parser()
    args = parser.parse_args()
    debugging = args.debug

    if args.list_sessions:
        debug("Entering sessions_list()")
        iterm2.run_until_complete(sessions_list)
        debug("sessions_list() has completed. Exiting.")
    else:
        if not args.destinations:
            parser.print_usage()
            sys.stderr.write(
                f"{basename(sys.argv[0])}: error: the following arguments are required: destination\n"
            )
            sys.exit(2)
        ssh_strings = args_to_ssh_strings(args)
        hostnames = args_to_hostnames(args)
        shell = get_user_shell()
        cols = args_to_columns(args)
        rows = ceil(len(ssh_strings) / cols)
        kill_defunct = args.kill
        sleep = args.sleep
        broadcast = not args.nobroadcast
        debug("Entering sessions_run()")
        iterm2.run_until_complete(sessions_run)
        debug("sessions_run() has completed. Exiting.")