misc-fred/eoptasks/eoptasks.py

143 lines
5.0 KiB
Python
Executable File

#! /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))
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)
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))
parser = argparse.ArgumentParser()
parser.add_argument('-l', '--list-servers', action='store_true')
parser.add_argument('-k', dest='keywords', type=str)
parser.add_argument('cmd', type=str, nargs='?', default=None)
args = parser.parse_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
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)
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)
tmux_session_name = 's%s' % random.randrange(1000)
os.system('tmux new-session -s %s -d /bin/bash -c "sleep 2"' % tmux_session_name)
tmux = libtmux.Server()
session = tmux.find_where({'session_name': tmux_session_name})
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.')
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)
if len(session.list_windows()) == 1:
# a single window remains but there are still unparallelizable
# servers to process, create temporary tmux window to avoid tmux
# being destroyed as the last window process is over.
session.new_window(attach=False, window_name='sleep',
window_shell='/bin/bash -c "sleep 2"')
while len(session.list_windows()) > 5:
time.sleep(0.1)