source: old/trunk/vtest/vm.py @ 2644

#!/usr/bin/python
# -*- coding: utf-8 -*-
# Copyright © 2009 by Karl Ramm
#
# All rights reserved.
#
# Permission to use, copy, modify, and distribute this software and
# its documentation for any purpose and without fee is hereby granted,
# provided that the above copyright notice appear in all copies and
# that both that copyright notice and this permission notice appear in
# supporting documentation, and that the name of Karl Ramm not be used
# in advertising or publicity pertaining to distribution of the
# software without specific, written prior permission.
#
# KARL RAMM DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS, IN
# NO EVENT SHALL KARL RAMM BE LIABLE FOR ANY SPECIAL, INDIRECT OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import os
import shlex
from os import execvp, _exit, dup2, fork, waitpid, close, chdir, \
    environ, WNOHANG, getgroups, kill

from time import sleep, time
from sys import exit
from pty import openpty
from itertools import count
from threading import Thread
from subprocess import Popen, PIPE
from logging import getLogger, basicConfig, DEBUG
from tempfile import NamedTemporaryFile
from grp import getgrnam
from signal import SIGTERM

environ['PAGER'] = 'cat'

basicConfig(level=DEBUG,
            format='%(asctime)s %(name)s.%(funcName)s:%(lineno)d %(message)s')
log = getLogger('vm')

runseq = count().next
def runcmd(cmd):
    this = runseq()
    if hasattr(cmd, 'split'):
        cmd = cmd.split()
    log.info('(%d) starting %s', this, cmd)
    p = Popen(cmd, stdout=PIPE, stderr=PIPE, close_fds=True)
    stdout, stderr = p.communicate()
    retcode = p.wait()
    log.info('(%d) returned %d', this, retcode)
    if stdout and stdout.strip():
        log.info('(%d) stdout = [%s]', this, stdout.strip())
    if stderr and stderr.strip():
        log.info('(%d) stderr = [%s]', this, stderr.strip())
    return retcode, stdout, stderr

vmcmd = 'qemu'
# probe for kvm_
retval, stdout, stderr = runcmd(['sh','-c','lsmod | grep kvm_'])
if retval == 0:
    try:
        gid = getgrnam('kvm').gr_gid
        if gid in getgroups():
            vmcmd='kvm'
    except:
        pass

class testmachine(object):
    testmachines = {}
    initparams = [
        ('basedir', os.getcwd()),
        ('domain', 'example.com'),
        ('hostname', 'test-%(index)d.%(domain)s'),
        ('base_image', '%(basedir)s/master.img'),
        ('image', '%(basedir)s/%(hostname)s.img'),
        ('mkimage', 'qemu-img create -b %(base_image)s -f qcow2 %(image)s'),
        ('qemu_binary', vmcmd),
        ('qemuopts', '-nographic'),
        ('memory', 256),
        ('netargs', '-net nic,vlan=1,macaddr=02:00:00:00:00:%(index)02d -net vde,vlan=1,sock=/var/run/vde2/tap0.ctl'),
        ('diskargs', '-hda %(image)s'),
        ('qemu', '%(qemu_binary)s -m %(memory)s %(netargs)s %(diskargs)s %(qemuopts)s'),
        ('logdir', '%(basedir)s'),
        ('logfile', '%(logdir)s/%(hostname)s.log'),
        ]
    def __init__(self, index=None, **kw):
        if index is not None:
            if index in self.testmachines:
                raise Exception('Test Machine #%d already exists' % index)
        else:
            for index in count():
                if index not in self.testmachines:
                    break
        self.testmachines[index] = self

        self.params = {'index': index}
        self.params.update(kw)
        for (name, template) in self.initparams:
            if name not in self.params:
                if isinstance(template, basestring): # we can't substitute, say, ints
                    self.params[name] = template % self.params
                else:
                    self.params[name] = template

        self.dead = False

        status, out, err = runcmd(self.params['mkimage'])
        if status != 0:
            raise Exception('Error creating image')

        self.qemuproc = loggedproc(self.params['qemu'], self.params['logfile'])
        self.boottime = time()

    def wait(self):
        if self.dead:
            raise Exception('dead testmachine')
        code = self.qemuproc.wait()
        self.dead = True
        log.debug('vm returned %d', code)
        del testmachine.testmachines[self.params['index']]
        return code

    def alive(self):
        if self.dead:
            return False
        alive = self.qemuproc.alive()
        if not alive:
            self.dead = True
        return alive

    def shoot(self):
        self.qemuproc.murder(SIGTERM)
        sleep(1)
        if self.qemuproc.alive():
            self.qemuproc.murder(SIGKILL)

class sshtestmachine(testmachine):
    initparams = testmachine.initparams + [
        ('sshsock', '%(basedir)s/ssh.%%h.%%r.%%p'),
        ('sshtarget', 'root@%(hostname)s'),
        ('sshmaster', 'ssh -NMS %(sshsock)s -o stricthostkeychecking=no -o userknownhostsfile=./known-hosts -i ./ssh_key %(sshtarget)s'),
        ('sshopt', '-o controlpath=%(sshsock)s -o stricthostkeychecking=no -o userknownhostsfile=./known-hosts -i ./ssh_key'),
        ('ssh', 'ssh %(sshopt)s %(sshtarget)s'),
        ('scp', 'scp %(sshopt)s'),
        ]

    def __init__(self, index=None, **kw):
        super(sshtestmachine, self).__init__(index, **kw)

        self.sshstate = 'go'
        self.sshrunning = None
        self.sshstart = None
        if self.params['sshmaster']:
            self.sshthread = Thread(target=self.sshmaster)
            self.sshthread.start()

    def sshmaster(self):
        log.info('starting sshmaster thread')
        wait = 20
        log.debug('waiting for %d seconds', wait)
        sleep(wait)
        while self.sshstate == 'go':
            self.sshstart = time()
            self.sshrunning = True
            runcmd(self.params['sshmaster'])
            self.sshrunning = False
            log.debug('sshmaster duration = %f, relative start = %f',
                      time() - self.sshstart,
                      self.sshstart - self.boottime)
            log.debug('pausing for a second')
            sleep(1)

    def run(self, cmd):
        log.debug('%s: running %s', self.params['hostname'], cmd)
        return runcmd(self.params['ssh'] + ' ' + cmd)

    def _scp(self, source, dest):
        retval, stdout, stderr = runcmd(' '.join([self.params['scp'],
                                        source, dest]))
        if retval != 0:
            raise Exception('scp failed', retval, stdout, stderr)

    def putfile(self, local, remote):
        log.debug('%s: copying local %s to remote %s', self.params['hostname'], local, remote)
        self._scp(local, '%s:%s' % (self.params['sshtarget'], remote))

    def getfile(self, remote, local):
        log.debug('%s: copying remote %s to local %s', self.params['hostname'], remote, local)
        self._scp('%s:%s' % (self.params['sshtarget'], remote), local)

    def putstr(self, remote, string):
        log.debug("%s: putting [%s] in remote %s", self.params['hostname'], string, remote)
        fp = NamedTemporaryFile()
        if string[-1] != '\n':
            string += '\n'
        fp.write(string)
        fp.flush()
        return self.putfile(fp.name, remote)

    def doscript(self, script):
        if hasattr(script, 'splitlines'):
            script = script.splitlines()
        for cmd in script:
            cmd = cmd.lstrip()
            if cmd and cmd[0] != '#':
                yield (cmd,) + self.run(cmd)

    def script(self, script, bail=True):
        results=[]
        for cmd, retval, stdout, stderr in self.doscript(script):
            if retval:
                raise(Exception(self.params['hostname'], cmd, retval, stdout, stderr))
            results.append((cmd, retval, stdout, stderr))
        return results #I'm not sure we should even bother

    def shutdown(self):
        log.debug('shutting down machine')
        self.sshstate = 'stop'
        self.run('poweroff')

    def shoot(self):
        self.sshstate='stop'
        super(sshtestmachine, self).shoot()

    def ready(self):
        t = time()
        log.debug('sshrunning = %s sshstart = %s t = %f',
                  self.sshrunning, self.sshstart, t)
        return self.sshrunning and self.sshstart < (t - 15.0)

class proc(object):
    def __init__(self, cmd):
        if hasattr(cmd, 'split'):
            self.cmdlist = shlex.split(cmd)
        else:
            self.cmdlist = cmd
        self.status = None
        self.start()
    def alive(self):
        if self.status is not None:
            return False
        pid, status  = waitpid(self.pid, WNOHANG)
        if pid:
            log.debug('%d returned %d in check', pid, status)
            self.status = status
            return False
        return True
    def start(self):
        log.debug('starting %s', self.cmdlist)
        self.pid = fork()
        if self.pid == 0:
            self._launch()
    def _launch(self):
        execvp(self.cmdlist[0], self.cmdlist)
        _exit(99)
    def wait(self):
        if self.status is None:
            pid, self.status = waitpid(self.pid, 0)
        return self.status
    def murder(self, sig=SIGTERM):
        kill(self.pid, sig)

class loggedproc(proc):
    def __init__(self, cmdlist, logfile):
        self.logfile = logfile
        self.master = None
        self.slave = None
        super(loggedproc, self).__init__(cmdlist)
    def start(self):
        self.master, self.slave = openpty()
        self.logfp = open(self.logfile, 'w')
        super(loggedproc, self).start()
        close(self.slave)
        self.logfp.close()
    def _launch(self):
        close(self.master)
        dup2(self.slave, 0) # stdin
        dup2(self.logfp.fileno(), 1) # stdout
        dup2(self.logfp.fileno(), 2) # stderr
        self.logfp.close()
        close(self.slave)
        super(loggedproc, self)._launch()

def main():
    pass

if __name__ == '__main__':
    main()