misc-fred/eoptasks/eoptasks.py

#! /usr/bin/python3
#
# This script provides parallel remote execution of commands, while having
# some special knownledge of servers that should *not* be handled in parallel.
#
# It defers terminal-handling to tmux(1).
#
# It has some targeting capacities using keywords. Commas for 'OR' and slashes
# for 'AND', ex: ext/test,saas/test/passerelle will select all external test
# servers + all passerelle servers on the SaaS.
#
# It takes any shell command and has some builtin shortcuts such as apt.update
# and apt.upgrade.
#
# Requirements: libtmux and pyyaml.

import argparse
import configparser
import os
import random
import re
import sys
import time

import libtmux
import yaml

class Server:
    def __init__(self, servername, group=''):
        self.name = servername
        self.keywords = set(re.split(r'[-_ \.]', servername + ' ' + group))

        # add all possible hostname parts as keywords,
        # ex: node1.dev.entrouvert.org will add:
        #     node1.dev, node1.dev.entrouvert, node1.dev.entrouvert.org,
        #     dev.entrouvert, dev.entrouvert.org, entrouvert.org
        parts = servername.split('.')
        for i in range(len(parts)-1):
            for j in range(i, len(parts)):
                if i != j:
                    self.keywords.add('.'.join(parts[i:j+1]))

    def __repr__(self):
        return '<Server %s %r>' % (self.name, self.keywords)

    def shell(self):
        return 'ssh %s' % self.name

    def cmd(self, cmd):
        return 'ssh %s "%s"' % (self.name, cmd)


def get_servers():
    servers = []
    config = configparser.ConfigParser()
    config.read(os.path.join(os.path.expanduser('~/.config/eoptasks.ini')))
    servergroup = config.get('config', 'servergroups', fallback=None)
    if servergroup is None:
        print("You need to create ~/.config/eoptasks.ini with such a content:\n"
              "\n"
              "  [config]\n"
              "  servergroups = /home/user/src/puppet/data/servergroups.yaml\n")
        sys.exit(1)

    servergroups = yaml.load(open(servergroup))['servergroups']
    for group in servergroups:
        for servername in servergroups[group]:
            servers.append(Server(servername, group))
    return servers

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-l', '--list-servers', action='store_true')
    parser.add_argument('--status-window', dest='session_name', type=str)
    parser.add_argument('-k', dest='keywords', type=str)
    parser.add_argument('cmd', type=str, nargs='?', default=None)
    args = parser.parse_args()
    return args

def filter_servers(servers, args):
    selected_servers = []
    if args.keywords:
        for keyword in args.keywords.split(','):
            keywords = set(keyword.split('/'))
            selected_servers.extend([
                x for x in servers
                if keywords.issubset(x.keywords) and not x in selected_servers])
        for keyword in args.keywords.split(','):
            if keyword.startswith('!') or keyword.startswith('-'):
                selected_servers = [x for x in selected_servers if keyword[1:] not in x.keywords]
    else:
        selected_servers = servers
    return selected_servers


args = parse_args()

if args.session_name:
    print('eoptasks, tmux session %s' % args.session_name)
    time.sleep(3600)
    sys.exit(0)

servers = get_servers()
selected_servers = filter_servers(servers, args)

if args.list_servers:
    for server in sorted(selected_servers, key=lambda x: x.name):
        print(server.name)
    sys.exit(0)

if not selected_servers:
    sys.stderr.write('No matching servers\n')
    sys.exit(1)

if not args.cmd:
    sys.stderr.write('Missing command\n')
    sys.exit(1)

def init_tmux_session():
    tmux_session_name = 's%s' % random.randrange(1000)
    os.system('tmux new-session -s %s -n 🌑 -d %s --status-window %s' % (
            tmux_session_name, sys.argv[0], tmux_session_name))
    return tmux_session_name

tmux_session_name = init_tmux_session()

pid = os.fork()
if pid:
    os.system('tmux attach-session -t %s' % tmux_session_name)
else:
    def cluster_name(server_name):
        return re.match(r'(.*?)(\d*)$', server_name).group(1).replace(
                '.rbx.', '.loc.').replace('.gra.', '.loc.').replace('.sbg.', '.loc.')

    tmux = libtmux.Server()
    session = tmux.find_where({'session_name': tmux_session_name})

    cmd = {
        'apt.update': 'sudo apt update',
        'apt.upgrade': 'sudo apt update && sudo apt upgrade -y',
        'collectstatic': '''sudo -u authentic-multitenant authentic2-multitenant-manage collectstatic --noinput;
                            sudo -u bijoe bijoe-manage collectstatic --noinput;
                            sudo -u chrono chrono-manage collectstatic --noinput;
                            sudo -u combo combo-manage collectstatic --noinput;
                            sudo -u corbo corbo-manage collectstatic --noinput;
                            sudo -u fargo fargo-manage collectstatic --noinput;
                            sudo -u hobo hobo-manage collectstatic --noinput;
                            sudo -u passerelle passerelle-manage collectstatic --noinput;
                            sudo -u wcs wcs-manage collectstatic;
                            /bin/true'''.replace('\n', ''),
    }.get(args.cmd, args.cmd)

    status_window = session.attached_window

    total_number = len(selected_servers)
    random.shuffle(selected_servers)
    while selected_servers:
        current_clusters = [cluster_name(x.name) for x in session.list_windows()]
        for server in selected_servers[:]:
            if cluster_name(server.name) in current_clusters:
                continue
            selected_servers.remove(server)
            session.new_window(
                    attach=False,
                    window_name=server.name,
                    window_shell=server.cmd(cmd))
            break
        else:
            time.sleep(0.1)
        while len(session.list_windows()) > 5:
            time.sleep(0.1)
        percentage = len(selected_servers) / total_number
        if percentage == 1:
            status_window.rename_window('🌕')
        elif percentage >= 0.75:
            status_window.rename_window('🌔')
        elif percentage >= 0.5:
            status_window.rename_window('🌓')
        elif percentage >= 0.25:
            status_window.rename_window('🌒')
    while len(session.list_windows()) > 1:
        time.sleep(0.1)
    try:
        status_window.kill_window()
    except libtmux.exc.LibTmuxException:
        # killing the last window will end the tmux session, and it's not
        # handled correctly by kill_window as it ends by updating the list
        # of windows (and hitting the dead tmux).
        pass