ansiblew

#!/usr/bin/env python

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import with_statement

import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib2
from contextlib import closing

DEFAULT_ANSIBLE_VERSION = "2.7.9"
DEFAULT_VIRTUALENV_VERSION = "16.4.3"

VERBOSE_LOGGING = os.getenv("ANSIBLEW_VERBOSE") in ("1", "true")
ANSIBLE_COMMAND = os.getenv("ANSIBLEW_COMMAND", "ansible-playbook")
ANSIBLEW_CONFIG_FILE = os.getenv("ANSIBLEW_CONFIG_FILE", "ansiblew.json")


def die(*args):
    """Prints the given messages and exits the application with an error."""
    print(*args, file=sys.stderr)
    sys.exit(1)


def info(*args):
    print(*args)


def debug(*args):
    if VERBOSE_LOGGING:
        print(*args)


def load_config():
    filename = os.path.join(os.path.dirname(sys.argv[0]), ANSIBLEW_CONFIG_FILE)
    if not os.path.exists(filename):
        die("Could not find config file '" + ANSIBLEW_CONFIG_FILE + "'")

    try:
        with open(filename, "rb") as fp:
            content = fp.read()

        config = json.loads(content)

        # create a canonicalized hash
        h = hashlib.md5()
        h.update("xy")
        h.update(json.dumps(config, sort_keys=True))

        return config, h.hexdigest()[:8]

    except OSError:
        die("Could not load config file 'ansiblew.json'")


def find_workspace():
    workspace = os.path.expanduser("~/.ansiblew")
    if "~" in workspace:
        workspace = "/tmp/ansiblew"

    return workspace


def mkdirs(path):
    if not os.path.exists(path):
        debug("Make directory at", path)
        os.makedirs(path)


def download(url):
    info("Download", url)
    with closing(urllib2.urlopen(url)) as fp:
        return fp.read()


def check_call(*args, **kwargs):
    fail_on_error = kwargs.pop("fail_on_error", True)
    capture_output = kwargs.pop("capture_output", False)

    # check if we need to send input to the process
    stdin = kwargs.pop("stdin", None)
    if stdin is not None:
        kwargs["stdin"] = subprocess.PIPE

    if capture_output:
        kwargs["stdout"] = subprocess.PIPE
        kwargs["stderr"] = subprocess.STDOUT

    if capture_output and stdin:
        raise ValueError("stdin can not be used with capture_output")

    debug("Executing command", " ".join(args))
    proc = subprocess.Popen(args, **kwargs)

    output = ""
    if capture_output:
        lines = []
        for line in proc.stdout:
            lines.append(line)
            sys.stdout.write(line)

        output = "".join(lines)

    else:
        # wait for the process to finish
        proc.communicate(stdin)

    # check if the process finished with an error
    retcode = proc.poll()
    if retcode and fail_on_error:
        die("Command %s failed with return code %d" % (args[0], retcode))

    return retcode, output


def atomically(func):
    def wrapper(target, *args, **kwargs):
        if os.path.exists(target):
            debug("Previous target found in", target)
            return

        temp_target = target + ".%d" % (int(1000 * time.time()) % 10000)
        try:
            mkdirs(temp_target)

            # call real code
            func(temp_target, *args, **kwargs)

            try:
                # atomic rename, might fail if someone else created the new directory concurrently
                os.rename(temp_target, target)
            except OSError:
                pass

        finally:
            if os.path.exists(temp_target):
                debug("Cleanup temporary directory at", temp_target)
                shutil.rmtree(temp_target, ignore_errors=True)

        # double check if directory exists now.
        if not os.path.exists(target):
            die("Renaming of temporary directory to %s failed" % target)

    return wrapper


@atomically
def install_virtualenv(target, url):
    tar_content = download(url)
    check_call("tar", "xz", "--strip-components=1", stdin=tar_content, cwd=target)


def analyze_install_error_output(output):
    """Analyzes the error output, prints a message and dies."""
    packages = []
    if "include <pyconfig.h>" in output or "include <Python" in output:
        packages.append("python2-devel")

    if "include <ffi.h>" in output:
        packages.append("libffi-devel")

    if "include <openssl" in output:
        packages.append("openssl-devel")

    if packages:
        die("\n\n[ERROR] Looks like packages are missing. Try: yum install %s" % " ".join(packages))
    else:
        die("\n\n[ERROR] Ansible could not be installed for some reason. Check log above.")


@atomically
def install_ansible(target, virtualenv_binary, ansible_version, extra_packages):
    info("Create virtualenv in", target)
    check_call(virtualenv_binary, "--system-site-packages", target)

    info("Install ansible in", target)
    python = os.path.join(target, "bin", "python")
    pip = os.path.join(target, "bin", "pip")
    rc, output = check_call(
        python, pip, "install", "--no-cache-dir", "--upgrade", "https://releases.ansible.com/ansible/ansible-%s.tar.gz" % ansible_version,
        capture_output=True, fail_on_error=False)

    if rc:
        raise analyze_install_error_output(output)

    if extra_packages:
        info("Install extra packages")
        check_call(python, pip, "install", "--no-cache-dir", "--upgrade", *extra_packages)

    info("Make venv relocatable to move it into its final place")
    check_call(virtualenv_binary, "--relocatable", target)

    info("Verify that the installed ansible executes")
    ansible = os.path.join(target, "bin", "ansible")
    check_call(ansible, "--version")

@atomically
def install_vars_plugins(target):
    info("Installing vars plugin to set ansible_python_interpreter.")
    with open(target + "/interpreter.py", "wb") as fp:
        fp.write(re.sub(r'^\s+[|] ', '', """
            | from ansible.plugins.vars import BaseVarsPlugin
            | class VarsModule(BaseVarsPlugin):
            |   def __init__(self, *args):
            |     pass
            |   def get_vars(self, loader, path, entities, cache=True):
            |     return {'ansible_python_interpreter': 'python'}
            """, 0, re.MULTILINE))

def update_wrapper():
    """Update the wrapper in place."""
    content = download("https://raw.githubusercontent.com/sumcumo/ansiblew/master/ansiblew")

    info("Replacing script at %s with new version" % sys.argv[0])
    with open(sys.argv[0], "wb") as fp:
        fp.write(content)


def main():
    if len(sys.argv) == 2 and sys.argv[1] == "update":
        return update_wrapper()

    config, config_hash = load_config()

    extra_requirements = config.get("requirements", [])
    ansible_version = config.get("ansible_version", DEFAULT_ANSIBLE_VERSION)
    venv_version = config.get("virtualenv_version", DEFAULT_VIRTUALENV_VERSION)

    debug("Run with virtualenv v%s, ansible v%s, packages=%s" % (
        venv_version, ansible_version, ", ".join(extra_requirements)))

    # where to put files
    workspace = os.path.join(find_workspace(), ansible_version, config_hash)
    workspace_venv = os.path.join(workspace, "venv")
    workspace_ansible = os.path.join(workspace, "ansible")

    # download and extract virtualenv
    install_virtualenv(workspace_venv, "https://github.com/pypa/virtualenv/archive/%s.tar.gz" % venv_version)

    # install ansible in a virtualenv
    virtualenv = os.path.join(workspace_venv, "virtualenv.py")
    install_ansible(workspace_ansible, virtualenv, ansible_version, extra_requirements)

    if "ANSIBLE_VARS_PLUGINS" not in os.environ:
        workspace_vars = os.path.join(workspace, "vars_plugins")
        install_vars_plugins(workspace_vars)
        os.putenv("ANSIBLE_VARS_PLUGINS", workspace_vars)

    # replace this process with the target process using os.exec
    ansible = os.path.join(workspace_ansible, "bin", ANSIBLE_COMMAND)
    os.execv(ansible, sys.argv)


if __name__ == '__main__':
    sys.exit(main() or 0)