scripts/startvms.py

#!/usr/bin/env python
'''Start virtual machines when a network service becomes available

Example OpenRC init script::

    #!/sbin/runscript

    depend() {
        need libvirtd
    }

    start() {
        ebegin "Starting additional virtual machines"
        startvms -l ${STARTVMS_LOG:-/var/log/startvms.log} ${STARTVMS_ARGS} \
            ${VIRTUAL_MACHINES}
        eend $?
    }
'''

import argparse
import datetime
import libvirt
import logging
import os
import signal
import socket
import sys
import time


LIBVIRT_DEFAULT_URI = os.environ.get('LIBVIRT_DEFAULT_URI', 'qemu:///session')
WAIT_FOR_HOST = "arcturus.pyrocufflink.jazz"
WAIT_FOR_PORT = 389
WAIT_TIMEOUT = 300  # seconds


log = logging.getLogger('startvms')


class Timeout(Exception):
    pass


class MaxLevelFilter(logging.Filter):

    def __init__(self, max_level=logging.WARNING, *args, **kwargs):
        logging.Filter.__init__(self, *args, **kwargs)
        self.max_level = max_level

    def filter(self, record):
        return record.levelno <= self.max_level


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true', default=False)
    parser.add_argument('--foreground', '-F', action='store_true',
                        default=False,
                        help='Do not fork into the background before waiting')
    parser.add_argument('--log-file', '-l', metavar='FILENAME',
                        help='Send log messages to FILENAME')
    parser.add_argument('--connect', '-c', metavar='URI', dest='uri',
                        default=LIBVIRT_DEFAULT_URI,
                        help='libvirt connection URI')
    parser.add_argument('--timeout', '-t', default=WAIT_TIMEOUT, type=int,
                        metavar='SECONDS',
                        help='Wait a maximum of SECONDS')
    parser.add_argument('--port', '-p', default=WAIT_FOR_PORT, type=int,
                        help='Wait for PORT on remote host')
    parser.add_argument('--host', '-H', default=WAIT_FOR_HOST,
                        help='Wait for HOST to come up before starting VMs')
    parser.add_argument('vms', nargs='+', metavar='VMNAME',
                        help='Virtual machine(s) to start')
    return parser.parse_args()


def _setup_logging(debug=False, log_file=None):
    root_log = logging.getLogger()
    root_log.setLevel(logging.NOTSET)
    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_handler.setLevel(logging.ERROR)
    root_log.addHandler(stderr_handler)
    if log_file:
        file_handler = logging.FileHandler(log_file)
        if debug:
            file_handler.setLevel(logging.DEBUG)
        else:
            file_handler.setLevel(logging.INFO)
        file_formatter = logging.Formatter(
            '%(asctime)s [%(name)s] %(levelname)s %(message)s'
        )
        file_handler.setFormatter(file_formatter)
        root_log.addHandler(file_handler)
    else:
        stdout_handler = logging.StreamHandler(sys.stdout)
        stdout_handler.addFilter(MaxLevelFilter())
        if debug:
            stdout_handler.setLevel(logging.DEBUG)
        else:
            stdout_handler.setLevel(logging.WARNING)
        root_log.addHandler(stdout_handler)


def _daemonize():
    if os.fork():
        raise SystemExit(0)
    signal.signal(signal.SIGHUP, signal.SIG_IGN)
    os.setsid()
    os.chdir('/')
    if os.fork():
        raise SystemExit(0)


def wait_for(host, port, timeout=None):
    log.info('Waiting {time} for port {port} on {host}'.format(
        time='{0} seconds'.format(timeout) if timeout else 'indefinitely',
        port=port,
        host=host,
    ))
    if timeout is not None:
        end = datetime.datetime.now() + datetime.timedelta(seconds=timeout)
        keep_going = lambda: datetime.datetime.now() < end
    else:
        keep_going = lambda: True
    while keep_going():
        try:
            log.debug('Attempting connection to {0}:{1}'.format(host, port))
            s = socket.create_connection((host, port), 1)
        except:
            log.debug('Connection failed')
            time.sleep(1)
        else:
            log.debug('Connection succeeded')
            s.shutdown(socket.SHUT_RDWR)
            s.close()
            break
    else:
        raise Timeout('Timed out waiting for port {port} on {host}'.format(
            port=port,
            host=host,
        ))


def start_vm(conn, name):
    domain = conn.lookupByName(name)
    if domain.isActive():
        log.debug('Domain {0} is already active, skipping'.format(name))
    else:
        log.info('Starting domain {0}'.format(name))
        domain.create()


def main():
    args = _parse_args()
    if not args.foreground:
        _daemonize()
    _setup_logging(args.debug, args.log_file)

    try:
        wait_for(args.host, args.port, args.timeout)
    except Timeout as e:
        log.error('{0}'.format(e))
        raise SystemExit(os.EX_NOHOST)

    try:
        conn = libvirt.open(args.uri)
    except libvirt.libvirtError as e:
        log.error('{0}'.format(e))
        raise SystemExit(os.EX_UNAVAILABLE)
    error = False
    for vm in args.vms:
        try:
            start_vm(conn, vm)
        except libvirt.libvirtError as e:
            error = True
            log.error('{0}'.format(e))
            continue
    raise SystemExit(int(error))


if __name__ == '__main__':
    main()