standard-inventory-qcow2

#!/usr/bin/env python

import argparse
import errno
import json
import os
import shutil
import shlex
import signal
import socket
import subprocess
import sys
import tempfile
import time
import distutils.util

IDENTITY = """
-----BEGIN RSA PRIVATE KEY-----
MIIEpQIBAAKCAQEA1DrTSXQRF8isQQfPfK3U+eFC4zBrjur+Iy15kbHUYUeSHf5S
jXPYbHYqD1lHj4GJajC9okle9rykKFYZMmJKXLI6987wZ8vfucXo9/kwS6BDAJto
ZpZSj5sWCQ1PI0Ce8CbkazlTp5NIkjRfhXGP8mkNKMEhdNjaYceO49ilnNCIxhpb
eH5dH5hybmQQNmnzf+CGCCLBFmc4g3sFbWhI1ldyJzES5ZX3ahjJZYRUfnndoUM/
TzdkHGqZhL1EeFAsv5iV65HuYbchch4vBAn8jDMmHh8G1ixUCL3uAlosfarZLLyo
3HrZ8U/llq7rXa93PXHyI/3NL/2YP3OMxE8baQIDAQABAoIBAQCxuOUwkKqzsQ9W
kdTWArfj3RhnKigYEX9qM+2m7TT9lbKtvUiiPc2R3k4QdmIvsXlCXLigyzJkCsqp
IJiPEbJV98bbuAan1Rlv92TFK36fBgC15G5D4kQXD/ce828/BSFT2C3WALamEPdn
v8Xx+Ixjokcrxrdeoy4VTcjB0q21J4C2wKP1wEPeMJnuTcySiWQBdAECCbeZ4Vsj
cmRdcvL6z8fedRPtDW7oec+IPkYoyXPktVt8WsQPYkwEVN4hZVBneJPCcuhikYkp
T3WGmPV0MxhUvCZ6hSG8D2mscZXRq3itXVlKJsUWfIHaAIgGomWrPuqC23rOYCdT
5oSZmTvFAoGBAPs1FbbxDDd1fx1hisfXHFasV/sycT6ggP/eUXpBYCqVdxPQvqcA
ktplm5j04dnaQJdHZ8TPlwtL+xlWhmhFhlCFPtVpU1HzIBkp6DkSmmu0gvA/i07Z
pzo5Z+HRZFzruTQx6NjDtvWwiXVLwmZn2oiLeM9xSqPu55OpITifEWNjAoGBANhH
XwV6IvnbUWojs7uiSGsXuJOdB1YCJ+UF6xu8CqdbimaVakemVO02+cgbE6jzpUpo
krbDKOle4fIbUYHPeyB0NMidpDxTAPCGmiJz7BCS1fCxkzRgC+TICjmk5zpaD2md
HCrtzIeHNVpTE26BAjOIbo4QqOHBXk/WPen1iC3DAoGBALsD3DSj46puCMJA2ebI
2EoWaDGUbgZny2GxiwrvHL7XIx1XbHg7zxhUSLBorrNW7nsxJ6m3ugUo/bjxV4LN
L59Gc27ByMvbqmvRbRcAKIJCkrB1Pirnkr2f+xx8nLEotGqNNYIawlzKnqr6SbGf
Y2wAGWKmPyEoPLMLWLYkhfdtAoGANsFa/Tf+wuMTqZuAVXCwhOxsfnKy+MNy9jiZ
XVwuFlDGqVIKpjkmJyhT9KVmRM/qePwgqMSgBvVOnszrxcGRmpXRBzlh6yPYiQyK
2U4f5dJG97j9W7U1TaaXcCCfqdZDMKnmB7hMn8NLbqK5uLBQrltMIgt1tjIOfofv
BNx0raECgYEApAvjwDJ75otKz/mvL3rUf/SNpieODBOLHFQqJmF+4hrSOniHC5jf
f5GS5IuYtBQ1gudBYlSs9fX6T39d2avPsZjfvvSbULXi3OlzWD8sbTtvQPuCaZGI
Df9PUWMYZ3HRwwdsYovSOkT53fG6guy+vElUEDkrpZYczROZ6GUcx70=
-----END RSA PRIVATE KEY-----
"""

USER_DATA = """#cloud-config
users:
  - default
  - name: root
    groups: sudo
    shell: /bin/bash
    ssh_authorized_keys:
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDUOtNJdBEXyKxBB898rdT54ULjMGuO6v4jLXmRsdRhR5Id/lKNc9hsdioPWUePgYlqML2iSV72vKQoVhkyYkpcsjr3zvBny9+5xej3+TBLoEMAm2hmllKPmxYJDU8jQJ7wJuRrOVOnk0iSNF+FcY/yaQ0owSF02Nphx47j2KWc0IjGGlt4fl0fmHJuZBA2afN/4IYIIsEWZziDewVtaEjWV3InMRLllfdqGMllhFR+ed2hQz9PN2QcapmEvUR4UCy/mJXrke5htyFyHi8ECfyMMyYeHwbWLFQIve4CWix9qtksvKjcetnxT+WWrutdr3c9cfIj/c0v/Zg/c4zETxtp standard-test-qcow2
ssh_pwauth: True
chpasswd:
  list: |
    root:foobar
  expire: False
runcmd:
  - mkfs.ext4 /dev/sdb
  - mount /dev/sdb /usr/local 
  - sudo mount -o remount,rw /usr
"""

def main(argv):
    parser = argparse.ArgumentParser(description="Inventory for a QCow2 test image")
    parser.add_argument("--list", action="store_true", help="Verbose output")
    parser.add_argument('--host', help="Get host variables")
    parser.add_argument("subjects", nargs="*", default=shlex.split(os.environ.get("TEST_SUBJECTS", "")))
    opts = parser.parse_args()

    try:
        if opts.host:
            data = host(opts.host)
        else:
            data = list(opts.subjects)
        sys.stdout.write(json.dumps(data, indent=4, separators=(',', ': ')))
    except RuntimeError as ex:
        sys.stderr.write("{0}: {1}\n".format(os.path.basename(sys.argv[0]), str(ex)))
        return 1

    return 0


def list(subjects):
    hosts = []
    variables = {}
    for subject in subjects:
        if subject.endswith((".qcow2", ".qcow2c")):
            vars = host(subject)
            if vars:
                hosts.append(subject)
                variables[subject] = vars
    return {"localhost": {"hosts": hosts, "vars": {}}, "subjects": {"hosts": hosts, "vars": {}},
            "_meta": {"hostvars": variables}}


def start_qemu(image, cloudinit, log, disk_directory=None, disk_size=None, portrange=(2222, 5555)):
    for port in xrange(*portrange):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            sock.bind(("127.0.0.3", port))

            if disk_size:
                return subprocess.Popen(["/usr/bin/qemu-system-x86_64", "-m", "1024", image,
                                         "-enable-kvm", "-snapshot", "-cdrom", cloudinit,
                                         "-hdb", disk_directory,
                                         "-net", "nic,model=virtio", "-net",
                                         "user,hostfwd=tcp:127.0.0.3:{0}-:22".format(port),
                                         "-device", "isa-serial,chardev=pts2", "-chardev", "file,id=pts2,path=" + log,
                                         "-display", "none"], stdout=open(os.devnull, 'w')), port
            else:
                return subprocess.Popen(["/usr/bin/qemu-system-x86_64", "-m", "1024", image,
                                         "-enable-kvm", "-snapshot", "-cdrom", cloudinit,
                                         "-net", "nic,model=virtio", "-net",
                                         "user,hostfwd=tcp:127.0.0.3:{0}-:22".format(port),
                                         "-device", "isa-serial,chardev=pts2", "-chardev", "file,id=pts2,path=" + log,
                                         "-display", "none"], stdout=open(os.devnull, 'w')), port
        except IOError:
            pass
        finally:
            sock.close()
    else:
        raise RuntimeError("unable to find free local port to map SSH to")


def host(image):
    null = open(os.devnull, 'w')

    try:
        tty = os.open("/dev/tty", os.O_WRONLY)
        os.dup2(tty, 2)
    except OSError:
        tty = None
        pass

    # A directory for temporary stuff
    directory = tempfile.mkdtemp(prefix="inventory-cloud")
    identity = os.path.join(directory, "identity")
    with open(identity, 'w') as f:
        f.write(IDENTITY)
    os.chmod(identity, 0o600)
    metadata = os.path.join(directory, "meta-data")
    with open(metadata, 'w') as f:
        f.write("")
    userdata = os.path.join(directory, "user-data")
    with open(userdata, 'w') as f:
        f.write(USER_DATA)

    # Create additional disk
    disk_size = None
    disk_directory = None
    try:
        disk_size = os.environ.get("EXTEND_DISK_SIZE")
        if disk_size:
            sys.stderr.write("\nCreate additional cloud init disk DISK SIZE {}\n".format(disk_size))

            disk_directory = "{}/atomic-host-disk2-{}".format(directory, disk_size)
            subprocess.check_call(["qemu-img", "create", "-f", "qcow2", disk_directory, disk_size], stdout=null)

    except KeyError:
        sys.stderr.write("\nCouldn't create additional cloud init disk DISK SIZE\n")
    pass

    # Create our cloud init so we can log in
    cloudinit = os.path.join(directory, "cloud-init.iso")
    subprocess.check_call(["/usr/bin/genisoimage", "-input-charset", "utf-8",
                           "-volid", "cidata", "-joliet", "-rock", "-quiet",
                           "-output", cloudinit, userdata, metadata], stdout=null)

    # Determine if virtual machine should be kept available for diagnosis after completion
    try:
        diagnose = distutils.util.strtobool(os.getenv("TEST_DEBUG", "0"))
    except ValueError:
        diagnose = 0

    sys.stderr.write("Launching virtual machine for {0}\n".format(image))

    # And launch the actual VM
    artifacts = os.environ.get("TEST_ARTIFACTS", os.path.join(os.getcwd(), "artifacts"))
    try:
        os.makedirs(artifacts)
    except OSError as exc:
        if exc.errno != errno.EEXIST or not os.path.isdir(artifacts):
            raise
    log = os.path.join(artifacts, "{0}.log".format(os.path.basename(image)))

    proc = None  # for failure detection
    cpe = None  # for exception scoping
    for tries in xrange(0, 5):
        try:
            proc, port = start_qemu(image, cloudinit, log, disk_directory=disk_directory, disk_size=disk_size)
            break
        except subprocess.CalledProcessError as cpe:
            time.sleep(1)
            continue
    if proc is None:
        raise RuntimeError("Could not launch VM for qcow2 image"
                           " '{0}':{1}".format(image, cpe.output))

    # The variables
    variables = {"ansible_ssh_port": "{0}".format(port),
                 "ansible_ssh_host": "127.0.0.3",
                 "ansible_ssh_user": "root",
                 "ansible_ssh_pass": "foobar",
                 "ansible_ssh_private_key_file": identity,
                 "ansible_ssh_common_args": "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"}

    # Write out a handy inventory file, for our use and for debugging
    args = " ".join(["{0}='{1}'".format(*item) for item in variables.items()])
    inventory = os.path.join(directory, "inventory")
    with open(inventory, "w") as f:
        f.write("[subjects]\nlocalhost {1}\n".format(image, args))

    # Wait for ssh to come up
    ping = ["/usr/bin/ansible", "--inventory", inventory, "localhost", "--module-name", "raw", "--args", "/bin/true"]

    for tries in xrange(0, 30):
        try:
            (pid, ret) = os.waitpid(proc.pid, os.WNOHANG)
            if pid != 0:
                raise RuntimeError("qemu failed to launch qcow2 image: {0}".format(image))
            subprocess.check_call(ping, stdout=null, stderr=null)
            break
        except subprocess.CalledProcessError:
            time.sleep(3)
    else:
        # Kill the qemu process
        try:
            os.kill(proc.pid, signal.SIGTERM)
        except OSError:
            pass
        raise RuntimeError("could not access launched qcow2 image: {0}".format(image))

    # Process of our parent
    ppid = os.getppid()

    child = os.fork()
    if child:
        return variables

    # Daemonize and watch the processes
    os.chdir("/")
    os.setsid()
    os.umask(0)

    if tty is None:
        tty = null.fileno()

    # Duplicate standard input to standard output and standard error.
    os.dup2(null.fileno(), 0)
    os.dup2(tty, 1)
    os.dup2(tty, 2)

    # Now wait for the parent process to go away, then kill the VM
    while True:
        time.sleep(3)
        try:
            os.kill(ppid, 0)
            os.kill(proc.pid, 0)
        except OSError:
            break  # Either of the processes no longer exist

    if diagnose:
        sys.stderr.write("\n")
        sys.stderr.write("DIAGNOSE: ssh -p {0} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null "
                         "root@{1} # password: {2}\n".format(port, "127.0.0.3", "foobar"))
        sys.stderr.write("DIAGNOSE: export ANSIBLE_INVENTORY={0}\n".format(inventory))
        sys.stderr.write("DIAGNOSE: kill {0} # when finished\n".format(os.getpid()))

        def _signal_handler(*args):
            sys.stderr.write("\nDIAGNOSE ending...\n")

        signal.signal(signal.SIGTERM, _signal_handler)
        signal.pause()

    # Kill the qemu process
    try:
        os.kill(proc.pid, signal.SIGTERM)
    except OSError:
        pass

    shutil.rmtree(directory)
    sys.exit(0)


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