csshnator

#!/usr/bin/env python
try:
    import gtk
except ImportError:
    from gi.repository import Gtk as gtk
import sys
import math
import subprocess
import argparse
from configobj.configobj import ConfigObj
from os.path import expanduser
from shutil import copyfile

def check_version():
    # Check terminator version first
    out = subprocess.check_output(["terminator", "-v"])
    terminator_ver = float(out.split()[1])
    if (terminator_ver < 0.98):
        print "Terminator version is too low"
        print "Please install version 0.98 or greater"
        sys.exit(2)


def parse_args(argv):
    parser = argparse.ArgumentParser(
        description='Open ClusterSSH-like session on Terminator'
    )

    parser.add_argument(
        '-l',
        '--login',
        dest='login',
        type=str,
        default="",
        help='Login username to pass to all hosts used.'
    )

    parser.add_argument(
        '-s',
        '--show',
        dest='show',
        action='store_true',
        help='Show all clusters available.'
    )

    parser.add_argument(
        '-c',
        '--cluster-name',
        dest='cluster_name',
        type=str,
        default="tmp",
        help='Cluster name is a collection of hosts available on ~/.csshnatorrc file'
    )

    parser.add_argument(
        'cluster_nodes',
        nargs='*',
        default=None,
        help='Hostnames or user@hostname to connect to, separated by space'
    )

    return parser.parse_args(argv)


def main(argv):
    check_version()
    args = parse_args(argv)
    home = expanduser("~")
    cssh_config = ConfigObj(home + "/.csshnatorrc")

    if args.show:
        max_size = max([len(k) for k in cssh_config]) + 2
        formater = "{:<" + str(max_size) + "}"
        for key, value in cssh_config.items():
            print formater.format(key), value
        exit()

    cluster_name, login = args.cluster_name, args.login

    width = gtk.gdk.screen_width()
    height = gtk.gdk.screen_height()
    config_path = home + "/.config/terminator/config_csshnator"
    copyfile(home + "/.config/terminator/config", config_path)
    terminator_config = ConfigObj(config_path)

    if cluster_name == "tmp":
        if args.cluster_nodes:
            cluster_nodes = args.cluster_nodes
        else:
            return "No hosts to connect. Use --help for more info."
    else:
        cluster_nodes = cssh_config[cluster_name].split()
    nnodes = len(cluster_nodes)

    # Compute geometry
    nrows = round(math.sqrt(nnodes))
    ncols = math.ceil(nnodes/nrows)

    nrows = int(nrows)
    ncols = int(ncols)

    # TODO calculate geometry better so don"t need to create extra splits for
    # less than 3 nodes
    nrows = max(nrows, 2)  # atlease 2 rows to make geometry easier
    ncols = max(ncols, 2)  # atlease 2 columns to make geometry easier

    # Construct layout
    vpane_name = "vpane"
    hpane_name = "hpane"
    term_name = "term"
    window_name = "window"

    cssh_layout = {
        window_name + "0": {
            "type": "Window",
            "parent": "",
        }
    }

    # VPaned/ Horizontal splits
    paneparent = window_name + "0"
    for vpane in range(0, nrows - 1):
        panepos = height/nrows  # not used as far as I can tell
        paneratio = float(1.0/(nrows - vpane))  # 1/n ... 1/3,1/2
        panename = vpane_name + str(vpane)
        cssh_layout[panename] = {
                "type": "VPaned",
                "order": min(vpane, 1),  # first pane order 0, all others order 1
                "position": panepos,
                "ratio": paneratio,
                "parent": paneparent,
                }
        paneparent = panename

    # HPaned / Veritcal Split
    for row in range(0, nrows):
        # First split, parent is VPaned
        order = 0
        paneparent = vpane_name + str(row)
        if (row == (nrows - 1)):  # last row order is 1 last row parent is second to last
            order = 1
            paneparent = vpane_name + str(row - 1)
        panepos = width/ncols
        paneratio = float(1.0/(ncols))
        panename = hpane_name + str(row) + str(0)  # hpaned00 ~ hpanedn0
        cssh_layout[panename] = {
            "type": "HPaned",
            "position": panepos,
            "ratio": paneratio,
            "order": order,
            "parent": paneparent,
        }
        # Other panes parent is previous pane
        paneparent = panename
        for hpane in range(1, ncols - 1):
            panepos = width/ncols
            paneratio = float(1.0/(ncols - hpane))
            panename = hpane_name + str(row) + str(hpane)  # hpaned00 ~ hpanednn
            cssh_layout[panename] = {
                "type": "HPaned",
                "position": panepos,
                "ratio": paneratio,
                "order": 1,
                "parent": paneparent,
            }
            paneparent = panename

    # Child terminals parents are Hpanes
    node_ind = 0
    for row in range(0, nrows):
        for col in range(0, ncols):
            if (node_ind >= nnodes):
                command = "exit"  # exit the terminal if there is no node for this
            else:
                node = cluster_nodes[node_ind]  # get the node for this terminal
                node_ind += 1
                command = "echo CSSHNATOR started &&"
                command += "echo BroadCast All default At+A &&"
                command += "echo connecting to " + login + ("@" if login else "") + node + " &&"  # ssh command
                command += "ssh " + login + ("@" if login else "") + node  # ssh command
                # command = "echo " + login + "@" + node  + " && bash "#for debugging

            order = 0
            termparent = hpane_name + str(row) + str(col)  # parent is hpane
            if (col == (ncols - 1)):  # last col order 1 and col-1 parent
                order = 1
                termparent = hpane_name + str(row) + str(col - 1)

            term = term_name + str(row) + str(col)
            cssh_layout[term] = {
                "command": command,
                "profile": "default",
                "type": "Terminal",
                "order": order,
                "parent": termparent,
            }

    # Write to config and launch terminator
    configname = "cssh_config_" + cluster_name
    terminator_config["layouts"][configname] = cssh_layout
    terminator_config["global_config"]["broadcast_default"] = "all"
    terminator_config.write()
    subprocess.Popen(
        [
            "terminator",
            "-u",
            "-g", config_path,
            "-l", configname,
            "--title", "CSSHNator - {}".format(cluster_name)
        ]
    )


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