#!/usr/bin/python

import signal
import os
import sys
import subprocess
import atexit
import socket
import datetime
import re
import time
import shutil
import itertools

try:
    import argparse
except:
    print "missing argparse. try: easy_install argparse"
    sys.exit(1)

try:
    import guestfs
    import hivex
except:
    print "missing guestfs or hivex"
    guestfs = None

def which(x):
    path = os.environ['PATH']
    for f in path.split(':'):
        full_path = os.path.join(f, x)
        if os.path.exists(full_path):
            return full_path
    return None

################################################################################
lddpath = which('ldd')
spicedump_exe = which('spicedump.py')
python_exe = which('python')
#python_exe = which('pypy')
vncviewer = which('vncviewer')

HOME = os.environ['HOME']
################################################################################
# Edit <config_path> instead of this default section.
config_path = os.path.join(HOME, '.spice_launcher')
initial_defaults = """# vim: set filetype=python :
import os
HOME = os.environ['HOME']
qemu_exec_name = "qemu-system-x86_64"
images = {
    'windevel': '/store/images/windbg_winxp.img',
    'winxp':    '/store/images/winxp_sp3_qxl_after_cirrus.img',
    'win7':     '/home/alon/images2/win7x86_qxl_tests.img',
    'win2008r2': '/media/PassportExt4/images/win2008r2_qxl.qcow2',
    'wlk':       '/media/PassportExt4/images/win2008r2_wlk.qcow2',
    'f14ccid':  '/media/PassportExt4/images/F14_CCID.testing.qcow2',
    'winxp.inst': '/store/images/winxp_installation_test.qcow2',
}

spicy_exe = os.path.join(HOME, 'spice/upstream/bin/spicy')
spicec_exe = os.path.join(HOME, 'spice/upstream/bin/spicec')
spicy_teuf_exe = os.path.join(HOME, 'src/spice_upstream/teuf-spice-gtk/gtk/spicy')
roots = {'rhel6':'rhel6', 'upstream':'upstream'}
root_default = 'upstream'

defaults = dict(
    winxp=dict(
        root='upstream',
        args='--port 9001 --clients 1 --tablet --guestdebug 1 --qxldebug 1 --windbg-client 9000'.split(),
    ),
    wlk=dict(
        root='upstream',
        args='--port 9002 --clients 1 --tablet'.split(),
    ),
    win7=dict(
        root='upstream',
        args='--port 9003 --clients 1 --tablet --guestdebug 2 --qxldebug 1 --windbg-client 9000'.split()
    ),
    windevel=dict(
        root='rhel6',
        args='--port 9004 --clients 1 --windbg-server 9000'.split()
    ),
    f14ccid=dict(
        root='upstream',
        args='--port 9005 --clients 1 --tablet --guestdebug 4 --qxldebug 1 --smartcard'.split()
    ),
    win2008r2=dict(
        root='upstream',
        args='--port 9006 --clients 1 --tablet --revision 2'.split(),
    ),
)

spice_src_path = os.path.join(HOME, 'src/spice_upstream/spice')

qxl_win_root = os.path.join(HOME, 'shared_win/qxl/dist')

# required utilities for automated windows tests
res_change_exe = HOME+'/shared_win/spice_upstream/util/reschange/reschange.exe'
suspend_exe = HOME+'/shared_win/spice_upstream/util/suspend/suspend.exe'
"""

if not os.path.exists(config_path):
    with open(config_path, 'w+') as fd:
        fd.write(initial_defaults)
    print "created config file %s" % config_path
with open(config_path, 'r') as fd:
    exec(fd)

for name, default_location in [
        ('res_change_exe', 'res_change.exe'),
        ('suspend_exe', 'suspend.exe')]:
    if not os.path.exists(globals()[name]) and os.path.exists(default_location):
        print "config file points to non existant %r (%r), changing to local copy" % (
            name, globals()[name])
        globals()[name] = default_location


################################################################################

win_tests = dict(
shutdown=dict(contents="""
shutdown /p
"""),
suspend=dict(contents=r"""
%s
shutdown /p
""" % ("""
c:\suspend.exe
rem ping -n 2 localhost
""" * 20), files=[suspend_exe]),
suspend2=dict(contents=r"""
set RESCHANGE="c:\reschange.exe"
%%RESCHANGE%% -attach -mon 1024,768,32,0,0 -mon 800,600,32,1024,0
%s
shutdown /p
""" % (r"""
c:\suspend.exe
rem ping -n 2 localhost
""" * 20), files=[res_change_exe, suspend_exe]),
suspend_for_life=dict(contents=r"""
set RESCHANGE="c:\reschange.exe"
%RESCHANGE% -attach -mon 1024,768,32,0,0 -mon 800,600,32,1024,0
loop:
c:\suspend.exe
goto loop
""", files=[res_change_exe, suspend_exe]),
resolution1=dict(contents=r"""
set RESCHANGE="c:\reschange.exe"
%s
shutdown /p
""" % ("""
%RESCHANGE% -mon 800,600,32,0,0
%RESCHANGE% -mon 600,800,32,0,0
%RESCHANGE% -mon 1024,768,32,0,0
%RESCHANGE% -mon 768,1024,32,0,0
%RESCHANGE% -mon 1024,768,32,0,0
""" * 60), files=[res_change_exe]),
resolution2=dict(contents=r"""
set RESCHANGE="c:\reschange.exe"
%s
shutdown /p
""" % ("""
%RESCHANGE% -attach -mon 1024,768,32,0,0 -mon 800,600,32,1024,0
%RESCHANGE% -attach -mon 1024,768,32,0,0
%RESCHANGE% -attach -mon 1024,768,32,0,0 -mon 800,600,32,1024,0
%RESCHANGE% -attach -mon 1024,768,32,0,0
""" * 60), files=[res_change_exe]),
)

################################################################################

"""
tests are run without any clients connected, right now they can only
create / kill clients, later they should:
 send events to the clients (mouse/keyboard)
 parse all messages going to client and coming from client
 parse all messages going from qemu to guest agent and from it
  (presumably this can be done using a proxy and two chardevs)
 send qmp messages to qemu (connect/disconnect stuff, do screenshots)

Yes, this is all autotest stuff. maybe I should ask them to turn autotest into
a library.
"""
def test_agent_connection_stress(server, args):
    display = 2
    run_xephyr(display=display)
    auto_conf = itertools.cycle([False, True])
    times = itertools.chain([120], itertools.repeat(10))
    for i in xrange(100):
        client = run_client(host=args.spice_host, port=args.spice_port,
                   oldclient=True, smartcard=False,
                   auto_conf=auto_conf.next(),
                   certs=[], dbdir=None,
                   display=display)
        t = times.next()
        print "sleeping %d" % t
        time.sleep(t)
        client.spice_kill()
        client.wait()

g = globals().keys()
tests = [x[5:] for x in g if x.startswith('test_')]

################################################################################

red = "\x1b[101m%s\x1b[0m"

################################################################################

def set_title(x):
    sys.stdout.write("\033]0;%s\007" % str(x))

def file_in_use(x):
    return not subprocess.call(['lsof', '-t', x],
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)

def oldness(f):
    if not os.path.exists(f):
        return 'file doesn\'t exist'
    d = datetime.datetime.now() - datetime.datetime.fromtimestamp(os.stat(f).st_mtime)
    if d.days > 0:
        return '%d days' % d.days
    elif d.seconds > 3600:
        return '%d hours' % int(d.seconds / 3600)
    elif d.seconds > 60:
        return '%d minutes' % int(d.seconds / 60)
    return '%d seconds' % d.seconds

################################################################################

def get_lib(fname, rexp):
    l = os.popen('%s "%s"' % (lddpath, fname)).readlines()
    refs_raw=[x.split('=>') for x in l if '=>' in x]
    refs=[(r[0].strip().split()[0], r[1].strip().split()[0]) for r in refs_raw]
    ref = [r[1] for r in refs if re.match(rexp, r[0])]
    if len(ref) == 0:
        return 'library not found'
    return ref[0]

def make_iso(iso_name, root, title):
    cmd = 'mkisofs -J -R -V %s_%s -o %s %s' % (title, datetime.datetime.now().strtime('%Y%m%d_%H%M'), iso_name, root)
    print cmd
    os.system(cmd)

def version_tuple(symbol):
    return tuple(map(int, symbol.strip().rsplit('_', 1)[1].split('.')))

def get_versioned_library_version(lib, version_prefix):
    versions = [version_tuple(l) for l in os.popen('nm %s' % lib).readlines()
               if version_prefix in l]
    return max(versions) if len(versions) > 0 else None

def print_qemu_and_spice_versions(which, qemu_exec, clientexe, bios_dir):
    base = which + ' '
    bios_bin = os.path.join(bios_dir, 'bios.bin')
    print red % (base * (80 / len(base)))
    print "QEMU age: %s (%s)" % (oldness(qemu_exec), qemu_exec)
    print "bios.bin age: %s (%s)" % (oldness(bios_bin), bios_dir)
    for exe, lib_name, version_prefix in [
        (qemu_exec, 'spice-server', 'SPICE_SERVER'),
        (clientexe, 'cacard', 'CACARD')]:
        lib = get_lib(exe, 'lib%s.*' % lib_name)
        print "%s age: %s (%s)" % (lib_name, oldness(lib), lib)
        print "%s version: %r" % (lib_name,
            get_versioned_library_version(lib, version_prefix))
    print red % (base * (80 / len(base)))

def get_temp_name(prefix, postfix):
    global temp_files
    i = [0]
    def make_name(i):
        name = '%s_%03d%s' % (prefix, i[0], postfix)
        i[0] += 1
        return name
    filename = make_name(i)
    while os.path.exists(filename):
        filename = make_name(i)
    return filename

def create_qcow_with_backing_file(target_base, backing_file, erase_at_exit):
    filename = get_temp_name(target_base, '.qcow2')
    subprocess.call(['qemu-img', 'create', '-f', 'qcow2', '-o', 'backing_file=%s' % backing_file,
                    filename])
    if erase_at_exit:
        temp_files.append(filename)
    return filename

def get_image(args):
    image_base_fullpath = args.image
    if not os.path.exists(image_base_fullpath):
        print "missing image %s" % image_base_fullpath
        sys.exit(1)
    if not args.second and file_in_use(image_base_fullpath):
        print "image in use and --second not in arguments"
        sys.exit(1)
    if args.backing_file:
        base_image_name = os.path.basename(args.image)
        image_fullpath = create_qcow_with_backing_file(
            target_base='/tmp/'+base_image_name,
            backing_file=image_base_fullpath,
            erase_at_exit=args.erase_temp_backing_file)
    else:
        image_fullpath = image_base_fullpath
    print "image %s" % image_fullpath
    return image_fullpath

test_progs = {
    'streaming': 'test_display_streaming',
    'display': 'test_display_no_ssl',
    'playback': 'test_playback',
    'sockets': 'test_just_sockets_no_ssl',
    'replay': 'replay',
}
def build_test_cmdline(args):
    parts = (['/usr/bin/libtool', '--mode=execute', 'cgdb', '--args'] if args.cgdb else []) + [os.path.join(spice_src_path, 'server/tests',
                test_progs[args.test_prog])]
    if args.test_prog == 'replay' and args.replay_file:
        parts.append(args.replay_file)
    # tests have this port hardcoded
    args.port = 5912
    return ' '.join(parts)

def build_spice_option(args):
    ret = ["disable-ticketing"]
    if args.spice_port:
        ret.append('port=%s' % args.spice_port)
    if args.tls_port:
        if not args.x509_dir:
            if not args.server_key:
                print "have tls-port but missing server-key"
                sys.exit(1)
            if not args.server_cert:
                print "have tls-port but missing server-cert"
                sys.exit(1)
            if not args.ca_cert:
                print "have tls-port but missing ca-cert"
                sys.exit(1)
        ret.append('tls-port=%s' % args.tls_port)
        if args.x509_dir:
            ret.extend(['x509-dir=%s' % args.x509_dir])
        else:
            ret.extend([
                'x509-key-file=%s' % args.server_key,
                'x509-cert-file=%s' % args.server_cert,
                'x509-cacert-file=%s' % args.ca_cert,
                ])
        if args.server_key_password:
            ret.append('x509-key-password=%s' % args.server_key_password)
    if args.sasl:
        ret.append('sasl')
    return ['-spice', ','.join(ret)]

def exit_unless_exists(msg, filename):
    if not os.path.exists(filename):
        print msg % filename
        sys.exit(1)
    return filename

def get_root_dir(args):
    exit_unless_exists("root-dir %r is missing", args.root_dir)
    if not args.root:
        print "root is None"
        sys.exit(1)
    install_root = roots[args.root]
    root = os.path.join(args.root_dir, install_root)
    exit_unless_exists("missing %r", root)
    return root

def get_bios_dir(args):
    """ if roots[root] is not there default to whatever is in --biosdir,
    which would be the fedora/RHEL x86_64 default.
    """
    if args.bios_dir:
        return args.bios_dir
    if not args.root:
        print "--bios-dir not supplied nor --root"
        sys.exit(1)
    return os.path.join(get_root_dir(args), 'share/qemu')

def get_qemu_exec(args):
    """
    First check --qemu-exec
     if that is an executable file path assume it is qemu and take it
     if not check the root dict, if it is there check for a --root-dir,
     if not complain, if it is there then:
      add root_dir/roots[root]/lib to LD_LIBRARY_PATH
      (need this?)
       add root_dir/roots[root]/bin to PATH
      use qemu = root_dir/roots[root]/bin/qemu_exec_name

     qemu_exec_name is hardcoded to qemu-system-x86_64
    """
    if args.qemu_exec:
        exit_unless_exists("--qemu-exec doesn't exist: %r",
                           args.qemu_exec)
        qemu_exec = args.qemu_exec
    else:
        root = get_root_dir(args)
        libdir = os.path.join(root, 'lib')
        os.environ['LD_LIBRARY_PATH'] = libdir
        bindir = os.path.join(root, 'bin')
        os.environ['PATH'] = bindir + ':' + os.environ['PATH']
        qemu_exec = os.path.join(bindir, qemu_exec_name)
    return exit_unless_exists("missing qemu-exec %r", qemu_exec)

def bridge_participants(brname):
    interfaces = []
    if not os.path.exists('/sys/class/net/%s' % brname):
        print "no such bridge %s" % brname
        sys.exit(1)
    return os.listdir('/sys/class/net/%s/brif/' % brname)

def interface_list():
    with open('/proc/net/dev') as f:
        lines = f.readlines()[2:]
    return [x.split(':', 1)[0].strip() for x in lines]

def fatal_error(msg):
    print red % msg
    sys.exit(1)

def build_qemu_cmdline(args):
    """ Build qemu command line.

    Most of it is straight forward translation of the arguments from command
    line of spice_launcher to the qemu related parts.

    The start hinges on deciding where the qemu executable is. This script is
    designed to be used either (originally only) on development environment where
    qemu is installed into a non standard location (originally HOME/spice/<name>),
    or on a standard system where qemu and the bios are at different ('/' common
    prefix) locations.

    So the algorithm is:
     '--qemu' is absolute path? take that, expect a --bios-dir too
     --qemu exists? require --dest-dir, and make sure --qemu is in qemus dict (set in
     the config file, .spice_launcher).
    """
    if args.test_prog:
        return build_test_cmdline(args)
    bios_dir = get_bios_dir(args)
    qemu_exec = get_qemu_exec(args)
    args.image_fullpath = image_fullpath = get_image(args)
    if args.record_cmd:
        print "recording cmd ring to %r" % args.record_cmd
        os.environ['SPICE_WORKER_RECORD_FILENAME'] = args.record_cmd
    os.system('echo $LD_LIBRARY_PATH')
    cmdline = [qemu_exec,
        "-chardev", "stdio,id=muxstdio,mux=on", "-mon", "chardev=muxstdio,mode=readline",
        "-vga", args.vga,
        "-drive", "file=%s,cache=%s" % (args.image_fullpath, args.cache),
        "-enable-kvm",
        "-L", bios_dir, "-m",
        str(args.memory),
        "-cpu", str(args.cpu)] + build_spice_option(args)
    if args.snapshot:
        cmdline.append("-snapshot")
    if args.qmp_port:
        cmdline.extend(["-chardev",
            "socket,id=qmpmon,host=localhost,port=%s,server" % args.qmp_port,
            "-mon", "chardev=qmpmon,mode=control",])
    if args.no_shutdown:
        cmdline.append('-no-shutdown')
    if args.cpus > 1:
        cmdline.extend(['-smp', str(args.cpus)])
    if args.guestdebug > 0:
        cmdline.extend(['-global', 'qxl-vga.guestdebug=%d' % args.guestdebug])
    if args.qxldebug > 0:
        cmdline.extend(['-global', 'qxl-vga.debug=%d' % args.qxldebug])
    if args.cmdlog > 0:
        cmdline.extend(['-global', 'qxl-vga.cmdlog=%d' % args.cmdlog])
    if args.qxl > 1:
        cmdline.extend((['-device', 'qxl,guestdebug=%d,cmdlog=%d,debug=%d' % (args.guestdebug,
            args.cmdlog, args.qxldebug)]) * (args.qxl - 1))
    if args.revision:
        cmdline.extend(['-global', 'qxl-vga.revision=%d' % args.revision])
    if args.tablet:
        cmdline.extend(['-usb', '-device', 'usb-tablet'])
    if args.vdagent:
        cmdline.extend(['-device', 'virtio-serial', '-chardev', 'spicevmc,name=vdagent,id=vdagent',
                        '-device', 'virtserialport,chardev=vdagent,name=com.redhat.spice.0'])
    # user networking
    if args.smb:
        cmdline.extend(['-net', 'user,smb=%s,smbserver=10.0.2.2' % args.smb])
    if args.network == 'user':
        cmdline.extend('-netdev user,id=hostnet0'.split())
    elif args.network == 'bridge':
        sudo_prefix = ''
        if os.getuid() != 0:
            print "using sudo to create device"
            sudo_prefix = 'sudo '
        if args.tun in interface_list():
            print "tun device %s already exists, deleting" % args.tun
        if os.system('%stunctl -u %s -g %s -b -t %s' % (sudo_prefix, args.uid, args.gid, args.tun)):
            fatal_error("failed to create tun device")
        if os.system('%sifconfig %s 0.0.0.0 promisc up' % (sudo_prefix, args.tun)):
            fatal_error("ifconfig up failed")
        if args.tun in bridge_participants(args.bridge):
            print "someone is already using tun in this bridge, hopefully it is us"
        else:
            if os.system('%sbrctl addif %s %s' % (sudo_prefix, args.bridge, args.tun)):
                print "error calling brctl"
                sys.exit(1)
        cmdline.extend(['-netdev', 'tap,id=hostnet0,ifname=%s,script=no,downscript=no' % args.tun])
    cmdline.extend(["-device", "virtio-net-pci,netdev=hostnet0,id=net0,mac=%s" % args.mac])
    #cmdline.extend(["-net", "nic,model=virtio,macaddr=%s" % args.mac])
    # windbg (TODO - better then serial)
    if args.windbg_server:
        cmdline.extend(('-serial tcp::%d,server,nowait' % args.windbg_server).split())
    if args.windbg_client:
        if port_is_available(args.windbg_client):
            cmdline.extend(('-serial tcp:localhost:%d' % args.windbg_client).split())
        else:
            print "nothing listening on %d, do you want to start something there? (p.s. maybe systemd service?)" % args.windbg_client
    # smartcard
    if args.smartcard != 'off':
        ccid_debug = 1
        passthru_debug = 1
        cmdline.extend(['-chardev', 'spicevmc,id=smartcard,debug=3,name=smartcard',
            "-device", "usb-ccid,debug=%s,id=ccid" % ccid_debug,
            "-device", "ccid-card-passthru,debug=%s,chardev=smartcard" % (passthru_debug),
            ])
    # bios debugging
    if args.debugbios:
        cmdline.extend('-device isa-debugcon,iobase=0x402,chardev=muxstdio'.split())
    # incoming migration
    if args.incoming:
        cmdline.extend(('-incoming tcp:localhost:%d' % args.incoming).split())
    if args.freeze:
        cmdline.extend(['-S'])
    # vnc server
    if args.vnc_port:
        if (args.vnc_port < 5900):
            print "vnc_port must be >= 5900"
        cmdline.extend(('-vnc %s:%s' % (args.spice_host, args.vnc_port - 5900)).split())

    # cdrom
    if args.winqxl_cdrom:
        if args.cdrom:
            print "winqxl-cdrom and cdrom are mutually exclusive"
            sys.exit(1)
        if not os.path.exists(qxl_win_root):
            print "you need to edit qxl_win_root path in %s for this option" % (
                sys.argv[0])
            sys.exit(1)
        make_iso('/tmp/winqxl.iso', qxl_win_root, title='qxl')
        args.cdrom = '/tmp/winqxl.iso'
    if args.cdrom:
        if not os.path.exists(args.cdrom):
            print "missing cdrom image %r" % args.cdrom
            sys.exit(1)
        cmdline.extend(['-cdrom', args.cdrom])

    ######

    print_qemu_and_spice_versions(qemu_exec=qemu_exec,
        clientexe=get_client_ext(args.oldclient),
        which=qemu_exec, bios_dir=bios_dir)
    # wrap in various tools
    if args.cgdb:
        cgdb = which('cgdb')
        if not cgdb:
            gdb = which('gdb')
            if gdb:
                print "using gdb instead of cgdb"
                cgdb = gdb
            else:
                print "missing both cgdb and gdb"
                sys.exit(1)
        cmdline = [cgdb, '--args'] + cmdline
    if args.memcheck:
        print "wrapping with valgrind memcheck"
        cmdline = ['valgrind', '--error-limit=no', '--leak-check=full', '--suppressions=%s/bin/kvm.supp' % HOME,
            '--log-file=/tmp/kvm_%s.valgrind' % time.strftime('%Y%m%d_%H%M%S',time.gmtime())] + cmdline
    if args.time:
        print "wrapping in time"
        cmdline = ['/usr/bin/time', '--verbose'] + cmdline
    return cmdline

def run_xephyr(display, size=(1024,768+40), show_output=False):
    # TODO -wait for it to start listening to socket
    return start_process(('Xephyr -screen %dx%d :%d.0' %
        (size[0], size[1], display)).split())

def rpdb2_break():
    import rpbd2
    rpdb2.rpdb2.start_embedded_debugger('a')

def get_client_ext(oldclient):
    return spicec_exe if oldclient else spicy_exe

def run_client(host, port, oldclient,
               smartcard, certs, dbdir,
               tls_port=None, ca_file=None, host_subject=None,
               display=None, auto_conf=False):
    args = [get_client_ext(oldclient), '-h', host]
    if port:
        args.extend(['-p', str(port)])
    old_display = None
    if display is not None:
        old_display = os.environ['DISPLAY']
        os.environ['DISPLAY'] = ':%d.0' % display
    if smartcard == 'sw':
        if oldclient:
            args.extend(['--smartcard',
                '--smartcard-cert', ','.join(certs), '--smartcard-db', dbdir])
        else:
            args.extend(['--spice-smartcard',
                '--spice-smartcard-certificates=%s' % (','.join(certs)),
                '--spice-smartcard-db=%s' % dbdir])
    if auto_conf and oldclient:
        args.append('--full-screen=auto-conf')
    if tls_port:
        if oldclient:
            args.extend(['-s', str(tls_port), '--host-subject',
                         host_subject, '--ca-file', ca_file])
        else:
            args.extend(['-s', str(tls_port),
                         '--spice-host-subject=%s' % host_subject,
                         '--spice-ca-file=%s' % ca_file])
    if dbdir and not os.path.exists(dbdir):
        print "error - dbdir supplied does not exist"
        sys.exit(1)
    p = start_process(args=args, kill=True)
    if old_display:
        os.environ['DISPLAY'] = old_display
    return p

def run_vnc_client(host, port):
    return start_process(args=('%s %s::%s' % (vncviewer, host, port)).split(), kill=True)

def port_is_available(port):
    s = socket.socket(socket.AF_INET)
    ret = False
    try:
        s.connect(('localhost', port))
        s.close()
        ret = True
    except:
        pass
    return ret

def wait_for_port(port):
    # TODO: do this without actually opening the port - maybe just look at /proc/qemu_process_id/fd?
    s = socket.socket(socket.AF_INET)
    while True:
        try:
            s.connect(('localhost', port))
            s.close()
            break
        except:
            time.sleep(1)
            pass

################################################################################



def parseargs():
    """ Parse command line. Some of the parameters are pre populated:
     if sys.argv[0] is a symlink of the form:
      <name> - lookup name in images dictionary (set in config)
      <name>.<root> - lookup name in images and root in qemus.
    """
    image, root = None, root_default
    name = os.path.basename(sys.argv[0])
    default_args = []
    parts = []
    if os.path.islink(sys.argv[0]) and name != os.path.basename(os.readlink(sys.argv[0])):
        # you can symlink spice2 -> win7.upstream or win7, it will run the win7
        # from the images dictionary
        parts = name.split('.')
        base_part = parts[0] if len(parts) >= 1 else None
        if len(parts) == 1:
            if base_part in defaults and 'root' in defaults[base_part]:
                root = defaults[base_part]['root']
        if len(parts) == 2:
            image, root = parts
            image = os.path.basename(image)
        if base_part:
            if base_part in defaults:
                default_args = defaults[base_part]['args']
            if base_part not in images:
                print "missing key in images %r" % base_part
                print images.keys()
                sys.exit(1)
            image = images[base_part]
    parser = argparse.ArgumentParser(description='SPICE Launcher script, mark 2. Run spice, qemu, spice tests, automated tests on pre prepared images. Still missing: automated driver installation.')
    parser.add_argument('--record-cmd', dest='record_cmd', help='record qxl command ring')
    parser.add_argument('--copy-in', dest='copy_in', action='append')
    parser.add_argument('--no-backing-file', default=True, action='store_false', dest='backing_file', help='use the image as a backing file for the actual image')
    parser.add_argument('--startup', dest='startup')
    parser.add_argument('--wintest', dest='wintest', choices=win_tests.keys())
    parser.add_argument('--cache', dest='cache', choices=['unsafe','off','writeback','writethrough'], default='unsafe')
    parser.add_argument('--shutdown', dest='no_shutdown', action='store_false', default=True)
    parser.add_argument('--time', dest='time', action='store_true', default=False)
    parser.add_argument('--runonce', dest='runonce')
    parser.add_argument('--clear-runonce', dest='clear_runonce', default=False, action='store_true')
    parser.add_argument('--showcmdline', dest='showcmdline', default=False, action='store_true')
    parser.add_argument('--vga', dest='vga', choices=['qxl', 'cirrus'], default='qxl')
    parser.add_argument('--memcheck', dest='memcheck', default=False, action='store_true')
    parser.add_argument('--oldclient', dest='oldclient', default=True, action='store_true')
    parser.add_argument('--clients', dest='clients', default=0, type=int)
    parser.add_argument('--port', dest='spice_port', type=int, action='store')
    parser.add_argument('--host', dest='spice_host', default='127.0.0.1')
    parser.add_argument('--qxl', dest='qxl', type=int, default=1)
    parser.add_argument('--qmp-port', dest='qmp_port', type=int, default=0)
    parser.add_argument('--tablet', dest='tablet', action='store_true', default=False)
    parser.add_argument('--vdagent', dest='vdagent', action='store_true', default=False)
    parser.add_argument('--no-vdagent', dest='vdagent', action='store_false')
    parser.add_argument('--no-tablet', dest='tablet', action='store_false')
    parser.add_argument('--xephyr', dest='xephyr', action='store_true', default=False)
    parser.add_argument('--tls-port', type=int)
    parser.add_argument('--x509-dir')
    parser.add_argument('--ca-cert')
    parser.add_argument('--server-cert')
    parser.add_argument('--server-key')
    parser.add_argument('--server-key-password')
    parser.add_argument('--host-subject')
    parser.add_argument('--sasl', dest='sasl', action='store_true', default=False)
    parser.add_argument('--smb', dest='smb', default=None)
    parser.add_argument('--cgdb', dest='cgdb', default=False, action='store_true')
    parser.add_argument('--guestdebug', dest='guestdebug', type=int, default=0)
    parser.add_argument('--qxldebug', dest='qxldebug', type=int, default=0)
    parser.add_argument('--cmdlog', dest='cmdlog', type=int, default=0)
    parser.add_argument('--windbg-server', dest='windbg_server', type=int, default=None)
    parser.add_argument('--windbg-client', dest='windbg_client', type=int, default=None)
    parser.add_argument('--memory', dest='memory', type=int, default=512)
    parser.add_argument('--cpus', dest='cpus', type=int, default=2)
    parser.add_argument('--smartcard', dest='smartcard', choices=['off','hw','sw'], default='off')
    parser.add_argument('--smartcard-dbdir', dest='dbdir')
    parser.add_argument('--smartcard-certs', dest='certs', action='append')
    parser.add_argument('--debugbios', dest='debugbios', action='store_true', default=False)
    parser.add_argument('--image', dest='image', default=image)
    parser.add_argument('--snapshot', action='store_true', default=True)
    parser.add_argument('--no-snapshot', dest='snapshot', action='store_false')
    parser.add_argument('--cdrom', dest='cdrom', default=None)
    parser.add_argument('--winqxl-cdrom', dest='winqxl_cdrom', default=False, action='store_true')
    parser.add_argument('--qemu-exec', dest='qemu_exec')
    parser.add_argument('--root', default=root, choices=roots.keys() + [None])
    parser.add_argument('--root-dir', default=os.path.join(HOME, 'spice'))
    parser.add_argument('--bios-dir', default=None)
    parser.add_argument('--incoming', dest='incoming', type=int, default=None)
    parser.add_argument('--title', dest='title', default=None)
    parser.add_argument('--second', dest='second', default=False, action='store_true')
    parser.add_argument('--revision', dest='revision', default=None, type=int)
    parser.add_argument('--freeze', dest='freeze', default=False, action='store_true')
    parser.add_argument('--vncport', dest='vnc_port', type=int)
    parser.add_argument('--vncclients', dest='vnc_clients', type=int, default=1)
    parser.add_argument('--cpu', dest='cpu', choices=['host'], default='host', help='qemu cpu flag')
    parser.add_argument('--keep-temp-backing-file', dest='erase_temp_backing_file', default=True, action='store_false')
    parser.add_argument('--spicedump', dest='spicedump', default=False, action='store_true')
    parser.add_argument('--spicedump-filter', dest='spicedump_filter')
    parser.add_argument('--network', choices=['user', 'bridge', None], default='user')
    parser.add_argument('--mac', default='01:02:03:04:05:06')
    parser.add_argument('--bridge', default='virbr0')
    parser.add_argument('--tun', default='tap0')
    parser.add_argument('--uid', type=int, default=int(os.environ.get('SUDO_UID', os.getuid())))
    parser.add_argument('--gid', type=int, default=int(os.environ.get('SUDO_GID', os.getgid())))
    # convenient to run tests via the same framework, though this is not using
    # any of of the qemu flags, just the client running parts.
    parser.add_argument('--test-prog', dest='test_prog',
                        choices=test_progs.keys())
    parser.add_argument('--test-func', dest='test_func', choices=tests)
    parser.add_argument('--replay-file', dest='replay_file')
    args = parser.parse_args(default_args + sys.argv[1:])
    if args.image in images.keys():
        args.image = images[args.image]
    if len(parts) == 1 and not args.test_prog and not args.image:
        if image not in defaults:
            print "symlink %r not in %r" % (image, defaults.keys())
            sys.exit(1)
        if not args.root:
            args.root = defaults[image]['root']
    elif not args.image:
        print "image not provided. use --image"
        parser.print_usage()
        sys.exit(1)
    # tls arguments
    if args.x509_dir:
        args.ca_cert = os.path.join(args.x509_dir, 'ca-cert.pem')
    return args

def set_runonce_registry(g, root, runonce_paths):
    print "setting RunOnce to %r" % runonce_paths
    systemroot = g.inspect_get_windows_systemroot(root)
    path = "%s/system32/config/software" % systemroot # equivalent to HKLM\\Software
    path = g.case_sensitive_path(path)
    g.download (path, "/tmp/software")
    h = hivex.Hivex("/tmp/software", write=True)
    p = h.root()
    for k in r'Microsoft\Windows\CurrentVersion\RunOnce'.split('\\'):
        p = h.node_get_child(p, k)
    # 1 - hive_t_REG_SZ - unknown
    h.node_set_values(p, [{'key': 'Spice2RunOnce', 't': 1, 'value': runonce_path.encode('utf-16le')} for runonce_path in runonce_paths])
    h.commit(None)
    g.upload("/tmp/software", path)

def get_guestfs(image):
    if not guestfs:
        print "missing guestfs"
        sys.exit(1)
    g = guestfs.GuestFS()
    g.set_trace(1)
    g.add_drive_opts(image)
    g.launch()
    roots = g.inspect_os ()
    root = roots[0]
    g.mount_options("", root, "/")
    return g

def set_runonce(image, program):
    g = get_guestfs(image)
    if program is None:
       values = []
    else:
        runonce_base = 'runonce.bat'
        if os.path.exists(program):
            runonce_base = os.path.basename(program)
            g.upload(filename=program, remotefilename='/' + runonce_base)
        else:
            g.write('/' + runonce_base, program)
        values = ['c:\\' + runonce_base]
    set_runonce_registry(g, root, values)
    g.sync()
    g.umount_all()
    del g

def copy_file_in(g, image, path, dest):
    if g == None:
        g = get_guestfs(image)
    if not os.path.exists(path):
        print "%s does not exist" % path
        sys.exit(1)
    print "copying %r contents to c:\\" % (path)
    g.upload(path, os.path.join(dest, os.path.basename(path)))

def copy_dir_in(g, image, path):
    if not os.path.isdir(path):
        copy_file_in(g, image, path, '/')
        return g
    if g == None:
        g = get_guestfs(image)
    print "copying %r contents to c:\\" % (path)
    temp_name = shutil.make_archive('/tmp/spice2.copy_in', 'tar', path + '/')
    dest = '/' + os.path.basename(path)
    if not g.is_dir(dest):
        print "mkdir %r in image" % dest
        g.mkdir(dest)
    g.tar_in(tarfile=temp_name, directory=dest)
    os.unlink(temp_name)
    return g

startup = '/ProgramData/Microsoft/Windows/Start Menu/Programs/Startup'
def set_startup_from_file(g, image, path):
    if not os.path.exists(path):
        print "startup parameter is either a filename or a path"
    print "copying file %s to Startup on %s" % (path, image)
    filename = os.path.basename(path)
    dest = os.path.join(startup, filename)
    if g == None:
        g = get_guestfs(image)
    g.upload(filename=path, remotefilename=dest)
    return g

def set_startup_from_contents(g, image, filename, contents):
    dest = os.path.join(startup, filename)
    if g == None:
        g = get_guestfs(image)
    g.write(dest, contents)
    return g

def update_certs(args):
    if args.smartcard != 'sw' or args.certs != []:
        return
    if os.environ['SPICE_LAUNCHER_CERTS'] != '':
        print "missing SPICE_LAUNCHER_CERTS"
        sys.exit(1)
    if os.environ['SPICE_LAUNCHER_DB']:
        args.dbdir = os.environ['SPICE_LAUNCHER_DB']

################################################################################
processes = []
temp_files = []

class MyProcess(subprocess.Popen):
    def spice_kill(self):
        processes.remove(self)
        return self.kill()

def start_process(args, kill=False, show_output=True, **kw):
    global processes
    print repr(args)
    if show_output:
        p = MyProcess(args, **kw)
    else:
        p = MyProcess(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                      **kw)
    p._spice_kill = kill
    processes.append(p)
    return p

def cleanup():
    for p in processes:
        print "killing pid %s" % p.pid
        try:
            p.kill()
            p.kill() # qemu started with --no-shutdown, needs two signals
            p.wait()
        except OSError, e:
            pass
    for m in temp_files:
        if not os.path.exists(m):
            continue
        print "removing %s" % m
        os.unlink(m)
    os.system('reset -I')

def if_processes_are_dead_then_exit(*args):
    if all((p.poll() is not None or p._spice_kill) for p in processes):
        cleanup()
        sys.exit(0)
    #print '|'.join(('%d: %s, %s' % (p.pid, p.poll(), p._spice_kill)) for p in processes)

def ignore_ctrlc():
    signal.signal(signal.SIGINT, if_processes_are_dead_then_exit)

def exit_on_hup():
    def on_exit(*args):
        cleanup()
        sys.exit(1)
    old = signal.signal(signal.SIGHUP, on_exit)
    print "SIGHUP old:", old

################################################################################
clients = []
vnc_clients = []

def default_get_display():
    while True:
        yield None
def xephyr_get_display():
    i = 2
    while True:
        yield i
        i += 1

def main():
    args = parseargs()
    update_certs(args)
    cmdline = build_qemu_cmdline(args)
    atexit.register(cleanup)
    if args.runonce:
        set_runonce(args.image_fullpath, args.runonce)
    elif args.clear_runonce:
        set_runonce(args.image_fullpath, None)
    g = None
    if args.startup:
        g = set_startup_from_file(g, args.image_fullpath, args.startup)
    if args.wintest:
        g = set_startup_from_contents(g, image=args.image_fullpath, filename='startup.bat',
            contents=win_tests[args.wintest]['contents'])
        # NOTE: can just use autoplay and a suitable cdrom iso for these, faster then guestfs by far.
        for filename in win_tests[args.wintest].get('files', []):
            g = copy_file_in(g, args.image_fullpath, filename, '/')
    if args.copy_in:
        for p in args.copy_in:
            g = copy_dir_in(g, args.image_fullpath, p)
    if g:
        g.sync()
        g.umount_all()
        del g
    set_title("%s%s" % (
        'I ' if args.incoming else '',
        args.title if args.title else os.path.basename(sys.argv[0])
        ))
    if args.showcmdline:
        print repr(cmdline)
        return
    ignore_ctrlc()
    exit_on_hup()
    # At this point drop privileges before executing any process. Networking and
    # anything else requiring super-user privileges should have already been done.
    if os.getuid() == 0:
        print "dropping super user privileges"
        os.setuid(args.uid)
        os.setgid(args.gid)
    test_process = start_process(args=cmdline, show_output=True)
    if args.xephyr:
        display = xephyr_get_display()
    else:
        display = default_get_display()
    if args.spicedump:
        real_spice_port = args.spice_port
        args.spice_port = proxy_port = real_spice_port + 1000
        # run a spicedump proxy in the middle
        start_process(args=('%s %s -p -l %s -r localhost:%s%s' %
            (python_exe, spicedump_exe, proxy_port, real_spice_port,
             ' -f %s' % args.spicedump_filter if args.spicedump_filter else ''
             )).split())
    if args.qmp_port:
        #start_process(
        print "TODO"
    if args.test_func:
        test_func = globals()['test_'+args.test_func]
        test_func(server=test_process, args=args)
    else:
        if args.clients > 0:
            print "starting clients"
            cur_display = display.next()
            if cur_display:
                run_xephyr(display=cur_display)
            if args.spice_port:
                wait_for_port(args.spice_port)
            else:
                wait_for_port(args.tls_port)
            for i in xrange(args.clients):
                client_process = run_client(
                    host=args.spice_host, port=args.spice_port,
                    oldclient=args.oldclient, smartcard=args.smartcard,
                    certs=args.certs, dbdir=args.dbdir,
                    display=cur_display,
                    tls_port=args.tls_port,
                    ca_file=args.ca_cert,
                    host_subject=args.host_subject)
                clients.append(client_process)
        if args.vnc_port and args.vnc_clients > 0:
            if not vncviewer:
                print "vncviewer not in path - not starting"
            else:
                print "starting vnc clients"
                wait_for_port(args.vnc_port)
                for i in xrange(args.vnc_clients):
                    client_process = run_vnc_client(host=args.spice_host, port=args.vnc_port)
                    vnc_clients.append(client_process)
    print "waiting for test_process"
    test_process.wait()

if __name__ == '__main__':
    main()