agent-bootstrap-script.template

# NOTE: don't use .format for string substitution since templating uses same to do
# variable substitution


import subprocess
import sys
import httplib
import urllib2
import shlex
import os
import platform
import socket
import ssl
import tempfile
import json
import traceback
import signal
import logging

# Set the time, so certs work below ALWAYS.
# Future NTP job will set it again,
# so preserving host's date is unecessary
ret = subprocess.call("date -s @{server_epoch_seconds} > /dev/null 2>&1", shell=True)

# create Logger with a single file handler
bootstrap_log = logging.getLogger("bootstrap")
bootstrap_log.setLevel(logging.DEBUG)

handler = logging.FileHandler("/var/log/chroma-agent.log")
handler.setLevel(logging.DEBUG)
handler.setFormatter(logging.Formatter("[%(asctime)s] bootstrap %(name)s %(levelname)s %(message)s"))
bootstrap_log.addHandler(handler)

IML_CONF_DIR = "/etc/iml/"

# The agent's private key, never shared
PRIVATE_KEY = os.path.join(IML_CONF_DIR, "private.pem")

# The agent's certificate, generated by the manager in response to a CSR
AGENT_CERT = os.path.join(IML_CONF_DIR, "self.crt")

# The root certificate used to authenticate the manager
AUTHORITY_CERT = os.path.join(IML_CONF_DIR, "authority.crt")

REPO_PATH = "/etc/yum.repos.d/Intel-Lustre-Agent.repo"

base_url = "{base_url}"
reg_url = "{reg_url}"
cert_str = """{cert_str}"""
repo_url = "{repo_url}"
repo_packages = "{repo_packages}"
profile_json = '{profile_json}'

REPO_CONTENT = """{repos}"""


class VerifiedHTTPSConnection(httplib.HTTPSConnection):
    def connect(self):
        sock = socket.create_connection((self.host, self.port), self.timeout)
        if self._tunnel_host:
            self.sock = sock
            self._tunnel()
        self.sock = ssl.wrap_socket(sock, cert_reqs=ssl.CERT_REQUIRED, ca_certs=AUTHORITY_CERT)


class VerifiedHTTPSHandler(urllib2.HTTPSHandler):
    def __init__(self, connection_class=VerifiedHTTPSConnection):
        self.specialized_conn_class = connection_class
        urllib2.HTTPSHandler.__init__(self)

    def https_open(self, req):
        return self.do_open(self.specialized_conn_class, req)


def launch_command(cmd):
    args = shlex.split(cmd)
    debug = 'Command returned %d: stdout: "%s" stderr: "%s"'

    bootstrap_log.info("Executing shell command with arguments: %s" % args)

    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = p.communicate()

    bootstrap_log.debug(debug % (p.returncode, out, err))

    if p.returncode != 0:
        raise Exception(debug % (p.returncode, out, err))

    return out.strip()


def setup_keys():
    if not os.path.exists(IML_CONF_DIR):
        os.makedirs(IML_CONF_DIR)

    with open(AUTHORITY_CERT, "w") as f:
        f.write(cert_str)

    return launch_command("openssl genrsa -out %s 2048" % PRIVATE_KEY)


def reg_manager():
    fqdn = socket.getfqdn()
    nodename = os.uname()[1]
    csr = launch_command("openssl req -new -subj /C=/ST=/L=/O=/CN=%s -key %s" % (fqdn, PRIVATE_KEY))

    data = json.dumps({{
        "address": socket.gethostbyname(socket.gethostname()),
        "fqdn": fqdn,
        "nodename": nodename,
        "version": 0,
        "csr": csr,
        "capabilities": [],
    }})

    try:
        https_handler = VerifiedHTTPSHandler()
        url_opener = urllib2.build_opener(https_handler)
        url_opener.addheaders.append(("Content-Type", "application/json"))
        handle = url_opener.open(reg_url, data)
        response = handle.readlines()
        handle.close()

        json_resp = json.loads(response[0])
        open(AGENT_CERT, "w").write(json_resp["certificate"])
        return json_resp
    except urllib2.URLError, e:
        bootstrap_log.exception("Cannot reach host/url %s: %s" % (reg_url, e))
        raise


def create_repo(contents):
    tmp = tempfile.NamedTemporaryFile()
    # this format needs to match chroma_agent.action_plugins.agent_setup.configure_repo
    tmp.write(contents)

    tmp.flush()
    launch_command("cp %s %s" % (tmp.name, REPO_PATH))
    os.chmod(REPO_PATH, 0644)


def install_agent():
    return launch_command("yum install -y %s" % repo_packages)


def configure_server():
    launch_command("chroma-agent set_server_url --url %s" % base_url)
    launch_command("chroma-agent set_profile --profile_json '%s'" % profile_json)


def start_services():
    launch_command("systemctl enable --now iml-storage-server.target")


def kill_zombies():
    try:
        # attempt to stop a pre-installed instance
        launch_command("systemctl stop iml-storage-server.target")
    except Exception:
        # Will fail if not pre-installed
        pass

    # ensure that there are no old agents hanging around
    for pid in [d for d in os.listdir("/proc") if d.isdigit()]:
        try:
            with open("/proc/%s/cmdline" % pid) as f:
                if "chroma-agent-daemon" in f.read():
                    os.kill(int(pid), signal.SIGKILL)
        except (OSError, IOError):
            # Don't bail if this fails -- it's not critical
            pass


def main():
    try:
        bootstrap_log.info("Agent bootstrap started...")

        # Get a clean deployment environment
        kill_zombies()

        # Set up SSL keys and register with the manager using our
        # embedded registration token
        setup_keys()
        registration_response = reg_manager()

        contents = REPO_CONTENT.format(AUTHORITY_CERT, PRIVATE_KEY, AGENT_CERT).strip()

        if len(contents) > 0:
            # Now that we're registered, we can download packages
            create_repo(contents)
            install_agent()

        # Persist the agent configuration file so that it can connect
        # to the manager
        configure_server()

        bootstrap_log.info("Agent bootstrap completed, starting agent...")

        # Finally start the services.
        start_services()

        print json.dumps(
            {{"host_id": registration_response["host_id"], "command_id": registration_response["command_id"]}}
        )

        return 0

    except Exception, err:
        bootstrap_log.exception("Error from agent-bootstrap-script main():")
        return 1


if __name__ == "__main__":
    sys.exit(main())