#!/usr/bin/python
from ConfigParser import ConfigParser
from argparse import ArgumentParser
import fcntl
import socket,select
import errno
import re
import os,sys,time
VERSION=0.1
def find_vm(name):
    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
        if ( name in os.listdir(dirname) and
             os.access(dirname+"/"+name+"/start",os.X_OK)):
            return dirname+"/"+name
    raise ValueError("Machine "+name+" not found.")


def connect_vm(vm_dir):
    sock=socket.socket(socket.AF_UNIX)
    if not os.access(vm_dir+"/monitor",os.W_OK):
        return None
    try:
        sock.connect(vm_dir+"/monitor")
    except IOError as e:
        if e.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)
    return sock

def send_command(sock,command):
    fcntl.flock(sock,fcntl.LOCK_EX)
    try:
        sock.send(command+"\n")
        answer=""
        while not answer.endswith("(qemu) "):
            chunk=sock.recv(1024)
            #print "Got chunk = ",repr(chunk)
            if chunk == '':
                raise IOError("Unexpected EOF From monitor")
            answer+=chunk
    finally:
        fcntl.flock(sock,fcntl.LOCK_UN)
    return answer

def spiceurl(options):
    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:]
            if url.startswith('*:'):
                url="localhost"+url[1:]
    if url is None:
        return None
    return "spice://"+url.rstrip('\r')      

#
# command implementation
#

def cmd_spiceuri(options):
    print spiceurl(options)


def cmd_start(options):
    if options.stopped:
        arg=""
        if options.cdrom:
            arg=" -cdrom "+options.cdrom[0]
        cwd=os.getcwd()
        os.chdir(options.dir)
        os.system("./start%s" % arg)
        os.chdir(cwd)
        time.sleep(2)
        options.sock = connect_vm(options.dir)
    if options.gui:
        uri = spiceurl(options)
        os.system("remote-viewer %s &" % uri)
    elif not options.stopped:
        print >>sys.stderr,"VM already running"
def cmd_stop(options):
    print send_command(options.sock,'system_powerdown')
def cmd_monitor(options):
    try:
        print "(qemu) ",
        sys.stdout.flush()
        while True:
            r,w,x=select.select([sys.stdin,options.sock],[],[])
            if sys.stdin in r:
                cmd = sys.stdin.readline().rstrip()
                answer=send_command(options.sock,cmd)
                n=answer.index('\n')
                print answer[n+1:],
                sys.stdout.flush()
            elif options.sock in r:
                print "UNSOLICITED MESSAGE %"+options.sock.readline().rstrip()
    except KeyboardInterrupt:
        print "Keyboard interrupt"
        sys.exit()
def cmd_reset(options):
    print send_command(options.sock,'system_reset')
def cmd_cdrom(options):
    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, 
                            options.file))
    print answer
def find_usb(options,devices):
    if hasattr("pattern",options):
        pass 
    elif not hasattr("address",options):
        print >>sys.stderr,"Addess or search pattern for device is not specified"
        sys.exit(1)
    else:
        return options.address
def cmd_usb_insert(options):
    address=find_usb(options)
    answer=send_command(options.sock,"usb_add host:%s" % address)
    print answer
def cmd_usb_list(options):
    os.system("lsusb")
def cmd_usb_remove(options):
    address=find_usb(options)
    answer=send_command(options.sock,"usb_del %s" % address)
    print answer
    
def cmd_usb_attached(options):
    answer=send_command(options.sock,"info usb")
    print answer
def cmd_list(options):
    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
        maxlen=0
        vms=[]
        for vmname in os.listdir(dirname):
            if os.access(dirname+"/"+vmname+"/start",os.X_OK):
                count += 1
                f=[vmname]
                if maxlen < len(vmname):
                    maxlen = len(vmname)
                if options.state:
                    sock = connect_vm(dirname+"/"+vmname)
                    if sock is None:
                        state = "stopped"
                    else:
                        sock.close()
                        state = "running"
                    f.append(state)
                vms.append(f)
        for f in sorted(vms):
            if len(f)==2:
                print "%*s %s" % (-maxlen,f[0],f[1])
            else:
                print f[0]
    if not count:
        sys.exit(1)
def cmd_version(options):
    print VERSION

def validate_size(size):
    return re.match('\d+[KMG]',size) is not None
def cmd_create(parsed_args):
    if not parsed_args.image and not validate_size(parsed_args.size):
        print >>sys.stderr,"Invalid size of disk specifed %s. Should have K, M or G suffix"%parsed_args.size
        sys.exit(1)
    if not validate_size(parsed_args.mem):    
        print >>sys.stderr,"Invalid size of memory specifed %s. Should have K, M or G suffix"%parsed_args.size
        sys.exit(1)
    libdir="/usr/local/lib/vws"
    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)))

    if parsed_args.shared:
        machinedir = config.get("directories","SharedVMs")+"/"+parsed_args.name
        dirmode = 0755
    else:
        machinedir = os.environ["HOME"]+"/VWs/"+parsed_args.name
        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
    if not parsed_args.arch in ('i386','x86_64'):
        print >>sys.stderr,"KVM acceleration disabled due to target architecture"
        options.accel = ''
    elif not os.access("/dev/kvm",os.W_OK):
        print >>sys.stderr,"KVM acceleration disabled due to 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


    #  Creating directory for VM
    os.makedirs(machinedir,dirmode)
    driveopts={"interface":parsed_args.diskif,"image":drivename}
    if parsed_args.image:
        # Copying image file
        print >>sys.stderr,"Copying %s to %s"%(parsed_args.image,machinedir+"/"+drivename) 
        os.system("qemu-img convert -O qcow2 -p %s %s/%s"%(parsed_args.image,machinedir+"/"+drivename))
        os.chdir(machinedir)
    else:
        print >>sys.stderr,"Creating new image file"
        os.chdir(machinedir)
        os.system("qemu-img create  -f qcow2 %s %s"%(drivename,parsed_args.size))

    options["drive"]=options["drive"].format(**driveopts) 
    if parsed_args.debug:
        print repr(driveopts),repr(options["drive"])
        print repr(options)
    with open(libdir+"/start.template","r") as f:
        template=f.read()
    with open("start","w") as script:
        script.write(template.format(**options))

    os.chmod('start',0755)

    # If installation media is specified vws start for new vm
    if parsed_args.install:
        start_opts=Namespace(command='start',cdrom=parsed_args.install,stopped=True)
        cmd_start(start_opts)




#
# Utility functions for arg parsing
#
def new_command(cmds,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
#
# arg parsing
#



config=ConfigParser({'SharedVMs':'/var/cache/vws/shared',  
    'AutoStartVMs':'/var/cache/vws/autostart'})
config.add_section('directories')
config.read(['/etc/vws.conf',os.environ['HOME']+'/.vwsrc'])

args=ArgumentParser()
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("--addr",action='store_const',const=True,default=False,
                dest='addr',help='Show mac address and spice port')
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")
# 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')
# Following commands don't need extra args
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=new_command(cmds,'revert',help='Revert to last snapshot')
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,'screenshoot',help='take a screenshot')
p.add_argument('filename',help='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')
# Create new VM
p=new_command(cmds,'create',help="Create new VM")
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="20G");
p.add_argument("--arch",metavar='cputype',help="Emulated architecture",dest="arch",default='x86_64');
p.add_argument("--no-sound",help="Disable sound card",action='store_const',const = None,default='hda',  dest="sound")
p.add_argument("--sound",metavar='cardtype',help="Specify sound card type",dest='sound',default='hda')
p.add_argument("--vga",metavar='cardtype',help="specify video card type (cirrus,std,vmwae,qxl) default qxl",dest="vga",default="qxl")
p.add_argument("--net",help="Network - 'user' or bridge name",dest='net',default="user")
p.add_argument("--mem",metavar='size',help="Size of memory",dest="mem",default="1024M")
p.add_argument("--diskif",metavar='interface-type',help="Disk interface (virtio, scsi, ide)",choices=['virtio','scsi','ide'],dest="diskif",default="virtio")
p.add_argument('--shared',help='Create shared VM instead of private one',action='store_const',const= True,dest='shared',default=False)
p.add_argument('--image',metavar='filename',help='Existing disk image to import',dest='image',default=None)
p.add_argument('--install',metavar='filename.iso',help='ISO image to install OS from',dest='install',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:])

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

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

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)
func(parsed_args)
