#!/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 base64
import binascii
import re

from winswitch.consts import VNC_TYPE, VNC_PORT_BASE, X11_TYPE, WINDOWS_TYPE, OSX_TYPE, DEFAULT_VNC_PORT, LOCALHOST
from winswitch.globals import WIN32, OSX
from winswitch.objects.session import Session
from winswitch.util.common import save_binary_file, is_valid_exe, visible_command, csv_list
from winswitch.util.file_io import get_session_password_filename
from winswitch.virt.server_util_base import ServerUtilBase
from winswitch.virt.options_common import CLIPBOARD, FULLSCREEN, JPEG_QUALITY_OPTION, COMPRESSION, READ_ONLY
from winswitch.util.process_util import get_output


BASE64_PREFIX = "base64:"

def obfuscate_password(vncpasswd_command, password, log_password=False):
	if WIN32:
		return	password		#not obfuscated! (still works somehow!)
	args = vncpasswd_command, '********'
	if log_password:
		args = vncpasswd_command, password
	from winswitch.util.simple_logger import Logger
	logger = Logger("vnc_server_util", log_colour=Logger.CYAN)
	cmd = [vncpasswd_command, "-f"]
	try:
		code, out, err = get_output(cmd)
		if err and err.startswith("usage: "):
			logger.serror("vncpasswd failed: looks like an old version, you should upgrade", *args)
		else:
			logger.sdebug("returncode=%s, hexlify(out)=%s, err=%s" % (code, visible_command(binascii.hexlify(out)), visible_command(err)), *args)
			logger.sdebug("plain len=%s, encrypted len(out)=%s" % (len(password), len(out)), *args)
			if code==0 and len(out)>0:
				return out
	except Exception, e:
		logger.serror("failed to run %s: %s" % (csv_list(cmd), e), *args)
	logger.serror("vncpasswd failed: returning raw password... this may or may not work", *args)
	return	password

class	VNCServerUtil(ServerUtilBase):

	IGNORE_RES = [re.compile(r"\s*Timer:\s*time has moved forwards!.*")]

	#SESSION_WAITING_RE = re.compile(r"\svncext:\sListening for VNC connections on port\s.*")
	SESSION_WAITING_RE = re.compile(r".*Listening(\sfor VNC connections)?\son\s.*port\s(\d){1,5}.*")
	SESSION_RUNNING_RE = re.compile(r"Connections:\saccepted:\s.*")
	SESSION_DISCONNECTED_RE = re.compile(r"\Connections:\sclosed:\s.*")

	def	__init__(self, config, add_session, remove_session, update_session_status, session_failed):
		ServerUtilBase.__init__(self, VNC_TYPE, VNC_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..
		self.log_password = False
		self.stop_shadows_on_close = not WIN32 and not OSX
		self.stop_shadows_on_connect_failure = True
		self.ignored_options_for_compare += [CLIPBOARD, FULLSCREEN, JPEG_QUALITY_OPTION, COMPRESSION, READ_ONLY]
		if OSX:
			from winswitch.virt.ard_control import ARD_TIMEOUT
			self.shadow_start_timeout = ARD_TIMEOUT			#ARD is very slow to react

	def get_config_options(self):
		return	self.get_config_options_base(detect=False, log=True, start=True)+[
					"# log the VNC password used - insecure!",
					"log_password",
					"# when the client closes the connection to a shadow session we created, stop the shadow automatically",
					"stop_shadows_on_close",
					"# when the client fails to connect to a shadow session, close it",
					"stop_shadows_on_connect_failure"]

	def stop(self):
		"""
		Override so we can stop winvnc
		"""
		ServerUtilBase.stop(self)
		if WIN32:
			sessions = self.config.get_sessions_by_type(self.session_type)
			self.sdebug("found %s sessions: %s" % (self.session_type, sessions))
			for session in sessions:
				if session.status==Session.STATUS_CLOSED:
					self.sdebug("session %s already closed" % session)
				else:
					try:
						self.stop_display(session, None, session.display)
					except Exception, e:
						self.exc(e)

	def new_password(self):
		"""
		VNC passwords *MUST* be 8 characters long on win32.
		"""
		p = ServerUtilBase.new_password(self)
		return	p[:8]

	def can_client_set_status(self, session, user_id, _from, _to):
		"""
		We now watch the Xvnc process output, so this should not be strictly necessary.
		"""
		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):
		"""
		The port we test to know if the session is ready is the X port.
		"""
		return	self.get_X_port(session)

	def	do_prepare_session_for_attach(self, session, user, disconnect, call_when_done):
		self.sdebug(None, session, user, disconnect, call_when_done)
		# update the keymap - causes more problems than it's worth?
		#self.set_keyboard_mappings(session, user)
		# ready?
		if session.status==Session.STATUS_AVAILABLE or (session.status in [Session.STATUS_CONNECTED, Session.STATUS_IDLE] and disconnect):
			call_when_done()
		else:
			session.add_status_update_callback(None, Session.STATUS_AVAILABLE, call_when_done)



	def stop_OSX_VNC(self, session, user, display):
		if not user or not user.password:
			self.serror("cannot stop ARD without a user and a valid password!", session, user, display)
			return
		self.update_session_status(session, Session.STATUS_CLOSED)					#We have to assume the ARD command above will do its work properly..
		from winswitch.virt.ard_control import stop_ARD
		stop_ARD(user.password)

	def start_OSX_VNC(self, session, user, vnc_password):
		from winswitch.virt.ard_control import start_ARD, is_ARD_in_use
		if is_ARD_in_use():
			self.early_failure(session, "Apple Remote Desktop is currently starting or stopping, please wait a few seconds and try again")
			return	False
		if not user.password:
			self.early_failure(session, "Password not available for this login but the password is required to enable 'Apple Remote Desktop'")
			return	False
		def ARDControl_done(proc):
			""" This is called when the kickstart process has returned """
			self.slog("returncode=%s" % proc.returncode, proc)
			def ARD_failed(message="Failed to enable the VNC Server"):
				self.update_session_status(session, Session.STATUS_CLOSED)					#Something went wrong...
				self.early_failure(session, message)
			def ARD_ready():
				self.update_session_status(session, Session.STATUS_AVAILABLE)
			def ARD_test_abort():
				return	session.status==Session.STATUS_CLOSED
			if proc.returncode!=0:
				ARD_failed()
			else:
				""" It should be ready but isn't.. wait for it """
				from winswitch.net.net_util import wait_for_socket
				wait_for_socket(LOCALHOST, DEFAULT_VNC_PORT, max_wait=30,
							success_callback=ARD_ready, error_callback=ARD_failed, abort_test=ARD_test_abort)

		self.slog("VNC password set to %s" % session.password, session, user, vnc_password)
		start_ARD(user.password, vnc_password, ARDControl_done)
		return	True

	def start_display(self, session, user, is_preload):
		self.sdebug(None, session, user, is_preload)
		if session.shadowed_display:
			if self.stop_shadows_on_close:
				#stop the shadow as soon as client disconnects
				def close_on_suspend():
					self.stop_display(session, user, session.display)
				session.add_status_update_callback(Session.STATUS_CONNECTED, Session.STATUS_AVAILABLE, close_on_suspend, clear_it=True, timeout=None)
			if self.stop_shadows_on_connect_failure:
				from winswitch.util.main_loop import callLater
				def check_connected():
					self.sdebug("status=%s" % session.status)
					if session.status!=Session.STATUS_CONNECTED:
						self.stop_display(session, user, session.display)
				callLater(30, check_connected)
		#Examples:
		#ie: /usr/local/bin/Xvnc :63 -rfbport 16064 -rfbauth /home/freebsd8/.winswitch/server/sessions/63.pass -desktop gnome-about-me
		#ie: x0vncserver --rfbport=12345 --display=:67 --PasswordFile=./passwordfile
		#ie: x11vnc --rfbport 12345 --display :64
		#ie: C:\...\winvnc4.exe -PasswordFile=./passwordfile -SecurityTypes=None

		#Xvnc/x0vnc need an obfuscated password:
		raw_password = session.password
		obfuscated_password = obfuscate_password(self.config.vncpasswd_command, session.password, self.log_password)
		session.password = "base64:%s " % base64.b64encode(obfuscated_password)

		if OSX:
			assert not is_preload
			return self.start_OSX_VNC(session, user, raw_password)

		password_file = get_session_password_filename(session.display, session.user)
		save_binary_file(password_file, obfuscated_password)
		if WIN32:
			assert not is_preload
			args_list = ['%s' % self.config.winvnc_command, '-PasswordFile=%s' % password_file, '-PortNumber=%s' % session.port]
			if not self.config.show_vnc_tray:
				args_list.append('-ShowTrayIcon=no')
		else:
			auth_param = "-rfbauth"
			if session.shadowed_display:
				assert self.config.vncshadow_command
				args_list = [self.config.vncshadow_command, "-display", session.shadowed_display]
				if self.config.vncshadow_command.find("x0vncserver")>=0:
					auth_param = "-PasswordFile"
			else:
				args_list = [self.config.xvnc_command, session.display]
				#-desktop is also supported by x11vnc, but not by x0vncserver... so we dont set it globally
				args_list += ["-desktop", "%s" % session.name]
				args_list += self.get_X_geometry_args(session.screen_size)
				if self.config.xvnc_options:
					for opt in self.config.xvnc_options:
						if opt!="-once":
							args_list.append(opt)
				#xvnc_type = get_xvnc_type(self.config.xvnc_command)
				#if xvnc_type == XVNC_TIGER:
				#	args_list.append("-log *:stderr:30")
				#	args_list.append(logfile)
			args_list += ["-rfbport", "%d" % session.port, auth_param, "%s" % password_file]
		return self.start_daemon(session, args_list, session.get_env())

	def get_unique_shadow(self, session, user, read_only):
		""" On win32 and osx, the display can only be exported once (one server instance - one server port)
			So we must ensure we return the existing instance if there is one """
		if self.config.export_local_displays and (\
			(WIN32 and session.session_type==WINDOWS_TYPE) or \
			(OSX and session.session_type==OSX_TYPE) \
			):
			for session in self.config.get_all_sessions_for_display(session.display, shadows=True):
				if session.session_type==VNC_TYPE:
					self.prepare_display(session, user)
					return session
		return	None	# No existing shadow found (or we dont require a unique one if on non-win32/non-osx)

	def create_shadow_session(self, session, user, read_only, screen_size, options):
		"""
		Override so we can validate the session we are supposed to shadow.
		"""
		if (WIN32 or OSX) and not self.config.export_local_displays:
			self.serror("export_local_displays is not set, cannot export display", session, user, read_only, screen_size, options)
			return	None
		if WIN32:
			if not self.config.supports_winvnc:
				raise Exception("exporting is not allowed: supports_winvnc is not set")
			if not is_valid_exe(self.config.winvnc_command):
				raise Exception("winvnc_command is invalid: %s" % self.config.winvnc_command)
			if session.session_type!=WINDOWS_TYPE:
				raise Exception("invalid session type: %s" % session.session_type)
			shadow = ServerUtilBase.new_shadow_session(self, session, user, read_only, screen_size, options)
			shadow.port = self.config.winvnc_port
			self.save_session(shadow)
			self.config.add_session(shadow)
			self.add_session(shadow)
			return shadow
		if OSX:
			if session.session_type!=OSX_TYPE:
				raise Exception("invalid session type: %s" % session.session_type)
			shadow = ServerUtilBase.new_shadow_session(self, session, user, read_only, screen_size, options)
			shadow.port = DEFAULT_VNC_PORT
			self.save_session(shadow)
			self.config.add_session(shadow)
			self.add_session(shadow)
			return shadow
		assert session.session_type in [VNC_TYPE, X11_TYPE]
		return ServerUtilBase.create_shadow_session(self, session, user, read_only, screen_size, options)

	def stop_display(self, session, user, display):
		"""
		OSX override so we can disable the VNC server when the session is closed.
		(on other platforms, the vnc server process is controlled by us, so we can just kill it)
		"""
		if OSX:
			self.stop_OSX_VNC(session, user, display)
			return	#nothing else to do: there is no process to kill on osx
		ServerUtilBase.stop_display(self, session, user, display)


	def process_log_line(self, session, line):
		if not line or line=='\n':
			return
		for ignore in VNCServerUtil.IGNORE_RES:
			if ignore.match(line):
				return
		trimmed = line.strip()
		if WIN32:
			if trimmed.startswith("Starting User-Mode VNC Server."):
				return	Session.STATUS_AVAILABLE
		if VNCServerUtil.SESSION_WAITING_RE.match(trimmed):
			return	Session.STATUS_AVAILABLE
		elif VNCServerUtil.SESSION_RUNNING_RE.match(trimmed):
			return	Session.STATUS_CONNECTED
		elif VNCServerUtil.SESSION_DISCONNECTED_RE.match(trimmed):
			if not trimmed.endswith("(Non-shared connection requested)"):		#previous connection closed, not current
				return	Session.STATUS_AVAILABLE
		return	None


	def start_session_command(self, session):
		"""
		We override so we can also start the vnc_session_commands (devilpsie, xloadimage..
		and the window manager)
		And also set the screen size (if needed).
		"""
		if not session.full_desktop:
			self.sdebug("full_desktop=%s, wm=%s, desktop session commands=%s, vnc session commands=%s" % (session.full_desktop, self.config.desktop_window_manager_command, str(self.config.desktop_session_commands), str(self.config.vnc_session_commands)), session)
			#Start a window manager:
			self.do_start_session_command(session, self.config.desktop_window_manager_command, True)
			#Start VNC specific commands (devilspie, xloadimage, ...) - no need to kill those on exit
			self.do_start_session_commands(session, self.config.vnc_session_commands, False)
		#screen size is already set in the start display command
		#self.do_start_session_command(session, self.get_randr_command(session), False)
		ServerUtilBase.start_session_command(self, session)
