#!/bin/python

 import os

 import time

 import dbus

 import sys

 import subprocess

 from qubes.qubes import QubesVmCollection

 from qubes.qubes import QubesDispVmLabels

 from qubes.qubes import QubesException

 from qubes.qubes import QubesVmLabels

 from qubes.qmemman_client import QMemmanClient

 from qubes.qubesutils import block_list,block_attach,block_detach,block_detach_all,block_check_attached

 from qubes.qubesutils import kbytes_to_kmg, bytes_to_kmg

 notify_object = None

 current_savefile = '/var/run/qubes/current_savefile'

 current_dvm_conf = '/var/run/qubes/current_dvm.conf'

 avir_appvm_template = 'avir-template'

 avir_dvm_template = 'avir-template-dvm'

 def get_vm_collection():

 	vm_collection = QubesVmCollection()

 	vm_collection.lock_db_for_reading()

 	vm_collection.load()

 	vm_collection.unlock_db()

 	return vm_collection

 def get_vm_by_name(vmname):

 	return get_vm_collection().get_vm_by_name(vmname)

 def gen_scanner_name():

 	vm_collection = get_vm_collection()

 	vms = [vm for vm in vm_collection.values()]

 	vm_names = []

 	for vm in vms:

 		if vm.is_appvm() and vm.template.name == avir_appvm_template:

 			vm_names.append (vm.name)

 	print vm_names

 	ids = range(100)

 	for id in ids:

 		if "scanner"+str(id) not in vm_names:

 			return "scanner"+str(id)

 	print >> sys.stderr, "Unable to generate new scanner vm name"

 	exit(1)

 class VMFactory:

 	def __init__(self):

 		self.tray_notify(str="Initialized")

 	def tray_notify(self, str, timeout = 3000):

 		notify_object.Notify("Qubes", 0, "red", "Qubes", str, [], [], timeout, dbus_interface="org.freedesktop.Notifications")

 	def tray_notify_error(self, str, timeout = 3000):

 		notify_object.Notify("Qubes", 0, "dialog-error", "Qubes", str, [], [], timeout, dbus_interface="org.freedesktop.Notifications")

 	def create_appvm(self):

 		#generate new avir-scanner name

 		scanner_vmname = gen_scanner_name()

 		print "Generating new scanner vm %s " % scanner_vmname

 		#create new vscanner virtual machine from template "avir-template"

 		qvm_collection = QubesVmCollection()

 		qvm_collection.lock_db_for_writing()

 		qvm_collection.load()

 		avir_template = qvm_collection.get_vm_by_name(avir_appvm_template)

 		if not avir_template.is_template():

 			print >> sys.stderr, "VM '{0}' is not a TemplateVM".format(avir_template)

 			exit (1)

 		vmtype = "QubesAppVm"

 		try:

 			scanner_vm = qvm_collection.add_new_vm(vmtype, name=scanner_vmname, template=avir_template, label = QubesVmLabels["orange"])

 		except QubesException as err:

 			print >> sys.stderr, "ERROR: {0}".format(err)

 			exit (1)

 		scanner_vm.memory = 512

 		scanner_vm.vcpus = 2

 		try:

 			scanner_vm.create_on_disk(verbose=True, source_template=avir_template)

 			#if options.root:

 			#       os.unlink(vm.root_img)

 			#       os.rename(options.root, vm.root_img)

 		except (IOError, OSError) as err:

 			print >> sys.stderr, "ERROR: {0}".format(err)

 			exit (1)

 		qvm_collection.save()

 		qvm_collection.unlock_db()

 		return scanner_vmname

 	def start_vm(self, vmname):

 		vm = get_vm_by_name(vmname)

 		if vm is None or vm.qid not in get_vm_collection():

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		try:

 			vm.verify_files()

 			xid = vm.start(verbose=True, preparing_dvm=False, start_guid=True, notify_function=None)

 		except (IOError, OSError, QubesException, MemoryError) as err:

 			#if options.tray:

 			#	tray_notify_error(str(err))

 			#else:

 			print >> sys.stderr, "ERROR: {0}".format(err)

 			exit (1)

 	def create_dvm(self):

 		qmemman_client = QMemmanClient()

 		if not qmemman_client.request_memory(400*1024*1024):

 			qmemman_client.close()

 			errmsg = 'Not enough memory to create DVM. '

 			errmsg +='Terminate some appVM and retry.'

 			if os.path.exists('/usr/bin/kdialog'):

 				subprocess.call(['/usr/bin/kdialog', '--sorry', errmsg])

 			else:

 				subprocess.call(['/usr/bin/zenity', '--warning', errmsg])

 			return None

 		self.tray_notify("Starting new DispVM...")

 		qvm_collection = QubesVmCollection()

 		qvm_collection.lock_db_for_writing()

 		qvm_collection.load()

 		avir_template = qvm_collection.get_vm_by_name(avir_dvm_template)

 		if avir_template is None:

 			sys.stderr.write( 'Domain ' + avir_dvm_template + ' does not exist ?')

 			qvm_collection.unlock_db()

 			qmemman_client.close()

 			return None

 		label = avir_template.label

 		retcode = subprocess.call(['/usr/lib/qubes/qubes_restore',

 			current_savefile,

 			current_dvm_conf,

 			'-u', str(avir_template.default_user),

 			'-c', label.color,

 			'-i', label.icon_path,

 			'-l', str(label.index)])

 		qmemman_client.close()

 		if retcode != 0:

 			if os.path.exists('/usr/bin/kdialog'):

 				subprocess.call(['/usr/bin/kdialog', '--sorry', 'DisposableVM creation failed, see qubes_restore.log'])

 			else:

 				subprocess.call(['/usr/bin/zenity', '--warning', 'DisposableVM creation failed, see qubes_restore.log'])

 			qvm_collection.unlock_db()

 			return None

 		f = open('/var/run/qubes/dispVM_xid', 'r');

 		disp_xid = f.readline().rstrip('\n')

 		disp_name = f.readline().rstrip('\n')

 		disptempl = f.readline().rstrip('\n')

 		f.close()

 		print >>sys.stderr, "getting vm %s" % disptempl	

 		avir_disptempl = qvm_collection.get_vm_by_name(disptempl);

 		if avir_disptempl is None:

 			sys.stderr.write( 'Domain ' + disptempl + ' does not exist ?')

 			qvm_collection.unlock_db()

 			return None

 		dispid = int(disp_name[4:])

 		dispvm = qvm_collection.add_new_disposablevm(disp_name, avir_disptempl.template, label=label, dispid=dispid, netvm=avir_disptempl.netvm)

 		print >> sys.stderr, "Waiting for the VM %s to startup...", dispvm.name

 		startup_counter = 0

 		startup_counter_max = 60

 		while startup_counter < startup_counter_max:

 			if dispvm.is_running():

 				break

 			shutdown_counter += 1

 			time.sleep (1)

 		qvm_collection.save()

 		qvm_collection.unlock_db()

 		return dispvm.name

 	def stop_vm(self, vmname):

 		vm = get_vm_by_name(vmname)

 		if vm is None or vm.qid not in get_vm_collection():

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		vm.shutdown(force=True)

 		print >> sys.stderr, "Waiting for the VM %s to shutdown...", vmname

 		shutdown_counter = 0

 		shutdown_counter_max = 60

 		while shutdown_counter < shutdown_counter_max:

 			if not vm.is_running():

 				return

 			shutdown_counter += 1

 			time.sleep (1)

 	def remove_vm(self, vmname):

 		qvm_collection = QubesVmCollection()

 		qvm_collection.lock_db_for_writing()

 		qvm_collection.load()

 		vm = qvm_collection.get_vm_by_name(vmname)

 		if vm is None or vm.qid not in qvm_collection:

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		if vm.is_running():

 			print >> sys.stderr, "Cannot remove a running VM, stop it first"

 			exit (1)

 		if vm.is_template():

 			dependent_vms = qvm_collection.get_vms_based_on(vm.qid)

 			if len(dependent_vms) > 0:

 				print >> sys.stderr, "The following AppVMs use '{0}' as a template:".format(vmname)

 				for vm in dependent_vms:

 					print >> sys.stderr, "{name:<12} (qid={qid})".format(qid=vm.qid, name=vm.name)

 				print >> sys.stderr, "Please remove those VMs first."

 				exit (1)

 			if qvm_collection.default_template_qid == vm.qid:

 				qvm_collection.default_template_qid = None

 		try:

 			vm.remove_from_disk()

 		except (IOError, OSError) as err:

 			print >> sys.stderr, "Warning: {0}".format(err)

 		# Do not exit, perhaps the VM files were somehow removed

 		# so just remove it from Qubes DB

 		qvm_collection.pop(vm.qid)

 		qvm_collection.save()

 		qvm_collection.unlock_db()

 	def attach_dev(self, vmname, devname):

 		vm = get_vm_by_name(vmname)

 		if vm is None or vm.qid not in get_vm_collection():

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		print "Attaching %s to vm - %s" % (devname, vm.name)

 		kwargs = {}

 		kwargs['auto_detach'] = True

 		try:

 			backend_vm = get_vm_by_name("dom0")

 			block_attach(vm, backend_vm, devname, **kwargs)

 		except QubesException as e:

 			print >> sys.stderr, "ERROR: %s" % str(e)

 			exit(1)

 	def detach_dev(self, vmname, devname):

 		vm = get_vm_by_name(vmname)

 		if vm is None or vm.qid not in get_vm_collection():

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		print "Detaching %s from vm - %s" % (devname, vmname)

 		kwargs = {}

 		#kwargs['frontend'] = "to_scan"

 		try:

 			block_detach(vm, **kwargs)

 		except QubesException as e:

 			print >> sys.stderr, "ERROR: %s" % str(e)

 			sys.exit(1)

 	def scan(self, vmname):

 		vm = get_vm_by_name(vmname)

 		if vm is None or vm.qid not in get_vm_collection():

 			print >> sys.stderr, "A VM with the name '{0}' does not exist in the system.".format(vmname)

 			exit(1)

 		subprocess.call(['/usr/lib/qubes/qrexec_client', '-d', vmname, 'root:mount /dev/xvdi /mnt/removable'])

 		#subprocess.call(['/usr/lib/qubes/qrexec_client', '-d', vmname, 'root:clamscan /mnt/removable'])

 		subprocess.call(['/usr/lib/qubes/qrexec_client', '-d', vmname, 'root:umount /mnt/removable'])

 def main():

 	global notify_object

 	print >>sys.stderr, "time=%s, qfile-daemon-dvm init" % (str(time.time()))

 	notify_object = dbus.SessionBus().get_object("org.freedesktop.Notifications", "/org/freedesktop/Notifications")

 	vm_factory = VMFactory()

 	#vmname = vm_factory.create_dvm()

 	vmname = vm_factory.create_appvm()

 	vm_factory.start_vm(vmname)

 	vm_factory.attach_dev(vmname, "sdb")

 	vm_factory.scan(vmname)

 	vm_factory.detach_dev(vmname, "sdb")

 	#vm_factory.stop_vm(vmname)

 	#vm_factory.remove_vm(vmname) 

 main()