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

    """ 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 e:
        if e.errno == errno.ECONNREFUSED:
    except IOError as ex:
        if ex.errno == errno.ECONNREFUSED:
            # virtal machine is not running
            return None
        else:
            raise e
    r, w, x = select.select([sock], [], [], 0.001)
    if sock in r:
        greeting = sock.recv(1024)
            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")

def spiceurl(options):
    """ Returns spice URI for given (as set of parsed args) VM """
    output = send_command(options.sock, "info spice")
    url = None
    for line in output.split("\n"):
        if url is not None:
            continue
        n = line.find("address:")
        if n != -1:
            url = line[n+9:]
        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:
            n = line.find('\t')
            if n <= 0:
            idx = line.find('\t')
            if idx <= 0:
                continue
            name = line[:n]
            name = line[:idx]
            if name == "bridge name":
                continue
            lst.append(name)
    return lst

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):
                m = re.search("file=([^,\\s]*)", line)
                if m:
                    result.append(m.group(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

#
# command implementation
#

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


def cmd_start(options):
    """ vws start """
    if options.stopped:
        arg = ""
        if options.cdrom:
            arg = " -cdrom " + options.cdrom[0]
        if options.snapshot:
            arg = arg+" -snapshot"
        if options.args:

    if options.gui:
        uri = spiceurl(options)
        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:
            r, w, x = select.select([sys.stdin, options.sock], [], [])
            if sys.stdin in r:
            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())
                n = answer.index('\n')
                print answer[n+1:],
                idx = answer.index('\n')
                print answer[idx+1:],
                sys.stdout.flush()
            elif options.sock in r:
            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

        answer = send_command(options.sock, "eject " + options.id)
    else:
        answer = send_command(options.sock, "change %s %s" %
                              (options.id, 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")
    l = []
    lst = []
    for dev in f:
        m = re.match('Bus (\\d+) Device (\\d+): (.*)$', dev)
        if m:
            if m.group(3).endswith("root hub"):
        match = re.match('Bus (\\d+) Device (\\d+): (.*)$', dev)
        if match:
            if match.group(3).endswith("root hub"):
                continue
            l.append((m.group(1) + "." + m.group(2), m.group(3)))
            lst.append((match.group(1) + "." + match.group(2), match.group(3)))
    f.close()
    return l
    return lst

def get_vm_devices(sock):
    """ Parses output of info usb monitor command into list of devices"""
    answer = send_command(sock, "info usb")
    l = []
    lst = []
    for dev in answer.split("\n"):
        m = re.match('Device (\\d+\\.\\d+), .*?, Product (.*)$', dev)
        if m:
            l.append((m.group(1), m.group(2)))
    return l
        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(options):
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 t in get_vm_devices(options.sock):
        print "Address %s : %s" % (t[0], t[1])
    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 = [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

                print "%*s %s" % (-maxlen, f[0], f[1])
            else:
                print f[0]
    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")
    m = re.search('\\[(\\d+)\\]: ', answer)
    if not m:
    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 " + m.group(1))
def cmd_version(options):
        print send_command(options.sock, "stopcapture " + match.group(1))
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))
    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 d in lst:
        print "%-30s  %+8s  %+8s" % (d["image"], d["virtual size"],
                                     d["disk size"])
    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:
            m = re.match("backing.file: (.*)$", line)
            if m:
                return m.group(1)
            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 "
      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))

def cmd_commit(options):
    #
    # Commits last snapshot changes into it's backing file
    """
    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

    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)

template = """#!/bin/sh
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

-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
"""

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
    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(map(lambda x: "%02x" % ord(x),
                           chr(0x52) + os.urandom(5)))
    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 = 0755
    else:
        machinedir = os.path.join(os.environ["HOME"], "VWs",
                                  parsed_args.machine)

        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))

    # 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))
        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=[parsed_args.install],
                               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(cmds, name, **kwargs):
def new_command(cmd_parser, name, **kwargs):
    """
    Adds a subparser and adds a machine name argument to it
    """
    p = cmds.add_parser(name, **kwargs)
    p.add_argument('machine', type=str, help='name of vm to operate on')
    return p
    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"