#!/usr/bin/python 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', '/scratch'), ('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', 128), ('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', 'vdeq %(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 %(sshtarget)s'), ('sshopt', '-o controlpath=%(sshsock)s -o stricthostkeychecking=no'), ('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 = cmd.split() 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()