orchestrate.py

#! /usr/bin/python3

import argparse
import importlib
import json
import os
import subprocess
import sys
import time

DEFAULT_ACTIONS = ['files', 'ping', 'init', 'wait', 'deployment']


def perform_action(action, commands):
    if component:
        plugin_commands = component.commands(action)
        if plugin_commands:
            commands += plugin_commands
    for cmd in commands:
        print('-> %s' % ' '.join(cmd))
        if not args.dry_run:
            subprocess.run(cmd)
    if component:
        component.action(action)


parser = argparse.ArgumentParser(description='terraform cluster orchestrator')
parser.add_argument('-c', '--config',
                    help='Terraform output for using by orchestrator (default: config.json)',
                    default='config.json')
parser.add_argument('actions', metavar='ACTIONS', nargs='*',
                    help='actions (default: %s)' % ' '.join(DEFAULT_ACTIONS),
                    default=DEFAULT_ACTIONS)
parser.add_argument('-n', '--dry-run', action='store_true',
                    help='simulated run')
parser.add_argument('-p', '--parameters',
                    help='orchestration parameters')
args = parser.parse_args()

j = None
if args.config == '-':
    j = json.load(sys.stdin)
else:
    with open(args.config) as f:
        j = json.load(f)

config = j['config']['value']
n = int(config['n'])
d = config['domain']
hosts = j['hosts']['value']
public_hosts = j['public_hosts']['value']
master_hostname = config['master_hostname']
master_ip = public_hosts[master_hostname]
user = config['image_user']
secrets = config['secrets']
t = config.get('type', None)

print('== plugin ==')
if t and os.path.exists('deployments/%s/plugin.py' % t):
    print('-> deployments/%s/plugin.py' % t)
    plugin = importlib.import_module('deployments.%s.plugin' % t)
    Component = getattr(plugin, 'Component')
    component = Component(args, config, hosts, public_hosts)
else:
    component = None

if 'files' in args.actions:
    print('== files ==')
    print('-> hosts')
    if not args.dry_run:
        with open('hosts', 'w') as f:
            f.write('''\
127.0.0.1	localhost

# The following lines are desirable for IPv6 capable hosts
::1     localhost ip6-localhost ip6-loopback
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters

''')
            for h, ip in hosts.items():
                f.write('%s	%s.%s	%s.%s.	%s\n' % (ip, h, d, h, d, h))

    print('-> public_hosts')
    if not args.dry_run:
        with open('public_hosts', 'w') as f:
            for h, ip in public_hosts.items():
                f.write('%s	%s.%s	%s.%s.	%s\n' % (ip, h, d, h, d, h))

    print('-> inventory')
    if not args.dry_run:
        with open('inventory', 'w') as f:
            f.write('''\
[masters]
%s ansible_user=%s ansible_become=true ansible_host=%s

[nodes]
''' % (master_hostname, user, master_ip))
            for h, ip in hosts.items():
                if h == master_hostname:
                    continue
                f.write('\
%s ansible_user=%s ansible_become=true ansible_host=%s ansible_ssh_common_args=\'\
-o ForwardAgent=yes -o ProxyCommand="ssh -W %s:22 -q %s@%s"\'\n' % (h, user, ip, ip, user,
                                                                    master_ip))
    if component:
        component.action('files')

if 'ping' in args.actions:
    print('== ping ==')
    if not args.dry_run:
        print('-> (ping) ', end='')
        ret = None
        while ret is None or ret != 0:
            print('.', end='')
            ret = subprocess.call(['ping', '-c', '1', '-i', '2', master_ip])
        print('')
        cmd = ['ssh-keygen', '-R', master_ip]
        print('-> %s' % ' '.join(cmd))
        subprocess.run(cmd)
        cmd = ['ssh', '%s@%s' % (user, master_ip), '-o', 'ConnectTimeout=5', '-o',
               'StrictHostKeyChecking=no', ':']
        ret = None
        while ret is None or ret != 0:
            print('-> %s' % ' '.join(cmd))
            ret = subprocess.call(cmd)
            time.sleep(2)
        print('')
        print('-> sleep %d' % (10 + 5 * n))
        time.sleep(10 + 5 * n)
    if component:
        component.action('ping')

if 'init' in args.actions:
    print('== init ==')

    h = list(hosts.keys()) + ['%s.%s.' % (host, d) for host in hosts.keys()] + \
        ['%s.%s' % (host, d) for host in hosts.keys()] + list(hosts.values()) + [master_ip]

    cycle_all = \
        'echo -n At:; uname -n; for h in %s; do ssh -o StrictHostKeyChecking=no $h :; done' \
        % (' '.join(h))
    commands = [
        ['ssh-keygen', '-R', master_ip],
        ['ssh', '-o', 'StrictHostKeyChecking=no', '%s@%s' % (user, master_ip), ':'],
        ['scp', '-p', 'hosts', '%s@%s:/tmp/' % (user, master_ip)],
        ['ssh', '%s@%s' % (user, master_ip), 'sudo mv -v /tmp/hosts /etc/hosts'],
        # the frontend to all nodes
        ['ssh', '-o', 'ForwardAgent=yes', '%s@%s' % (user, master_ip), cycle_all]
    ]
    for ip in hosts.values():
        # check the special case if it is not already frontend ('ssh-keygen -R' would be a problem)
        if ip == master_ip:
            continue
        commands += [
            ['ssh-keygen', '-R', ip],
            ['ssh', '-o', 'ForwardAgent=yes', '-o', 'ProxyCommand=ssh -W %s:22 -q %s@%s' %
             (ip, user, master_ip), '-o', 'StrictHostKeyChecking=no', '%s@%s' % (user, ip), ':']
        ]
    commands += [
        ['ansible', '-i', './inventory', '-m', 'copy', '-a', 'src=hosts dest=/etc/hosts', 'nodes'],
    ]
    # node to node ssh keys
    # (after the /etc/hosts is distributed)
    for node in hosts.keys():
        # skip the frontend - already covered
        if node == master_hostname:
            continue
        # the node to all nodes
        commands += [
            ['ssh', '-o', 'ForwardAgent=yes', '%s@%s' % (user, master_ip),
             'ssh -o ForwardAgent=yes %s \'%s\'' % ('%s.%s' % (node, d), cycle_all)],
        ]
    perform_action('init', commands)

if 'wait' in args.actions:
    print('== wait ==')
    commands = [
        ['ansible', '-i', './inventory', '-m', 'command', '-a', 'uname -a', 'all'],
        ['ansible', '-i', './inventory', '-m', 'shell', '-a',
         'while ! test -f /var/lib/cloud/instance/boot-finished; do sleep 2; done', 'all'],
    ]
    perform_action('wait', commands)

if 'deployment' in args.actions:
    print('== deployment ==')
    commands = []
    perform_action('deployment', commands)