#!/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()