#!/usr/bin/python
""" vws - script to control QEMU/KVM virtual workstations """
# pylint: disable=bad-builtin
from ConfigParser import ConfigParser
from argparse import ArgumentParser, Namespace
import fcntl
import socket, select
import errno
import re
import os, sys, time, os.path

VERSION = 0.4
def find_vm(name):
    """ Search and return VM directory """
    search_path = [os.path.join(os.environ['HOME'], "VWs"),
                   config.get("directories", "SharedVMs"),
                   config.get("directories", "AutostartVMs")]
    for dirname in search_path:
        if not os.access(dirname, os.X_OK):
            continue
        if (name in os.listdir(dirname) and
                os.access(os.path.join(dirname, name, "start"), os.X_OK)):
            return os.path.join(dirname, name)
    raise ValueError("Machine %s not found." % name)


def connect_vm(vm_dir):
    """ Connects to monitor of VM in vm_dir and returns connected socket"""
    sock = socket.socket(socket.AF_UNIX)
    monitor_path = os.path.join(vm_dir, "monitor")
    if not os.access(monitor_path, os.W_OK):
        return None
    try:
        sock.connect(monitor_path)
    except IOError as ex:
        if ex.errno == errno.ECONNREFUSED:
            # virtal machine is not running
            return None
        else:
            raise ex
    readfd, dummy_w, dummy_x = select.select([sock], [], [], 0.001)
    if sock in readfd:
        dummy_greeting = sock.recv(1024)
    return sock

def send_command(sock, command):
    """ Sends monitor command to given socket and returns answer """
    fcntl.flock(sock, fcntl.LOCK_EX)
    try:
        sock.send(command + "\n")
        answer = ""
        while not answer.endswith("(qemu) "):
            chunk = sock.recv(1024)
            if chunk == '':
                raise IOError("Unexpected EOF From monitor")
            answer += chunk
    finally:
        fcntl.flock(sock, fcntl.LOCK_UN)
    return answer

def spiceurl(sock):
    """ Returns spice URI for given (as set of parsed args) VM """
    output = send_command(sock, "info spice")
    url = None
    for line in output.split("\n"):
        if url is not None:
            continue
        idx = line.find("address:")
        if idx != -1:
            url = line[idx+9:]
            if url.startswith('*:'):
                url = "localhost"+url[1:]
    if url is None:
        return None
    return "spice://" + url.rstrip('\r')



def list_bridges():
    """ Return list of bridge network interfaces present in the system """
    lst = []
    with os.popen(config.get('tools', 'bridge_list'), "r") as f:
        for line in f:
            idx = line.find('\t')
            if idx <= 0:
                continue
            name = line[:idx]
            if name == "bridge name":
                continue
            lst.append(name)
    return lst

def parse_arp(iface):
    """
    Returns map which maps mac addresses to IPs for specified interface"
    """
    addr_map = {}
    pipe = os.popen(config.get("tools","arp")+" -n -i "+iface, "r")
    for line in pipe:
        data = line.split() 
        mac=data[2]
        if mac == "HWAddress":
            continue
        if mac == iface:
            # Foind line with (incomplete) entry
            continue
        addr_map[data[2]]=data[0]
    return addr_map

def validate_size(size):
    """ Checks if size argument has proper format """
    return re.match('\\d+[KMG]', size) is not None

def get_drives(vm_dir):
    """ Return list of drive files in the VW directory """
    result = []
    with open(vm_dir + "/start") as f:
        for line in f:
            if (re.match("\\s*-drive .*", line) and
                    line.find("media=disk") > -1):
                match = re.search("file=([^,\\s]*)", line)
                if match:
                    result.append(match.group(1))
    return result

def snapshot_mode(sock):
    """ Returns True if VM is running in snapshot mode """
    answer = send_command(sock, "info block")
    return re.search(": /tmp", answer) is not None

def read_netinfo(filename):
    """ Reads network information from start script """
    with open(filename,"r") as f:
       for line in f:
          match=re.search("-net nic,macaddr=(\\S+) -net ([^,]+)",line)
          if match:
              f={"mac":match.group(1)}
              if match.group(2) == "user":
                   f["iface"]="user"
              elif match.group(2) == "bridge":
                   f["iface"]=re.search("br=(\\S+)",line).group(1)
              else:
                   f["iface"]="unknown";
              return f
    return {"iface":"unknown","mac":"?","ip":"?"}
def get_netinfo(sock):
    """ Gets network information from the running VM """
    answer = send_command(sock, "info network")
    match=re.search("bridge\\.0:.*,br=(\\S+).*macaddr=(\\S+)", answer, re.S)
    if match:
        return {"iface":match.group(1), "mac":match.group(2)}
    else:
        match = re.search("user.0:.*net=([^,]+).*\n.*macaddr=(\\S+)",answer)
        if match:
           return {"iface":"user", "ip":match.group(1), 
                    "mac":match.group(2)}
        else:
           print >>sys.stderr,answer
           return {"iface":"unknown","ip":"?","mac":"?","card":"?"}
#
# command implementation
#

def cmd_spiceuri(options):
    """ vws spiceuri """
    print spiceurl(options.sock)


def cmd_start(options):
    """ vws start """
    if options.stopped:
        arg = ""
        if options.cdrom:
            arg = " -cdrom " + os.path.abspath(options.cdrom[0])
        if options.snapshot:
            arg = arg+" -snapshot"
        if options.args:
            arg = arg + " " + "".join(options.args)
        print arg
        cwd = os.getcwd()
        os.chdir(options.dir)
        # Check for snapshot
        nxt = 0
        snapshot_id = None
        with os.popen("qemu-img info \"%s\"" %
                      (get_drives(options.dir)[0]), "r") as f:
            for line in f:
                if line == 'Snapshot list:\n':
                    nxt = 2
                elif nxt == 2 and line.startswith('ID'):
                    nxt = 1
                elif nxt == 1:
                    nxt = 0
                    snapshot_id = line[:line.index(' ')]
                    arg = arg + " -loadvm " + snapshot_id
                    break

        os.system("./start%s" % arg)
        os.chdir(cwd)
        time.sleep(2)
        options.sock = connect_vm(options.dir)
        if snapshot_id:
            send_command(options.sock, "delvm " + snapshot_id)
    else:
        if options.snapshot or options.args:
            print >>sys.stderr, ("Cannot change qemu options. " +
                                 "VM is already running")
        if options.cdrom:
            options.file = options.cdrom[0]
            options.id = None
            cmd_cdrom(options)
    if options.gui:
        uri = spiceurl(options.sock)
        os.system((config.get('tools', 'viewer') + "&") % uri)
    elif not options.stopped:
        print >>sys.stderr, "VM already running"

def cmd_stop(options):
    """ vws stop """
    if snapshot_mode(options.sock) or options.hard:
        print send_command(options.sock, 'quit')
    else:
        print send_command(options.sock, 'system_powerdown')

def cmd_monitor(options):
    """ vws monitor """
    try:
        print "(qemu) ",
        sys.stdout.flush()
        while True:
            readfd, dummy_w, dummy_x = select.select([sys.stdin, options.sock],
                                                     [], [])
            if sys.stdin in readfd:
                cmd = sys.stdin.readline()
                # Check for eof
                if len(cmd) == 0:
                    break
                answer = send_command(options.sock, cmd.rstrip())
                idx = answer.index('\n')
                print answer[idx+1:],
                sys.stdout.flush()
            elif options.sock in readfd:
                print "UNSOLICITED MESSAGE %" + options.sock.readline().rstrip()
    except KeyboardInterrupt:
        print "Keyboard interrupt"
        sys.exit()

def cmd_reset(options):
    """ vws reset """
    print send_command(options.sock, 'system_reset')

def cmd_save(options):
    """ vws save """
    answer = send_command(options.sock, 'savevm')
    if re.search("Error", answer):
        print >>sys.stderr, answer
        sys.exit(1)
    else:
        send_command(options.sock, 'quit')

def cmd_cdrom(options):
    """ vws cdrom """
    if options.id  is None:
        # Search for devices which could be interpreted as CDROM
        devlist = send_command(options.sock, "info block")
        for dev in  re.findall("([-\\w]+): [^\n]+\n    Removable device:",
                               devlist):
            if re.search("cd", dev):
                options.id = dev
                break
    if options.id is None:
        print >>sys.stderr, "No CDROM device found among:\n" + devlist
        sys.exit(1)
    if options.file == "":
        print >>sys.stderr, "Please specify either --eject or iso image"
        sys.exit(1)
    if options.file is None:
        answer = send_command(options.sock, "eject " + options.id)
    else:
        answer = send_command(options.sock, "change %s %s" %
                              (options.id, os.path.abspath(options.file)))
    print answer

def find_usb(options, devices):
    """ Search for pattern or address given in options in the
        given list of devices.
        List should be produced by get_host_devices() or
        get_vm_devices()
    """
    if hasattr("pattern", options):
        for dev in devices:
            if re.search(options.pattern, dev[1]):
                options.address = dev[0]
                break
    elif not hasattr("address", options):
        print >>sys.stderr, ("Address or search pattern for device " +
                             "is not specified")
        sys.exit(1)
    else:
        return options.address

def get_host_devices():
    """ Parses output of lsusb into list of tuples "address" "descr" """
    global config
    f = os.popen(config.get('tools', "lsusb"), "r")
    lst = []
    for dev in f:
        match = re.match('Bus (\\d+) Device (\\d+): (.*)$', dev)
        if match:
            if match.group(3).endswith("root hub"):
                continue
            lst.append((match.group(1) + "." + match.group(2), match.group(3)))
    f.close()
    return lst

def get_vm_devices(sock):
    """ Parses output of info usb monitor command into list of devices"""
    answer = send_command(sock, "info usb")
    lst = []
    for dev in answer.split("\n"):
        match = re.match('Device (\\d+\\.\\d+), .*?, Product (.*)$', dev)
        if match:
            lst.append((match.group(1), match.group(2)))
    return lst

def cmd_usb_insert(options):
    """ vws usb insert """
    address = find_usb(options, get_host_devices())
    answer = send_command(options.sock, "usb_add host:%s" % address)
    print answer

def cmd_usb_list(dummy_options):
    """ vws usb list - just list all host devices """
    for addr, descr in get_host_devices():
        print addr, ": ", descr

def cmd_usb_attached(options):
    """ vws usb attached - list devices assigned to given vm """
    for dev in get_vm_devices(options.sock):
        print "Address %s : %s" % (dev[0], dev[1])

def cmd_usb_remove(options):
    """ vws usb remove """
    address = find_usb(options, get_vm_devices(options.sock))
    answer = send_command(options.sock, "usb_del %s" % address)
    print answer

def cmd_list(options):
    """ vws list """
    count = 0
    search_path = [("private",os.environ['HOME'] + "/VWs"),
                   ("shared",config.get("directories", "SharedVMs")),
                   ("autostart",config.get("directories", "AutostartVMs"))]

    maxlen = 0
    vms = []
    bridges = set()
    for (vmtype,dirname) in search_path:
        if not os.access(dirname + "/.", os.X_OK):
            continue
        for vmname in os.listdir(dirname):
            if os.access(dirname + "/" + vmname + "/start", os.X_OK):
                count += 1
                f = {"name":vmname}
                if maxlen < len(vmname):
                    maxlen = len(vmname)
                if options.state:
                    f["type"]=vmtype
                    sock = connect_vm(dirname + "/" + vmname)
                    if sock is None:
                        f.update({"state":"stopped","uri":"-","ip":"-"})
                        f.update(read_netinfo(dirname + "/" + vmname + "/start"))
                    else:
                        uri=spiceurl(sock)
                        f["uri"]=uri[uri.rindex(":")+1:]
                        f.update(get_netinfo(sock))
                        if "ip" not in f:
                            bridges.add(f["iface"])
                        sock.shutdown(socket.SHUT_RDWR)
                        sock.close()
                        f["state"] = "running"
                vms.append(f)
    arp_data={}
    for bridge in bridges:
        arp_data.update(parse_arp(bridge))
    for f in sorted(vms,key=lambda x: x["name"]):
        if "state" in f:
            if "mac" in f and not "ip" in f:
                if f["mac"] in arp_data:
                    f["ip"] = arp_data[f["mac"]]
                else:
                    f["ip"] = "-"
            f["name"] = f["name"].ljust(maxlen)
            print "%(name)s %(state)s   %(type)-9s %(uri)-4s %(iface)-5s %(mac)s %(ip)s " % f
        else:
            print f["name"]
    if not count:
        sys.exit(1)

def cmd_screenshot(options):
    """ vws screenshot """
    from os.path import abspath
    filename = abspath(options.filename)
    print send_command(options.sock, "screendump " + filename)

def cmd_record(options):
    """ vws record """
    from os.path import abspath
    filename = abspath(options.filename)
    print send_command(options.sock, "wavcapture " + filename)

def cmd_stoprecord(options):
    """ vws stoprecord """
    answer = send_command(options.sock, "info capture")
    match = re.search('\\[(\\d+)\\]: ', answer)
    if not match:
        print >>sys.stderr, "No sound recording in progress"
        sys.exit(1)
    else:
        print send_command(options.sock, "stopcapture " + match.group(1))

def cmd_sendkey(options):
    """ vws sendkey """
    print send_command(options.sock, "sendkey " + options.keyspec)

def cmd_version(dummy_options):
    """ vws cersion """
    print VERSION


def cmd_snapshot(options):
    """ vws snapshot - create snapshot """
    if not options.stopped:
        print >>sys.stderr, "Cannot make snapshot of running VW"
        sys.exit(1)
    drives = get_drives(options.dir)
    os.chdir(options.dir)
    newnames = {}
    for i in drives:
        name, ext = os.path.splitext(i)
        newnames[i] = name + "." + options.snapname + ext
        if os.path.exists(newnames[i]):
            print >>sys.stderr, "Snapshot %s already exists", options.snapname
            return 1
    for i in drives:
        os.rename(i, newnames[i])
        os.system("qemu-img create -f qcow2 -b \"%s\" \"%s\"" %
                  (newnames[i], i))
        os.chmod(i,0664)
    return 0

def cmd_snapshots(options):
    """ vws snapshots - list existing snapshots """
    os.chdir(options.dir)
    drives = get_drives(options.dir)
    lst = []
    info = {}
    with os.popen("qemu-img info --backing-chain " + drives[0], "r") as f:
        for line in f:
            if line.find(": ") != -1:
                var, val = line.strip().split(": ")
                if val != "":
                    info[var] = val
            elif line[0] == '\n':
                lst.append(info)
                info = {}
        lst.append(info)
    for snap in lst:
        print "%-30s  %+8s  %+8s" % (snap["image"], snap["virtual size"],
                                     snap["disk size"])


def get_backing(drive):
    """ find if partucular virtual drive has backing file and returns
        its name
    """
    with os.popen('qemu-img info "%s"' % drive, "r") as f:
        for line in f:
            match = re.match("backing.file: (.*)$", line)
            if match:
                return match.group(1)
    return None

def cmd_revert(options):
    """ reverts to last snapshot:
      Removes latest snapshot images and creates new ones instead
      number of snapshots in the stack is not changed by this command
    """
    if not options.stopped:
        print >>sys.stderr, "Cannot revert running VW to snapshot"
        sys.exit(1)
    os.chdir(options.dir)
    for drive in get_drives(options.dir):
        # Check if first  has backing file
        backing = get_backing(drive)
        if not backing:
            print >>sys.stderr, "Drive %s has no snapshots" % drive
            continue
        # Unlink current image
        os.unlink(drive)
        # create new image with same backing file
        os.system('qemu-img create -f qcow2 -b "%s" "%s"' % (backing, drive))
        os.chmod(drive,0664)

def cmd_commit(options):
    """
    Commits last snapshot changes into it's backing file
    There would be one snapshot less for virtual machine
    """
    if options.stopped:
        #
        # Stoppend vm - last snapshot is commited into its backing file.
        # Backing file is made current drive image
        #
        os.chdir(options.dir)
        found = 0
        for drive in get_drives(options.dir):
            backing = get_backing(drive)
            if backing is None:
                continue
            found = 1
            os.system('qemu-img commit "%s"'%drive)
            os.unlink(drive)
            os.rename(backing, drive)
        if not found:
            print >>sys.stderr, "No snapshots exist for this VM"
            sys.exit(1)
    else:
        if snapshot_mode(options.sock):
            send_command(options.sock, "commit")
        else:
            print >>sys.stderr, "VM is not running in snapshot mode"
            sys.exit(1)

def cmd_autostart(options):
    """
    Starts all VMs which are in the autostart directory
    """
    import pwd
    dirname =  config.get("directories", "AutostartVMs")
    os.seteuid(pwd.getpwnam(
                            config.get("permissions","autostart_user")
                            ).pw_uid)
    if not os.access(dirname,os.R_OK):
        return 
    for name in os.listdir(dirname):
        if not os.access(os.path.join(dirname,name,"start"), os.X_OK):
            continue
        machine_dir = os.path.join(dirname,name)
        sock = connect_vm(machine_dir)
        if sock:
            # Machine already running
            sock.shutdown(socket.SHUT_RDWR)
            sock.close()
            continue
        start_opts = Namespace(machine = name,
                command = 'start',
                dir = machine_dir,
                snapshot = False,
                stopped = True,
                atrgs = "",
                gui = False,
                cdrom = None)
        try:
            cmd_start(start_opts)
            print name," ",
        finally:
            # pylint: disable=no-member
            start_opts.sock.shutdown(socket.SHUT_RDWR)
    print ""

def cmd_shutdown(options):
    """ Search for all running machines and stops all of them """
    dirlist=[config.get("directories","AutostartVMs"),
             config.get("directories","SharedVms")]
    if os.getresuid()[1] == 0:  
        import grp
        dirlist += map(lambda x: os.path.expanduser("~"+x)+"/VWs",
        grp.getgrnam(config.get("permissions","vm_group")).gr_mem)
    else:
        dirlist.append(os.path.expanduser("~")+"/VWs")
    count = 1 #Fake positive values for there is not postcondition loop in the python
    forced_finish=time.time()+options.timeout
    while count > 0:
        count = 0
        if time.time() < forced_finish:
            command = "system_powerdown"
        else:
            command = "quit"
            
        for dirname in dirlist:
            if not os.access(dirname, os.X_OK):
                continue
            for vm in os.listdir(dirname):
                if not os.access(os.path.join(dirname,vm,"start"),os.X_OK):
                    # Not a VM
                    continue
                sock=connect_vm(os.path.join(dirname,vm))
                if not sock:
                    # not running
                    continue
                count += 1
                try:
                    send_command(sock,command)
                except IOError:
                    #When hard_stopping,socket might be closed by way
                    pass
                    sock.shutdown(socket.SHUT_RDWR)
                    sock.close()
        if not options.wait:
            return
        if count>0:
            time.sleep(10)


TEMPLATE = """#!/bin/sh
# Get machine name from current directory name
NAME=$(basename $(pwd))
# if remote access is enabled, then there should be
# SPICE_PASSWORD=password
QEMU_AUDIO_DRV=spice
export QEMU_AUDIO_DRV
if [ -n "$SPICE_PASSWORD" ]; then
   SPICE_AUTH="password=$SPICE_PASSWORD"
else
   SPICE_AUTH="disable-ticketing,addr=127.0.0.1"
fi
SPICE_PORT=$(find_free_port 5900)
if [ "$1" = '-cdrom' ]; then
	shift
	CDROM=",file=$1"
	shift
fi
#set umask to make machine group-accessable
umask 002
{qemubinary} -name $NAME {accel} \\
-m {memory} \\
{drive} \\
{cdrom}$CDROM \\
{net} \\
{usb} \\
{sound} \\
-chardev socket,server,nowait,path=monitor,id=monitor \\
-mon chardev=monitor,mode=readline \\
-vga {vga} \\
-spice port=$SPICE_PORT,$SPICE_AUTH \\
-device virtio-serial -chardev spicevmc,id=vdagent,name=vdagent \\
-device virtserialport,chardev=vdagent,name=com.redhat.spice.0 \\
-device ich9-usb-ehci1,id=usb \\
-device ich9-usb-uhci1,masterbus=usb.0,firstport=0,multifunction=on \\
-chardev spicevmc,name=usbredir,id=usbredirchardev1 \\
-device usb-redir,chardev=usbredirchardev1,id=usbredirdev1 \\
-daemonize -pidfile pid
chmod 0660 monitor pid
"""

def cmd_create(parsed_args):
    """ vws create - create new VM """
    BADSIZE = "Invalid size of %s specifed %s. Should have K, M or G suffix"
    global TEMPLATE
    if not parsed_args.image and not validate_size(parsed_args.size):
        print >>sys.stderr, BADSIZE % ("disk", parsed_args.size)
        sys.exit(1)
    if not validate_size(parsed_args.mem):
        print >>sys.stderr, BADSIZE % ("memory", parsed_args.size)
        sys.exit(1)
    drivename = "drive0.qcow2"
    options = {'qemubinary':'qemu-system-x86_64',
               "accel":"-enable-kvm",
               "memory":"1024M",
               "vga":'qxl',
               "drive":"-drive media=disk,index=0,if={interface},file={image}",
               "cdrom":"-drive media=cdrom,index=2,if=ide",
               "sound":"-soundhw hda",
               "usb":"-usb"}
    macaddr = ":".join(["%02x" % ord(x) for x in  chr(0x52) + os.urandom(5)])
    if parsed_args.shared:
        machinedir = os.path.join(config.get("directories", "SharedVMs"),
                                  parsed_args.machine)
        dirmode = 0775
    else:
        machinedir = os.path.join(os.environ["HOME"], "VWs",
                                  parsed_args.machine)
        dirmode = 0775

    if parsed_args.net != 'user':
        bridges = list_bridges()
        if not parsed_args.net in bridges:
            print >>sys.stderr, ("No such bridge %s. Available ones %s" %
                                 (parsed_args.net, ", ".join(bridges)))
            sys.exit(1)
        options["net"] = ("-net nic,macaddr=%s -net bridge,br=%s" %
                          (macaddr, parsed_args.net))
    else:
        options["net"] = "-net nic,macaddr=%s -net user" % (macaddr,)
    options["qemubinary"] = 'qemu-system-' + parsed_args.arch
    options["vga"] = parsed_args.vga
    NOACCEL = "KVM acceleration disabled due to "
    if not parsed_args.arch in ('i386', 'x86_64'):
        print >>sys.stderr, NOACCEL + "target architecture"
        options.accel = ''
    elif not os.access("/dev/kvm", os.W_OK):
        print >>sys.stderr, NOACCEL + "unavailability on the host system"
        options.accel = ''

    if not parsed_args.usb:
        options["usb"] = ''

    if not parsed_args.sound:
        options["sound"] = ''
    else:
        options["sound"] = '-soundhw ' + parsed_args.sound

    options["memory"] = parsed_args.mem

    if os.path.exists(machinedir):
        if os.path.exists(os.path.join(machinedir, "start")):
            print >> sys.stderr, ("Virtual Worstation %s already exists" %
                                  parsed_args.machine)
        else:
            print >> sys.stderr, ("Cannot create VW directory, " +
                                  "something on the way")
        sys.exit(1)
    #  Creating directory for VM
    os.makedirs(machinedir, dirmode)
    if parsed_args.shared:
        import grp
        gid=grp.getgrnam(config.get("permissions","vm_group")).gr_gid
        uid=os.getuid() 
        os.chown(machinedir,uid,gid)
        if config.getboolean("permissions","setgid_vm"):  
            os.chmod(machinedir,02775)
    driveopts = {"interface":parsed_args.diskif, "image":drivename}
    if parsed_args.install:
        install_image = os.path.abspath(parsed_args.install)

    if parsed_args.image:
        # Copying image file
        print >>sys.stderr, ("Copying %s to %s" %
                             (parsed_args.image,
                              os.path.join(machinedir, drivename)))
        os.system("qemu-img convert -O qcow2 -p %s %s" %
                  (parsed_args.image,
                   os.path.join(machinedir, drivename)))
        os.chdir(machinedir)
    else:
        print >>sys.stderr, "Creating new image file of %s" % parsed_args.size
        os.chdir(machinedir)
        os.system("qemu-img create -f qcow2 %s %s" %
                  (drivename, parsed_args.size))
    os.chmod(drivename,0664)
    # pylint: disable=star-args
    options["drive"] = options["drive"].format(**driveopts)
    if hasattr(parsed_args, "debug") and parsed_args.debug:
        print repr(driveopts), repr(options["drive"])
        print repr(options)
    with open("start", "w") as script:
        script.write(TEMPLATE.format(**options))
    os.chmod('start', dirmode)
    # If installation media is specified vws start for new vm
    if parsed_args.install:
        start_opts = Namespace(machine=parsed_args.machine,
                               command='start', cdrom=[install_image],
                               dir=machinedir, stopped=True, snapshot=False,
                               args="", gui=True)
        try:
            cmd_start(start_opts)
        finally:
            # pylint: disable=no-member
            start_opts.sock.shutdown(socket.SHUT_RDWR)

#
# Utility functions for arg parsing
#
def new_command(cmd_parser, name, **kwargs):
    """
    Adds a subparser and adds a machine name argument to it
    """
    parser = cmd_parser.add_parser(name, **kwargs)
    parser.add_argument('machine', type=str, help='name of vm to operate on')
    return parser
#
# prepare defaults for config
#

arch = os.uname()[4]
if re.match("i[3-9]86", arch):
    arch = "i386"
elif arch.startswith("arm"):
    arch = "arm"

config = ConfigParser({'SharedVMs':'/var/cache/vws/shared',
                       'AutoStartVMs':'/var/cache/vws/autostart'})
config.add_section('directories')
config.add_section('create options')
for option, value in [('net', 'user'), ('size', '20G'), ('mem', '1G'),
                      ('diskif', 'virtio'), ('sound', 'hda'), ('arch', arch),
                      ('vga', 'qxl')]:
    config.set('create options', option, value)
config.add_section('tools')
config.set('tools', 'viewer', 'remote-viewer %s')
config.set('tools', 'bridge_list', '/sbin/brctl show')
config.set('tools', 'lsusb', 'lsusb')
config.set('tools', 'arp', '/usr/sbin/arp')
config.add_section('permissions')
config.set('permissions','vm_group','kvm')
config.set('permissions','autostart_user','root')
config.set('permissions','setgid_vm','yes')
# Read configration files
config.read(['/etc/vws.conf', os.environ['HOME'] + '/.vwsrc'])
# Parse argument
args = ArgumentParser(description="Manage Virtual Workstations")
cmds = args.add_subparsers(dest='command', help="sub-command help")
p = cmds.add_parser("list", help="List existing VWs",
                    description="List existing VWs")
p.add_argument("--state", action='store_const', const=True, default=False,
               dest='state', help='Show state of the machine')
p.add_argument("--usb", action='store_const', const=True, default=False,
               dest='usb', help='Show connected USB devices')
p = cmds.add_parser("version", help="show vws version")
p = cmds.add_parser("autostart",help="Autostart all VMs marked as autostartable")
p = cmds.add_parser("shutdown",help="shut down all running VMs")
p.add_argument("--wait",help="wait until all machines would be shoutdown",
    action="store_const", const=True, default=False, dest="wait")
p.add_argument("--timeout",type=int,default=90,
   help="how long to way for VMs shutdown before forcing it off")
# Power management
p = new_command(cmds, 'start', help='Start VW and connect to console',
                description="Start VW if not running and connect " +
                "to the console")
p.add_argument('--no-gui', dest='gui', action='store_const', const=False,
               default=True, help='do not open console window')
p.add_argument('--cdrom', metavar='filename.iso', dest='cdrom', nargs=1,
               help='connect specified iso image to VMs cdrom on start')
p.add_argument('--args', metavar='string', dest='args', nargs=1, default="",
               help="Specify extra QEMU options")
p.add_argument("--snapshot", action='store_const', const=True, default=False,
               help="Run without modifying disk image")
p = new_command(cmds, 'stop', help='Shut down virtual machine',
                description="Terminate the VW, gracefully or ungracefully")
p.add_argument('--hard', help='Power off immediately', action='store_const',
               dest='hard', const=True, default=False)
new_command(cmds, 'save', help='Save VW state and stop emulation',
            description="Save VW state and stop emulation")
new_command(cmds, 'reset', help='Reboot a guest OS',
            description="Reboot othe guuest OS")
# Removable devices management
p = new_command(cmds, 'cdrom', help='manage CDROM Drive',
                description='"insert" an ISO image into emulated CD-ROM ' +
                "or eject it")
p.add_argument('--id', type=str, default=None,
               help='Identifier of CDROM drive if VM has more than one')
p.add_argument('file', nargs="?", default='',
               help='ISO image or special file to connect to drive')
p.add_argument('--eject', dest='file', action='store_const', const=None)
usb = cmds.add_parser('usb', help='manage USB devices'
                     ).add_subparsers(dest='subcommand',
                                      help='manage USB devices')
p = new_command(usb, 'insert', help='attach device to the virtual machine')
p.add_argument('pattern', help='Pattern of device name to look up in lsusb')
p.add_argument('--address', type=str, dest='address', nargs=1,
               help='exact address bus:device')
p = new_command(usb, 'remove', help='detach connected usb device')
p.add_argument('pattern', help='Pattern of device name to look up in lsusb')
p.add_argument('--address', type=str, dest='address', nargs=1,
               help='exact address bus:device')
p = new_command(usb, 'attached', help='list devices attached to vm')
usb.add_parser('list', help='list devices available in the host system')
# Snapshot management
p = new_command(cmds, 'snapshot', help='Create new snapshot')
p.add_argument('snapname', help='snapshot name')
p = new_command(cmds, 'revert', help='Revert to snapshot')
p.add_argument('snapname', help='name of snapshot to revert to')
p = new_command(cmds, 'commit',
                help='Commit snapshot changes into backing file')
p = new_command(cmds, 'snapshots', help='List existing snapshots')
# Screenshoits and recording
p = new_command(cmds, 'screenshot', help='take a screenshot')
p.add_argument('filename', help='PPM image filename to write screenshot to')
p = new_command(cmds, 'record', help='Record audio output from VM')
p.add_argument('filename', help='wav file to record autdio to')
new_command(cmds, 'stoprecord', help='stop recording audio')
p = new_command(cmds, 'sendkey', help='Send a keystroke to VM')
p.add_argument('keyspec', help='key specification like ctrl-alt-delete')
# Create new VM
p = new_command(cmds, 'create', help="Create new VW")
p.add_argument("--no-usb", help="Disable USB controller", action='store_const',
               const=False, default=True, dest="usb")
p.add_argument("--size", metavar='size', help="Size of primary disk images",
               dest="size", default=config.get('create options', 'size'))
p.add_argument("--arch", metavar='cputype', help="Emulated architecture",
               dest="arch", default=config.get('create options', 'arch'))
p.add_argument("--no-sound", help="Disable sound card", action='store_const',
               const=None, default=config.get('create options', 'sound'),
               dest="sound")
p.add_argument("--sound", metavar='cardtype', help="Specify sound card type",
               dest='sound', default=config.get('create options', 'sound'))
p.add_argument("--vga", metavar='cardtype',
               help="specify video card type (cirrus,std,vmware,qxl) default " +
               config.get('create options', 'vga',), dest="vga",
               default=config.get('create options', 'vga'))
p.add_argument("--net", help="Network - 'user' or bridge name",
               dest='net', default=config.get('create options', 'net'))
p.add_argument("--mem", metavar='size', help="Size of memory",
               dest="mem", default=config.get('create options', 'mem'))
p.add_argument("--diskif", metavar='interface-type',
               help="Disk interface (virtio, scsi, ide)",
               choices=['virtio', 'scsi', 'ide'],
               dest="diskif", default=config.get('create options', 'diskif'))
p.add_argument('--shared', help='Create shared VW instead of private one',
               action='store_const', const=True, dest='shared', default=False)
p.add_argument('--image', metavar='filename', dest='image', default=None,
               help='Existing disk image to import')
p.add_argument('--install', metavar='filename.iso', dest='install',
               help='ISO image to install OS from', default=None)
# Miscellenia
p = new_command(cmds, 'monitor', help='connect stdin/stdout to monitor of VM')
p = new_command(cmds, 'spiceuri', help='Output spice URI of machine')

parsed_args = args.parse_args(sys.argv[1:])

os.umask(002)
# Create command is totally different, so it is handled separately
if parsed_args.command == 'create':
    cmd_create(parsed_args)
    sys.exit(0)

funcname = "cmd_" + parsed_args.command
if hasattr(parsed_args, "subcommand"):
    funcname += "_" + parsed_args.subcommand
try:
    func = globals()[funcname]
except KeyError:
    print >>sys.stderr, "Operation %s is not implemented" % funcname
    sys.exit(3)

parsed_args.stopped = False
stopped_vm_commands = ['start', 'snapshot', 'revert', 'commit', 'snapshots']
if hasattr(parsed_args, 'machine'):
    parsed_args.dir = find_vm(parsed_args.machine)
    parsed_args.sock = connect_vm(parsed_args.dir)
    if parsed_args.sock is None:
        if not parsed_args.command in stopped_vm_commands:
            print >>sys.stderr, ("Virtual machine %s is not running." %
                                 parsed_args.machine)
            sys.exit(1)
        else:
            parsed_args.stopped = True

try:
    func(parsed_args)
finally:
    if hasattr(parsed_args, 'sock') and parsed_args.sock is not None:
        parsed_args.sock.shutdown(socket.SHUT_RDWR)
        parsed_args.sock.close()

