recover a lost Juju agent

recover-lost-agent.py

#!/usr/bin/env python
"""
This is a tool for recovering a lost juju unit

Usage:
{0} controller-ip unit-from unit-to basedir
"""

import subprocess
import shlex
import sys
import os
import logging
import re
import shutil
import tempfile
import yaml

DEFAULT_PASSWORD = "thisisasupersecurepasswordforthelostagents"
DEFAULT_PASSWORD_ENCODED = "9v8BrI6odD9agvYw2UKU+Ziya"

DIRS_TO_RENAME = [
    "{basedir}/agents/{unit_from} {basedir}/agents/{unit_to}",
    "{basedir}/tools/{unit_from} {basedir}/tools/{unit_to}",
    "{basedir}/init/jujud-{unit_from} {basedir}/init/jujud-{unit_to}",
    "{basedir}/init/jujud-{unit_to}/jujud-{unit_from}.service {basedir}/init/jujud-{unit_to}/jujud-{unit_to}.service"
]

SED_FILES = [
    "{basedir}/init/jujud-{unit_to}/jujud-{unit_to}.service",
    "{basedir}/init/jujud-{unit_to}/exec-start.sh",
    "{basedir}/agents/{unit_to}/agent.conf",
]


def sed_inplace(filename, pattern, repl):
    if not os.path.exists(filename):
        return

    print("Replacing file:%s pattern:%s replacement:%s" %
          (filename, pattern, repl))

    pattern_compiled = re.compile(pattern)
    with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp_file:
        with open(filename) as src_file:
            for line in src_file:
                tmp_file.write(pattern_compiled.sub(repl, line))

    shutil.copystat(filename, tmp_file.name)
    shutil.move(tmp_file.name, filename)


def update_mongodb(controller, unit, password=DEFAULT_PASSWORD_ENCODED):
    TEMPLATE = """use juju
db.units.update({'name': "%s"}, { $set: { 'passwordhash': "%s"}})
""" % (unit, password)

    with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp_file:
        tmp_file.write(TEMPLATE)

    run('juju scp {0} {1}:~'.format(tmp_file.name, controller), output=False)
    run("juju ssh %s \"sudo /usr/lib/juju/mongo3.2/bin/mongo --sslAllowInvalidCertificates --ssl -u admin -p $(grep oldpassword /var/lib/juju/agents/machine-0/agent.conf | awk -e '{print $2}') localhost:37017/admin < /home/ubuntu/%s\"" % (controller, os.path.basename(tmp_file.name)))


def run(cmd, output=True):
    cmd = shlex.split(cmd)
    print(cmd)
    if output:
        return subprocess.check_output(cmd)
    return subprocess.call(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE)


def _format_unit(unit, prefix="unit"):
    unit = unit.replace("/", "-")
    if prefix:
        return "{0}-{1}".format(prefix, unit)
    return unit


def rename_dirs(destdir, unit_from, unit_to):
    for directory in DIRS_TO_RENAME:
        run("mv {0}".format(
            directory.format(
                basedir=destdir,
                unit_from=unit_from,
                unit_to=unit_to,
            )), output=False)


def sed_files(destdir, unit_from, unit_to, unit_from_fmt, unit_to_fmt):
    for filename in SED_FILES:
        sed_inplace(
            filename.format(
                basedir=destdir,
                unit_to=unit_to_fmt,
                unit_from=unit_from_fmt,
            ),
            unit_from_fmt, unit_to_fmt)

        sed_inplace(
            filename.format(
                basedir=destdir,
                unit_to=unit_to_fmt,
                unit_from=unit_from_fmt,
            ),
            unit_from, unit_to)


def modify_agentconf(destdir, unit_to_fmt):
    agent_conf = os.path.join(destdir, "agents", unit_to_fmt, "agent.conf")
    sed_inplace(agent_conf, "^apipassword\:.*",
                "apipassword: {0}".format(DEFAULT_PASSWORD))
    sed_inplace(agent_conf, "^statepassword\:.*",
                "statepassword: {0}".format(DEFAULT_PASSWORD))


def main():
    unit_from, unit_to = sys.argv[2:4]
    controller = sys.argv[1]

    unit_to_fmt, unit_from_fmt = _format_unit(unit_to), _format_unit(unit_from)

    basedir = os.path.abspath(sys.argv[4])
    destdir = os.path.join(basedir, unit_to_fmt)

    if os.path.exists(destdir):
        shutil.rmtree(destdir)

    run('juju ssh {0} "sudo chmod -R a+r /var/lib/juju"'.format(unit_from))
    run('juju scp -- -r {0}:/var/lib/juju {1}'.format(unit_from, destdir),
        output=False)

    rename_dirs(destdir, unit_from_fmt, unit_to_fmt)
    sed_files(destdir, unit_from, unit_to, unit_from_fmt, unit_to_fmt)
    modify_agentconf(destdir, unit_to_fmt)
    update_mongodb(controller, unit_to)

    run('juju ssh {0} "sudo chown -R ubuntu:ubuntu /var/lib/juju"'.format(
       unit_to))
    run('juju scp -- -r {0}/ {1}:~'.format(destdir, unit_to), output=False)
    run('juju ssh {0} "sudo cp -rp ~/{1}/* /var/lib/juju"'.format(
        unit_to, unit_to_fmt), output=False)
    run('juju ssh {0} "sudo systemctl link \'/var/lib/juju/init/jujud-{unit}/jujud-{unit}.service\'"'.format(
       unit_to, unit=unit_to_fmt),
       output=False)
    run('juju ssh {0} "sudo systemctl daemon-reload"'.format(unit_to))
    run('juju ssh {0} "sudo systemctl restart jujud-{unit}.service'.format(unit_to,
                                                                           unit=unit_to_fmt))

if __name__ == "__main__":
    if len(sys.argv) < 5:
        print(__doc__.format(sys.argv[0]))
        sys.exit(-1)

    main()