#!/usr/bin/env python

# This file is part of Window-Switch.
# Copyright (c) 2009-2013 Antoine Martin <antoine@nagafix.co.uk>
# Window-Switch is released under the terms of the GNU GPL v3

import os.path
import time
import re

from winswitch.consts import VIRTUALBOX_TYPE, VIRTUALBOX_PORT_BASE, MAGIC_NO_PASSWORD
from winswitch.globals import USERNAME
from winswitch.objects.session import Session
from winswitch.objects.server_session import ServerSession
from winswitch.virt.server_util_base import ServerUtilBase
from winswitch.util.common import is_valid_exe, visible_command, csv_list, no_newlines, is_valid_dir
from winswitch.util.process_util import twisted_exec
from winswitch.util.main_loop import callLater, callFromThread


COMMON_SESSION_NAMES = {"mac os x server" : "osx", "macos" : "osx",
						"windows vista" : "win32", "windowsvista" : "windows", "windows 7" : "win32", "windows7" : "win32",
						"windowsxp" : "win32", "windows xp" : "win32", "windows" : "win32", "vista" : "win32", "xp" : "win32",
						"opensuse" : "opensuse",
						"centos" : "centos", "red hat" : "redhat", "redhat" : "redhat",
						"debian" : "debian", "wheezy" : "debian", "lenny" : "debian", "squeeze" : "debian", "sid" : "debian",
						"ubuntu" : "ubuntu", "hardy" : "ubuntu", "karmic" : "ubuntu", "lucid" : "ubuntu", "maverick" : "ubuntu", "natty" : "ubuntu", "oneiric" : "ubuntu", "precise" : "ubuntu",
						"bsd" : "bsd", "freebsd" : "bsd", "pcbsd" : "bsd", "netbsd" : "bsd", "openbsd" : "openbsd",
						"opensolaris" : "opensolaris", "solaris" : "solaris",
						}

USE_PASSWORD = True
TLS_PROPERTY = "RDP"
# = None			#leave unchanged
# = "negotiate"		#enable TLS if supported
# = "RDP"			#only Standard RDP Security is accepted.
# = "TLS"			#only Enhanced RDP Security is accepted. The client must support TLS.


class	VirtualboxServerUtil(ServerUtilBase):

	def	__init__(self, config, add_session, remove_session, update_session_status, session_failed):
		ServerUtilBase.__init__(self, VIRTUALBOX_TYPE, VIRTUALBOX_PORT_BASE, config, add_session, remove_session, update_session_status, session_failed)
		self.connecting_timeout = 20			#20 seconds to go from connecting to connected
		self.prelaunch_enabled = False			#preload broken at present..
		dotvb = os.path.expanduser("~/.VirtualBox")
		if is_valid_dir(dotvb):
			self.watched_directories.append(dotvb)
		self.batch_detect_sessions_delay = 2	#wait 2 seconds

	def get_config_options(self):
		return	self.get_config_options_base(detect=True)

	def can_client_set_status(self, session, user_id, _from, _to):
		return	ServerUtilBase.can_client_set_status(self, session, user_id, _from, _to) \
			or (_from==Session.STATUS_AVAILABLE and _to==Session.STATUS_CONNECTED) \
			or (_from==Session.STATUS_CONNECTING and _to==Session.STATUS_CONNECTED) \
			or (_from==Session.STATUS_CONNECTED and _to==Session.STATUS_AVAILABLE) \
			or (_from==Session.STATUS_IDLE and _to==Session.STATUS_AVAILABLE)

	def get_test_port(self, session):
		return	session.port

	def watch_existing_log(self, session):
		#not used with vbox
		pass

	def	do_prepare_session_for_attach(self, session, user, disconnect, call_when_done):
		""" this is the main entry point """
		self.slog(None, session, user, disconnect, call_when_done)
		def session_ready(msg=None):
			self.sdebug("will call %s" % call_when_done, msg)
			call_when_done()

		def err_cb(err):
			self.serror(None, err)
			msg = visible_command(err).replace("\\'", "'")
			if err and err.find("WARNING: The vboxdrv kernel module is not loaded.")>=0:
				msg = "VirtualBox is not setup correctly, please ensure the vboxdrv kernel module is loaded"
			self.early_failure(session, "Failed to setup the session: %s" % msg, Session.STATUS_UNAVAILABLE)
		def session_updated(updated_session):
			######## what state are we in?
			if session.status==Session.STATUS_SUSPENDED:
				self.start_headless(session, user, session_ready, err_cb)
			else:
				self.configure_headless_session(session, user, session_ready, err_cb)
		self.detect_existing_session(session.ID, session, callback=session_updated)

	def configure_headless_session(self, session, user, ok_cb, err_cb):
		""" sets the new password and calls configure_vrde """
		#first we need a password:
		def next_step(*args):
			self.configure_vrde(session, user, ok_cb, err_cb)
		if not USE_PASSWORD and not session.password:
			session.password = MAGIC_NO_PASSWORD
		if session.password:
			next_step()
			return
		def got_password_hash(password_hash):
			self.slog(None, password_hash)
			s = no_newlines(password_hash)
			HASH_HEADER = "Password hash: "
			if not s.startswith(HASH_HEADER):
				err_cb("failed to hash the password")
				return
			encrypted_password = password_hash[len(HASH_HEADER):].strip()
			if session.encrypted_password != encrypted_password:
				session.encrypted_password = encrypted_password
				#now we need to save it:
				cmd = [self.config.vboxmanage_command, "setextradata", session.ID, "VBoxAuthSimple/users/%s" % session.user, session.encrypted_password]
				twisted_exec(cmd, ok_callback=next_step, err_callback=err_cb)
			else:
				next_step("password already configured")
		session.password = self.new_password()[:32]
		cmd = [self.config.vboxmanage_command, "internalcommands", "passwordhash", session.password]
		twisted_exec(cmd, got_password_hash, err_callback=err_cb)

	def configure_vrde(self, session, user, ok_cb, err_cb):
		""" sets all the session options
			(if they need setting - and update the session object, propagate changes to the clients)
			and calls ok_cb() """
		new_options = {}
		port = session.port
		if port is None or port<=0:
			port = self.get_free_port()
			#send the updated session to the clients:
			new_options["vrdeport"] = port
		if not session.vrde:
			new_options["vrde"] = "on"
		if USE_PASSWORD:
			if session.vrdeauthtype!="external":
				new_options["vrdeauthtype"] = "external"
			vrdeauthtype = "external"
			if session.vrdeauthlibrary!="VBoxAuthSimple":
				new_options["vrdeauthlibrary"] = "VBoxAuthSimple"
			vrdeauthlibrary = "VBoxAuthSimple"
		else:
			if session.vrdeauthtype!="null":
				new_options["vrdeauthtype"] = "null"
			vrdeauthtype = "null"
			vrdeauthlibrary = ""
		if session.host!=session.vrdeaddress:
			new_options["vrdeaddress"] = session.host
		#if SUSPENDED, we will be the one starting it and this fails if 3D is enabled!?
		if session.status==Session.STATUS_SUSPENDED and session.accelerate3d:
			new_options["accelerate3d"] = "off"

		if len(new_options)==0:
			#nothing to do!
			ok_cb("ready: already configured!")
			return

		#these are optional but since we are configuring, set them:
		if TLS_PROPERTY:
			new_options["vrdeproperty"] = "Security/Method=%s" % TLS_PROPERTY
		if not session.vrdemulticon:
			new_options["vrdemulticon"] = "on"						#support concurrent clients

		self.sdebug("setting new options: %s" % new_options, session, user, ok_cb, err_cb)
		cmd = [self.config.vboxmanage_command, "modifyvm", session.ID]
		for k,v in new_options.items():
			cmd.append("--%s" % k)
			cmd.append(str(v))
		def configured(msg):
			if port!=session.port:
				session.port = port
				self.sdebug("new session port=%s" % session.port, msg)
			else:
				self.sdebug(None, msg)
			session.vrde = True
			session.vrdeauthtype = vrdeauthtype
			session.vrdeauthlibrary = vrdeauthlibrary
			session.vrdemulticon = True
			self.add_session(session)	#update clients
			ok_cb("instance configured")
		twisted_exec(cmd, configured, err_callback=err_cb)

	def start_headless(self, session, user, ok_cb=None, err_cb=None):
		""" calls configure_headless_session then start_daemon """
		self.sdebug(None, session, user, ok_cb, err_cb)
		def configure_failed(err):
			self.sdebug(None, err)
			if err_cb:
				err_cb(err)
		def configured_ok(*args):
			self.slog(None, *args)
			#now we can start the VBoxHeadless process:
			cmd = [self.config.vboxheadless_command, "--startvm", session.ID]
			env = session.get_env()
			del env["DISPLAY"]
			session.start_time = int(time.time())
			self.update_session_status(session, Session.STATUS_STARTING)
			if self.start_daemon(session, cmd, env):
				ok_cb("session started")
			else:
				err_cb("failed to start!")
		self.configure_headless_session(session, user, configured_ok, configure_failed)

	def session_process_ended(self, pid, condition, session):
		""" override so we set the session state to SUSPENDED (not CLOSED) """
		self.do_session_process_ended(pid, condition, session, Session.STATUS_SUSPENDED)

	def process_log_line(self, session, line):
		VERSION_LINE = "Oracle VM VirtualBox Headless Interface "
		if line.startswith(VERSION_LINE):
			vbox_version = no_newlines(line[len(VERSION_LINE):])
			self.slog("identified VirtualBox version %s" % vbox_version)
			return
		START_ERROR_LINE = "Error: failed to start machine."
		if line.startswith(START_ERROR_LINE):
			self.serror(None, session, line)
			callFromThread(self.early_failure, session, "VirtualBox Error starting %s: %s" % (session.name, line[len(START_ERROR_LINE):].strip()), Session.STATUS_UNAVAILABLE)
			return
		self.sdebug(None, session, line)




	def start_display(self, session, user, is_preload):
		raise Exception("virtualbox cannot be used to start new sessions!")

	def stop_display(self, session, user, display):
		"""
		Send shutdown signal and if that fails, kill it..
		"""
		self.slog(None, session, user, display)
		start_update_count = session.status_update_count
		def force_poweroff():
			def poweroff_ok(msg):
				self.slog(None, msg)
				self.detect_existing_session(session.ID, session)
			def poweroff_error(msg):
				self.serror(None, msg)
				self.detect_existing_session(session.ID, session)
			cmd = [self.config.vboxmanage_command, "controlvm", session.ID, "poweroff"]
			twisted_exec(cmd, poweroff_ok, err_callback=poweroff_error)
		def acpipowerbutton_ok(msg):
			self.sdebug(None, msg)
			#check in 30s and poweroff if still running:
			def check_power_off(*args):
				self.sdebug("session=%s, status_update_count=%s, initial=%s" % (session, session.status_update_count, start_update_count), *args)
				if session.status_update_count==start_update_count:
					def session_updated(new_session):
						self.sdebug("session=%s, status_update_count=%s, initial=%s" % (new_session, new_session.status_update_count, start_update_count), *args)
						if new_session.status_update_count==start_update_count:
							force_poweroff()
					self.detect_existing_session(session.ID, session, callback=session_updated)
			callLater(30, check_power_off)
		def acpipowerbutton_error(msg):
			self.serror(None, msg)
			force_poweroff()
		cmd = [self.config.vboxmanage_command, "controlvm", session.ID, "acpipowerbutton"]
		twisted_exec(cmd, acpipowerbutton_ok, err_callback=acpipowerbutton_error)




	def can_capture(self, session):
		return True

	def can_capture_now(self, session):
		if	(time.time()-session.start_time)<10:
			return	False
		return	session.status in [Session.STATUS_CONNECTED, Session.STATUS_CONNECTING]

	def take_screenshot(self, ID, display, env, filename, ok_callback, err_callback):
		cmd = [self.config.vboxmanage_command, "controlvm", ID, "screenshotpng", filename]
		from winswitch.ui.capture_util import ExecCaptureDisplay
		c = ExecCaptureDisplay(cmd, display, env, filename, ok_callback, err_callback)
		c.exec_capture(cmd, wait_for_file=5)





	def detect_sessions(self):
		""" We are here because something has changed, figure out what:
			Go through all the existing sessions and update their status.
			If we find a new session, add it.
		"""
		self.sdebug()
		sessions = {}
		for session in self.config.get_sessions_by_type(self.session_type):
			if session.is_status_final():
				continue		#cant be updated
			self.detect_existing_session(session.ID, session)
			sessions[session.ID] = session
		def err_callback(err):
			self.serror()
		def got_sessions(vbox_sessions):
			for uuid,_ in vbox_sessions.items():
				if uuid in sessions:
					continue		#already handled
				self.detect_existing_session(uuid, None, callback=self.found_live_session)
		self.with_session_list(got_sessions, err_callback)

	def detect_existing_sessions(self, reloaded_sessions):
		""" given a list of existing sessions, either:
			* call session_cleanup if don't find it in the vbox session list
			* call detect_existing_session otherwise
			We get the vbox session list via with_session_list()
		"""
		def cleanup(*args):
			self.slog("cleanup(%s) ignoring the following sessions: %s" % (args, csv_list(reloaded_sessions)))
			for session in reloaded_sessions:
				self.session_cleanup(session)

		def found_vbox_sessions(vbox_sessions):
			found_sessions = {}
			for uuid,_ in vbox_sessions.items():
				found_sessions[uuid] = None
				for session in reloaded_sessions:
					if session.ID==uuid:
						found_sessions[uuid] = session
						reloaded_sessions.remove(session)
			#all the ones left have not been found:
			cleanup()
			#all the ones found must now be detected:
			for uuid,session in found_sessions.items():
				self.detect_existing_session(uuid, session, callback=self.found_live_session)

		self.with_session_list(found_vbox_sessions, err_callback=cleanup)

	def with_session_list(self, ok_callback, err_callback):
		""" executes "vboxmanage list vms
			and calls ok_callback({uuid:name}) or err_callback(msg)
			once the output has been received and parsed
		"""
		if not is_valid_exe(self.config.vboxmanage_command):
			err_callback("no VBoxManage command!")
			return

		def parse_manage_list_vms(output):
			self.slog(None, output)
			#"OSX 10.6.8" {9e78d384-f697-4c86-ba0d-0475439218e3}
			#"..."
			vm_name_id_re = re.compile("\"(.*)\" \{([0-9abcdefABCDEF\-]*)\}")
			vbox_sessions = {}
			for vm in output.splitlines():
				m = vm_name_id_re.match(vm)
				if not m:
					self.serror("unrecognized vm entry: %s" % vm, visible_command(output))
					continue
				name = m.group(1)
				uuid = m.group(2)
				if not uuid or not name:
					continue
				vbox_sessions[uuid] = name
			ok_callback(vbox_sessions)

		cmd = [self.config.vboxmanage_command, "list", "vms"]
		twisted_exec(cmd, parse_manage_list_vms, err_callback=err_callback)


	def detect_existing_session(self, uuid, session, callback=None):
		def props_err(msg):
			if session:
				self.session_cleanup(session)
		def props_ok(props):
			new_session = self.process_session_properties(uuid, session, props)
			if new_session and callback:
				callback(new_session)
		self.with_session_properties(uuid, props_ok, props_err)

	def with_session_properties(self, uuid, ok_callback=None, error_callback=None, parse_extra=True):
		assert uuid
		cmd = [self.config.vboxmanage_command, "showvminfo", "--machinereadable", uuid]
		def vminfo_error(err):
			if error_callback:
				error_callback(err)
		def vminfo_ok(vminfo_output):
			self.sdebug(None, visible_command(vminfo_output))
			def parse(extra=None):
				self.sdebug(None, visible_command(extra))
				props = self.parse_vminfo(vminfo_output, extra)
				if ok_callback:
					ok_callback(props)
			if not parse_extra:
				parse()
				return
			def extra_error(err):
				self.serror(None, err)
				parse()
			def extra_ok(extradata_output):
				parse(extradata_output)
			cmd = [self.config.vboxmanage_command, "getextradata", uuid, "enumerate"]
			twisted_exec(cmd, extra_ok, err_callback=extra_error)
		twisted_exec(cmd, vminfo_ok, err_callback=vminfo_error)

	def parse_vminfo(self, vminfo_output, extradata_output):
		def noquotes(s):
			if not s or len(s)<2:
				return	s
			if (s[0]=='"' and s[-1]=='"') or (s[0]=="'" and s[-1]=="'"):
				return	s[1:-1]
			return s
		def parseval(s):
			v = noquotes(s)
			if v=="on":
				return	True
			if v=="off":
				return	False
			if v=="none":
				return	None
			try:
				return	int(v)
			except:
				return v
		#first parse vminfo:
		props = {}
		for line in vminfo_output.splitlines():
			parts = line.split("=", 1)
			if len(parts)!=2:
				self.serror("unable to parse line: %s" % line, visible_command(vminfo_output), visible_command(extradata_output))
				continue
			k = noquotes(parts[0])
			v = parseval(parts[1])
			props[k] = v
		#parse extra:
		if extradata_output:
			KEY="Key: "
			VALUE=" Value: "
			for line in extradata_output.splitlines():
				if not line:
					continue
				npos = line.find(KEY)
				vpos = line.find(VALUE)
				if npos<0 or vpos<0:
					self.sdebug("ignored invalid extradata line: %s" % line, visible_command(vminfo_output), visible_command(extradata_output))
					continue
				k = line[(npos+len(KEY)):vpos].strip()
				v = line[vpos+len(VALUE):].strip()
				props[k] = v
		return props

	def process_session_properties(self, uuid, session, props):
		#self.sdebug(None, uuid, session, props)
		if ("name" not in props) or ("UUID" not in props):
			self.serror("cannot find name of session in vminfo!", uuid, session, props)
			if session:
				self.session_cleanup(session)
			return	None
		if session:
			if props.get("UUID")!=session.ID:
				self.serror("session UUID=%s does not match!" % props.get("UUID"), uuid, session, props)
				self.session_cleanup(session)
				return	None
			self.slog("re-using existing session object", uuid, session, visible_command(str(props)))
		else:
			name = props.get("name")
			self.slog("creating new session object for %s" % name, uuid, session, visible_command(str(props)))
			session = ServerSession(self.config)
			session.ID = uuid
			session.name = name
			session.password = ""
			session.user = USERNAME
			session.display = "VBOX-"+name.replace(" ", "_")
			(_, tunnel) = self.get_session_bind_address()
			session.requires_tunnel = tunnel
			session.session_type = self.session_type
			session.start_time = 0
			session.env = session.get_env()
			session.owner = self.config.local_session_owner
			session.command = "VirtualBox"
			session.preload = False
			session.full_desktop = True
			session.command_uuid = uuid
			session.can_export_sound = False
			session.uses_sound_out = False
			session.uses_sound_in = False
			ostype = props.get("GuestOSType") or props.get("ostype", "")
			osname = ostype.lower().replace("-", " ").replace("_64", "").replace("_", " ").replace("(64 bit)", "").strip()
			#convert the ostype into a common icon name
			icon_name = COMMON_SESSION_NAMES.get(osname)
			icon = None
			from winswitch.util.icon_cache import guess_icon_from_name
			if icon_name:
				icon = guess_icon_from_name(icon_name)
			if not icon:
				for x in osname.split(" "):
					icon_name = COMMON_SESSION_NAMES.get(x)
					if icon_name:
						icon = guess_icon_from_name(icon_name)
						if icon:
							break
				if not icon:
					for x in osname.split(" "):
						icon = guess_icon_from_name(x)
						if icon:
							break
			if icon:
				session.set_default_icon_data(icon)
			#guess resolution:
			mode = props.get("/VirtualBox/GuestAdd/Vbgl/Video/SavedMode")
			screen_size = None
			if mode:
				try:
					#ie: 744x480x32
					parts = mode.split("x")
					if len(parts)==3:
						screen_size = "%sx%s-%s" % (int(parts[0]), int(parts[1]), int(parts[2]))
				except:
					pass
			if not screen_size:
				mode = props.get("GUI/LastGuestSizeHint")
				if mode:
					try:
						#ie: 720,400
						parts = mode.split(",")
						if len(parts)==2:
							screen_size = "%sx%s" % (int(parts[0]), int(parts[1]))
					except:
						pass
			if screen_size:
				session.screen_size = screen_size
			else:
				session.screen_size = "1024x768"
		#parse rdp port info - aka vrde:
		#vrde="on"
		#vrdeauthtype="external"
		try:
			session.port = int(props.get("vrdeport"))
		except:
			pass
		if session.port is None or session.port<=0:
			try:
				session.port = int(props.get("vrdeports"))
			except:
				pass
		if session.port>0:
			#ensure we don't end up locking this port!
			self.sdebug("found valid vrde port: %s" % session.port, uuid, session, visible_command(str(props)))
			self.port_mapper.add_taken_port(session.port)
		(host, _) = self.get_session_bind_address()
		session.host = host
		session.vrdeaddress = props.get("vrdeaddress", "")
		session.vrde = props.get("vrde", False)
		session.vrdereusecon = props.get("vrdereusecon", False)
		session.vrdemulticon = props.get("vrdemulticon", False)
		session.vrdeauthtype = props.get("vrdeauthtype", "")
		session.vrdeauthlibrary = props.get("vrdeauthlibrary", "")
		session.accelerate3d = props.get("accelerate3d", False)
		if not session.vrdeauthlibrary and session.vrdeauthtype=="external":
			#the property is not returned to us??
			session.vrdeauthlibrary = "VBoxAuthSimple"

		vmstate = props.get("VMState", "")
		status = self.VMState_to_status(session, vmstate)
		if status:
			self.sdebug("current status=%s, new status=%s" % (session.status, status), uuid, session, visible_command(str(props)))
			if session.status in [Session.STATUS_CONNECTING, Session.STATUS_CONNECTED] and status==Session.STATUS_AVAILABLE:
				#we set connected/connecting from the client side, so don't reset it here
				pass
			else:
				session.set_status(status)
		self.save_session(session)
		return session

	def VMState_to_status(self, session, vmstate):
		if not vmstate:
			return	None
		if vmstate=="running":
			#we can't change vrde settings on a running session:
			if not session.vrde:
				self.serror("this session is running with vrde off - therefore unavailable", session, vmstate)
				return	Session.STATUS_UNAVAILABLE
			if USE_PASSWORD:
				if session.vrdeauthtype!="external":
					self.serror("this session is running with vrdeauthtype=%s - therefore unavailable" % session.vrdeauthtype, session, vmstate)
					return	Session.STATUS_UNAVAILABLE
			else:
				if session.vrdeauthtype!="null":
					self.serror("this session is running with vrdeauthtype=%s - therefore unavailable" % session.vrdeauthtype, session, vmstate)
					return	Session.STATUS_UNAVAILABLE

			if session.port is None or session.port<=0:
				self.serror("this session is running without a valid vrde port: %s - therefore unavailable" % session.port, session, vmstate)
				return	Session.STATUS_UNAVAILABLE
		return {
				"saved"		: Session.STATUS_SUSPENDED,
				"aborted"	: Session.STATUS_SUSPENDED,
				"poweroff"	: Session.STATUS_SUSPENDED,
				"running"	: Session.STATUS_AVAILABLE,
				}.get(vmstate.lower())





def detect():
	from winswitch.util.simple_logger import Logger
	logger=Logger("virtualbox_server_util", log_colour=Logger.CYAN)
	from winswitch.util.config import get_local_server_config
	config = get_local_server_config(False, False)
	def add_session(*args):
		logger.slog(None, *args)
	def remove_session(*args):
		logger.slog(None, *args)
	def update_session_status(*args):
		logger.slog(None, *args)
	def session_failed(*args):
		logger.slog(None, *args)
	def detect():
		vbsu = VirtualboxServerUtil(config, add_session, remove_session, update_session_status, session_failed)
		vbsu.detect_existing_sessions([])

def main():
	from winswitch.util.main_loop import test_loop
	test_loop(detect)


if __name__ == "__main__":
	main()
