#!/usr/bin/env python
"""
Spice Migration test

Somewhat stressfull test of continuous migration with spice in VGA mode or QXL mode,
depends on supplying an image in IMAGE variable (if no image is supplied then
VGA mode since it will just be SeaBIOS).

Dependencies:
either qmp in python path or running with spice and qemu side by side:
qemu/QMP/qmp.py
spice-tests/spice_migrate.py

Will create two temporary unix sockets in /tmp
Will leave a log file, migrate.log, in current directory.

Has a gui now:
    ./spice_migrate.py --gtk --sleep 5
    launch gtk gui
        wait for "start migration" button press
        update button label when migrations happen
    sleep 5 seconds between migrations

"""

"""
uses glib event loop for timers
"""

#
# start one spiceclient, have two machines (active and target),
# and repeat:
#  active wait until it's active
#  active client_migrate_info
#  active migrate tcp:localhost:9000
#  _wait for event of quit
#  active stop, active<->passive
#
# wait until it's active
#  command query-status, if running good
#  if not listen to events until event of running

try:
    import qmp
except:
    import sys
    sys.path.append("../../qemu/QMP")
    try:
        import qmp
    except:
        print "can't find qmp"
        raise SystemExit
import sys
from subprocess import Popen, PIPE
import os
import time
import socket
import datetime
import atexit
import argparse
import resource
import glib

def get_args():
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--qmp1', dest='qmp1', default='/tmp/migrate_test.1.qmp')
    parser.add_argument('--qmp2', dest='qmp2', default='/tmp/migrate_test.2.qmp')
    parser.add_argument('--spice_port1', dest='spice_port1', type=int, default=5911)
    parser.add_argument('--spice_port2', dest='spice_port2', type=int, default=6911)
    parser.add_argument('--migrate_port', dest='migrate_port', type=int, default=8000)
    parser.add_argument('--client_count', dest='client_count', type=int, default=1)
    parser.add_argument('--qemu', dest='qemu', default='qemu-system-x86_64')
    parser.add_argument('--log_filename', dest='log_filename', default='migrate.log')
    parser.add_argument('--image', dest='image', default='')
    parser.add_argument('--client', dest='client', default='spicy')
    parser.add_argument('--vdagent', choices=['on', 'off'], default='on')
    parser.add_argument('--usbtablet', choices=['on', 'off'], default='on')
    parser.add_argument('--secure', choices=['on', 'off'], default='off')
    parser.add_argument('--host-subject', default=None)
    parser.add_argument('--sleep', default=20, type=float)
    parser.add_argument('--glib', default=False, action='store_true')
    parser.add_argument('--gtk', default=False, action='store_true')
    args = parser.parse_args(sys.argv[1:])
    if os.path.exists(args.qemu):
        args.qemu_exec = args.qemu
    else:
        args.qemu_exec = os.popen("which %s" % args.qemu).read().strip()
    if not os.path.exists(args.qemu_exec):
        print "qemu not found (qemu = %r)" % args.qemu_exec
        sys.exit(1)
    orig_client = args.client
    if not os.path.exists(args.client):
        args.client = os.popen('which %s' % args.client).read().strip()
    if not os.path.exists(args.client):
        print "client %r not found" % orig_client
        sys.exit(1)
    print "qemu = %s" % args.qemu_exec
    print "LD_LIBRARY_PATH=%s" % os.environ.get('LD_LIBRARY_PATH', None)
    #print "ldd qemu:"
    #os.system('ldd %s' % args.qemu_exec)
    return args

def start_qemu(qemu_exec, image, spice_port, qmp_connection,
               incoming_port=None, extra_args=[], secure=False,
               debug=None, verbose=False):
    incoming_args = []
    if incoming_port:
        incoming_args = ("-incoming tcp::%s" % incoming_port).split()
    spice_params = ["disable-ticketing"]
    if secure:
        spice_params.extend(['tls-port=%s' % spice_port, 'x509-dir=%s' % os.getcwd()])
    else:
        spice_params.append("port=%s" % spice_port)
    if ':' not in qmp_connection:
        qmp_argument = "unix:%s,server,nowait" % qmp_connection
    else:
        host, port = qmp_connection.split(':')
        qmp_argument = "socket:host=%s,port=%s,server,nowait" % (host, port)
    args = ([qemu_exec, "--enable-kvm", "-qmp", qmp_argument,
             '-spice', ','.join(spice_params)] + extra_args + incoming_args)
    if os.path.exists(image):
        args += ["-m", "512", "-drive",
                 "file=%s,index=0,media=disk,cache=unsafe" % image, "-snapshot"]
    if debug:
        executable = debug[0]
        args = list(debug) + args
    else:
        executable = qemu_exec
    if verbose:
        print repr(args)
    proc = Popen(args, executable=qemu_exec, stdin=PIPE, stdout=PIPE)
    while not os.path.exists(qmp_connection):
        time.sleep(0.1)
    proc.qmp_connection = qmp_connection
    proc.qmp = qmp.QEMUMonitorProtocol(qmp_connection)
    while True:
        try:
            proc.qmp.connect()
            break
        except socket.error, err:
            pass
    proc.spice_port = spice_port
    proc.incoming_port = incoming_port
    return proc

def client_old_commandline(spice_port, secure, host_subject):
    cmdline = ['-h', 'localhost']
    if secure:
        cmdline.extend(['-s', str(spice_port), '--host-subject', host_subject, '--ca-file',
            os.path.join(os.getcwd(), 'ca-cert.pem')])
    else:
        cmdline.extend(['-p', str(spice_port)])
    return cmdline

def client_url_commandline(spice_port, secure, host_subject):
    if secure:
        raise Exception("don't know how a spice:// works for secure connections")
    return ['spice://localhost:%s' % (spice_port)]

def start_client(client, spice_port, secure, host_subject, verbose=False):
    client = str(client)
    if (client.endswith('remote-viewer')):
        cmdline = client_url_commandline(spice_port, secure, host_subject)
    else:
        cmdline = client_old_commandline(spice_port, secure, host_subject)
    cmdline = [client] + cmdline
    if verbose:
        print cmdline
        print ' '.join(cmdline)
    return Popen(cmdline, executable=client)

def _wait_active(q, active):
    events = ["RESUME"] if active else ["STOP"]
    while True:
        try:
            ret = q.cmd("query-status")
        except:
            # ValueError
            yield 'sleep', (0.1)
            continue
        if ret and ret.has_key("return"):
            if ret["return"]["running"] == active:
                break
        for e in q.get_events():
            if e["event"] in events:
                break
        yield 'sleep', (0.5)

def _wait_for_event(q, event):
    while True:
        for e in q.get_events():
            if e["event"] == event:
                return
        yield 'sleep', (0.5)

def cleanup(migrator):
    print "doing cleanup"
    migrator.close()

known_commands = {
    'wait_active': _wait_active,
    'wait_for_event': _wait_for_event
}

def blocker(it):
    for cmd, args in it:
        if not hasattr(args, '__len__'):
            args = (args,)
        if cmd in known_commands:
            blocker(known_commands[cmd](*args))
        elif cmd == 'sleep':
            time.sleep(*args)
        else:
            import pdb; pdb.set_trace()
            raise Exception('unhandled cmd %s' % cmd)

def glib_wrap_iter(it):
    def sleep_done(*args):
        glib_wrap_iter(it)
        return False
    try:
        cmd, args = it.next()
    except StopIteration:
        return
    if not hasattr(args, '__len__'):
        args = (args,)
    if cmd in known_commands:
        def handle_and_continue(cmd1, args1):
            it2 = known_commands[cmd1](*args1)
            for cmd, args in it2:
                yield cmd, args
            glib_wrap_iter(it)
        return glib_wrap_iter(handle_and_continue(cmd1=cmd, args1=args))
    elif cmd == 'sleep':
        assert(len(args) == 1)
        glib.timeout_add(int(1000 * args[0]), sleep_done)
    else:
        raise Exception("unhandled cmd")

def test_glib_wrap_iter():
    mainloop = glib.MainLoop()
    def it():
        t = time.time()
        yield 'sleep', 0.1
        t2 = time.time()
        yield 'sleep', 0.3
        t3 = time.time()
        print ('%2.2f, %2.2f' % (t2 - t, t3 - t2))
        mainloop.quit()
        raise SystemExit
    glib_wrap_iter(it())
    mainloop.run()

class Migrator(object):

    migration_count = 0
    on_migrate = lambda self, mig_count: None

    def __init__(self, log, client, qemu_exec, image, monitor_files, client_count=1,
                 spice_ports=[17000, 17001], migration_port=17002, vdagent=False,
                 usbtablet=False, secure=False, host_subject='',
                 extra_args=[], extra_per_vm=([], []), sound=False):
        self.client = client
        self.log = log
        self.qemu_exec = qemu_exec
        self.image = image
        self.migration_port = migration_port
        self.client_count = client_count
        self.monitor_files = monitor_files
        self.spice_ports = spice_ports
        self.vdagent = vdagent
        self.usbtablet = usbtablet
        self.secure = secure
        self.host_subject = host_subject
        if self.vdagent:
            extra_args.extend(['-device', 'virtio-serial',
                '-chardev', 'spicevmc,name=vdagent,id=vdagent',
                '-device', 'virtserialport,chardev=vdagent,name=com.redhat.spice.0'])
        if self.usbtablet:
            extra_args.extend(['-usb', '-device', 'usb-tablet'])
        if sound:
            extra_args.extend(['-device', 'intel-hda,id=sound0', '-device', 'hda-duplex,id=sound0-codec0,bus=sound0.0,cad=0'])
        self.extra_args = extra_args
        self.extra_per_vm = extra_per_vm
        self.active = self.start_qemu(which=0, incoming=False)
        self.target = self.start_qemu(which=1, incoming=True)
        self.remove_monitor_files()
        self.clients = []

    def start_qemu(self, which, incoming):
        kw = {}
        if incoming:
            kw['incoming_port'] = self.migration_port
        extra_args = self.extra_args + self.extra_per_vm[which]
        return start_qemu(qemu_exec=self.qemu_exec, image=self.image,
                          spice_port=self.spice_ports[which],
                          qmp_connection=self.monitor_files[which],
                          extra_args=extra_args,
                          secure=self.secure,
                          **kw)

    def close(self):
        self.remove_monitor_files()
        self.kill_qemu()

    def kill_qemu(self):
        for p in [self.active, self.target]:
            print "killing and waiting for qemu pid %s" % p.pid
            p.kill()
            p.wait()

    def remove_monitor_files(self):
        for x in self.monitor_files:
            if os.path.exists(x):
                os.unlink(x)

    def start_clients(self):
        if len(self.clients) == self.client_count:
            return 0
        start_count = len(self.clients)
        for i in range(len(self.clients), self.client_count):
            self.log.write("launching a client %d\n" % (len(self.clients) + 1))
            self.clients.append(start_client(client=self.client,
                spice_port=self.spice_ports[0], secure=self.secure, host_subject=self.host_subject))
        return len(self.clients) - start_count

    def _iterate(self, wait_for_user_input=False):
        yield 'wait_active', (self.active.qmp, True)
        yield 'wait_active', (self.target.qmp, False)
        started = self.start_clients()
        if started > 0:
            # wait for all clients events
            for i in xrange(started):
                yield 'wait_for_event', (self.active.qmp, 'SPICE_INITIALIZED')
            if wait_for_user_input:
                print "waiting for Enter to start migrations"
                raw_input()
        migrate_info_arguments = {'protocol':'spice', 'hostname':'localhost'}
        if self.secure:
            migrate_info_arguments['port'] = -1
            migrate_info_arguments['tls-port'] = self.target.spice_port
            migrate_info_arguments['cert-subject'] = self.host_subject
        else:
            migrate_info_arguments['port'] = self.target.spice_port
        self.active.qmp.cmd('client_migrate_info', migrate_info_arguments)
        self.active.qmp.cmd('migrate', {'uri': 'tcp:localhost:%s' % self.migration_port})
        yield 'wait_active', (self.active.qmp, False)
        yield 'wait_active', (self.target.qmp, True)
        yield 'wait_for_event', (self.target.qmp, 'SPICE_CONNECTED')
        dead = self.active
        dead.qmp.cmd("quit")
        dead.qmp.close()
        dead.wait()
        new_spice_port = dead.spice_port
        new_qmp_connection = dead.qmp_connection
        self.log.write("# STDOUT dead %s\n" % dead.pid)
        self.log.write(dead.stdout.read())
        del dead
        self.active = self.target
        which = 0 if new_spice_port == self.spice_ports[0] else 1
        self.target = self.start_qemu(which=which, incoming=True)
        self.migration_count += 1
        self.on_migrate(self.migration_count)

    def blocking_iterate(self, wait_for_user_input=False):
        blocker(self._iterate(wait_for_user_input))

    def loop(self, dt):
        dt = int(dt)
        self.count = 0
        def single():
            def single_iter():
                for cmd, args in self._iterate(False):
                    yield cmd, args
                self.count += 1
                glib.timeout_add(dt, single)
            glib_wrap_iter(single_iter())
            return False
        single()

def start_migrator(log, args, extra_args, extra_per_vm):
    migrator = Migrator(client=args.client, qemu_exec=args.qemu_exec,
        image=args.image, log=log, monitor_files=[args.qmp1, args.qmp2],
        migration_port=args.migrate_port, spice_ports=[args.spice_port1,
        args.spice_port2], client_count=args.client_count, vdagent=(args.vdagent=='on'),
        usbtablet=(args.usbtablet=='on'),
        secure=(args.secure=='on'), host_subject=args.host_subject,
        extra_args=extra_args, extra_per_vm=extra_per_vm, sound=False)
    atexit.register(cleanup, migrator)
    return migrator

def qmp_gui(args, migrator):
    import gtk
    mainloop = glib.MainLoop()
    def on_click(button):
        migrator.loop(args.sleep*1000)
    def on_quit(button):
        mainloop.quit()
    def on_migrate(mig_count):
        b.set_label("migrated %d" % mig_count)
    vbox = gtk.VBox()
    w = gtk.Window()
    w.add(vbox)
    bquit = gtk.Button('quit')
    bquit.connect("clicked", on_quit)
    vbox.add(bquit)
    b = gtk.Button("start migration")
    vbox.add(b)
    b.connect("clicked", on_click)
    w.show_all()
    migrator.on_migrate = on_migrate
    return mainloop

def main_gui(log, args, extra_args, extra_per_vm=([], [])):
    migrator = start_migrator(log=log, args=args, extra_args=extra_args,
                              extra_per_vm=extra_per_vm)
    migrator.start_clients()
    mainloop = qmp_gui(args, migrator)
    mainloop.run()

def main_console(log, args, extra_args):
    migrator = start_migrator(log=log, args=args, extra_args=extra_args,
                              extra_per_vm=([], []))
    if args.glib:
        migrator.loop(args.sleep*1000)
        mainloop = glib.MainLoop()
        mainloop.run()
    else:
        for i in xrange(50):
            migrator.blocking_iterate()
            print "sleeping %d seconds" % args.sleep
            time.sleep(args.sleep)

if __name__ == '__main__':
    # Limit core dumps to 10GB (default is no core dumps)
    resource.setrlimit(resource.RLIMIT_CORE, (10*1000*1000*1000, -1))
    args = get_args()
    extra_args = ['-vga', 'qxl', '-device', 'qxl']
    host_subject = None
    if args.secure == 'on':
        if any([not os.path.exists(f) for f in "ca-key.pem ca-cert.pem server-cert.pem server-key.pem".split()]):
            os.system('./spice_make_certs.sh')
        args.host_subject = ','.join(os.popen('openssl x509 -noout -text -in server-cert.pem | grep Subject: | cut -f 10- -d " "').read().strip().split(', '))
    print "log file %s" % args.log_filename
    log = open(args.log_filename, "a+")
    log.write("# "+str(datetime.datetime.now())+"\n")
    if args.gtk:
        main_gui(log, args, extra_args)
    else:
        main_console(log, args, extra_args)