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

from winswitch.consts import NX_TYPE, NX_PORT_BASE, VNC_TYPE, X11_TYPE
from winswitch.objects.session import Session
from winswitch.util.process_util import send_kill_signal
from winswitch.util.common import generate_UUID
from winswitch.virt.server_util_base import ServerUtilBase
from winswitch.virt.nx_common import get_line_speed_str, nx_display_value, SET_COOKIE, USE_OPTIONS_FILE
from winswitch.virt.reserve_port import ReservePort
from winswitch.virt.options_common import CLIPBOARD, FULLSCREEN, BANDWIDTH
from winswitch.util.main_loop import callLater

RESTRICT_ACCESS = True			#listen on loopback only when server or session is always tunnelled
#disabled as in some cases we want to publish the server as being tunneled, but allow non-tunneled connections too


class	NXServerUtil(ServerUtilBase):
	"""
	Utility class used by the server to start sessions and monitor them.
	(not to be confused with the NX client/server terminology which is reversed)
	"""
	DEFAULT_NXAGENT_OPTIONS = {
					"keyboard": "evdev/us",
					"cache": "32M",
					"images": "16M",
					"clipboard": "both",
					"render": 1,
					"composite": 1,
					"load": 1,
					"save": 1,
					"cleanup": 0,
					"core": 0,
					"backingstore": 1,
					"shmem": "1",
					"shpix": "1",
					"strict": 0,
					"media": 0
					}

	# Regular expressions loosely based on the ones from neatx's agent.py
	SESSION_STARTING_RE = re.compile(r"^Session:\s+Starting\s+session\s+at\s+")
	SESSION_WAITING_RE = re.compile(r"Info:\s+Waiting\s+for\s+connection\s+from\s+")
	SESSION_RUNNING_RE = re.compile(r"^Session:\s+Session\s+(started|resumed)\s+at\s+")
	SESSION_SUSPENDING_RE = re.compile(r"^Session:\s+Suspending\s+session\s+at\s+")
	SESSION_SUSPENDED_RE = re.compile(r"^Session:\s+Session\s+suspended\s+at\s+")
	SESSION_WATCHDOG_PID_RE = re.compile(r"^Info:\s+Watchdog\s+running\s+with\s+pid\s+'(?P<pid>\d+)'\.")
	SESSION_WATCHDOG_WAIT_RE = re.compile(r"^Info:\s+Waiting\s+the\s+watchdog\s+process\s+to\s+complete\.")
	SESSION_TERMINATING_RE = re.compile(r"^Session:\s+(Terminat|Abort)ing\s+session\s+at\s+")
	SESSION_TERMINATED_RE = re.compile(r"^Session:\s+Session\s+(terminat|abort)ed\s+at\s+")

	#IGNORED_RES = [re.compile(r"Warning:\sCouldn't\sstart\s'/usr/NX/bin/nxclient'."),
	#				re.compile(r"Warning:\sTrying\swith\spath")]

	def	__init__(self, config, add_session, remove_session, update_session_status, session_failed):
		ServerUtilBase.__init__(self, NX_TYPE, NX_PORT_BASE, config, add_session, remove_session, update_session_status, session_failed)

		self.prelaunch_enabled = False			#currently broken for NX
		self.ignored_options_for_compare += [CLIPBOARD, FULLSCREEN]
		self.state_map = [
			#(regular_expression, new_state, callback)
			# The callback takes precedence and may return the new state
			(self.SESSION_STARTING_RE,		Session.STATUS_STARTING,	None),
			(self.SESSION_WAITING_RE,		Session.STATUS_AVAILABLE,	None),
			(self.SESSION_RUNNING_RE,		Session.STATUS_CONNECTED,	None),
			(self.SESSION_SUSPENDING_RE,	Session.STATUS_SUSPENDING,	self.suspending),
			(self.SESSION_SUSPENDED_RE,		Session.STATUS_SUSPENDED,	None),
			(self.SESSION_WATCHDOG_PID_RE,	None,						self.save_watchdog_pid),
			(self.SESSION_WATCHDOG_WAIT_RE,	None,						self.kill_watchdog),
			(self.SESSION_TERMINATING_RE,	Session.STATUS_CLOSED,		None),
			(self.SESSION_TERMINATED_RE,	Session.STATUS_CLOSED,		None)
			]
		self.default_nxagent_options = dict(NXServerUtil.DEFAULT_NXAGENT_OPTIONS)

	def get_config_options(self):
		return	self.get_config_options_base(detect=False, log=True, start=True)+[
					"# default options for the nxagent configuration file",
					"default_nxagent_options"]

	def new_password(self):
		return	generate_UUID()

	def can_capture(self, session):
		""" NX can only capture screenshots when in full desktop mode, seamless mode does not show anything!? """
		return	ServerUtilBase.can_capture(self, session) and session.full_desktop

	def stop_preloaded_session(self, session):
		"""
		Preload sessions use a fake process that understands stop()
		"""
		assert session.port_reserve is not None
		session.port_reserve.stop()

	def is_prelaunched_session_compatible(self, session, screen_size, opts):
		"""
		NX prelaunch sessions consist just of a reserved port number.
		So we can *always* use them, no matter what the options are since they'll be used for starting the real session.
		"""
		return	True

	def suspending(self, session, match):
		self.sdebug("scheduling check_suspended()", session, match)
		callLater(30, self.check_suspended, session, session.status_update_count)
		return	Session.STATUS_SUSPENDING

	def check_suspended(self, session, prev_update_count):
		stuck = session.status==Session.STATUS_SUSPENDING and session.status_update_count==prev_update_count
		self.sdebug("current count=%s, stuck in suspending state=%s" % (session.status_update_count, stuck), session, prev_update_count)
		if stuck:
			if session.watchdog_pid:
				self.kill_watchdog(session, None)
			else:
				self.nxagent_sighup(session)

	def save_watchdog_pid(self, session, match):
		session.watchdog_pid = int(match.group("pid"))
		self.slog("pid=%s" % session.watchdog_pid, session, match)
		return	None

	def kill_watchdog(self, session, match):
		self.slog("pid=%s" % session.watchdog_pid, session, match)
		if session.watchdog_pid:
			send_kill_signal(session.watchdog_pid)
		return	None	#	Session.STATUS_CLOSED

	def nxagent_sighup(self, session):
		self.sdebug(None, session)
		if session.server_process_pid>1:
			send_kill_signal(session.server_process_pid, signal.SIGHUP)

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

	def	do_prepare_session_for_attach(self, session, user, disconnect, call_when_done):
		"""
		Updates the options file for the new session user, runs xmodmap.
		Then we hook up a bunch of session state callbacks to make sure it will end up
		firing "call_when_done" no matter what state it is currently in.
		"""
		self.slog(None, session, user, disconnect, call_when_done)
		if USE_OPTIONS_FILE:
			# update options:
			options = self.get_nxagent_options(session, user)
			nx_display_value(options, session.user, session.display, agent=True, optionsfile_dir=self.get_session_dir(session))
		self.set_xauth(session, user)
		# ready?
		if session.status == Session.STATUS_AVAILABLE:
			call_when_done()
		else:
			session.add_status_update_callback(None, Session.STATUS_AVAILABLE, call_when_done)

		# unsuspend it
		if session.status == Session.STATUS_SUSPENDED:
			self.nxagent_sighup(session)
		else:
			def suspended_sighup_nxagent():
				self.slog()
				self.nxagent_sighup(session)
			session.add_status_update_callback(None, Session.STATUS_SUSPENDED, suspended_sighup_nxagent)

		# if we need to force disconnect it... then force suspend it (will fire suspend -> sighup -> available)
		if session.status in [Session.STATUS_IDLE, Session.STATUS_CONNECTED] and disconnect:
			#could give it a bit more time but would need to be careful using a timer that we don't kill the new one! (use a status update counter to ensure)
			self.nxagent_sighup(session)

	def set_keyboard_mappings(self, session, user):
		pass

	def get_nxagent_options(self, session, user=None):
		nx_options = self.default_nxagent_options
		nx_options["listen"] = "%d" % session.port
		if RESTRICT_ACCESS and (self.config.ssh_tunnel or session.requires_tunnel):
			nx_options["accept"] = "127.0.0.1"
		if SET_COOKIE:
			nx_options["cookie"] = session.password
		if user:
			if BANDWIDTH in session.options:
				nx_options["link"] = session.options.get(BANDWIDTH)
			else:
				nx_options["link"] = get_line_speed_str(user.line_speed)
			if user.platform and user.platform.startswith("linux"):
				nx_options["client"] = "linux"
		if session.full_desktop and session.shared_desktop:
			if session.read_only:
				nx_options["shadowmode"] = "0"
			else:
				nx_options["shadowmode"] = "1"
			assert session.shadowed_display is not None
			nx_options["shadow"] = session.shadowed_display
		self.slog("=%s" % nx_options, session, user)
		return	nx_options

	def prelaunch_start_command(self, session, server_command, user, screen_size, opts, filenames):
		"""
		pre-launch sessions with NX start a ReservePort,
		we must stop it and start the real nx display in its place.
		"""
		self.sdebug("current server_process_pid=%s" % session.server_process_pid, session, server_command, user, screen_size, opts, filenames)
		session.port_reserve.stop()
		#initialise some session fields:
		ServerUtilBase.prelaunch_start_command(self, session, server_command, user, screen_size, opts, filenames)
		#start the real NX process:
		self.set_xauth(session, user)
		self.start_display(session, user, False)
		self.sdebug(None, session, server_command, user, screen_size, opts, filenames)

	def get_unique_shadow(self, session, user, read_only):
		return	None		#there are no "Unique" shadow session restrictions with NX, always create a new one

	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 session.session_type==NX_TYPE:
			assert session.full_desktop is True		#cannot shadow rootless sessions
		else:
			assert session.session_type in [VNC_TYPE, X11_TYPE]
		return ServerUtilBase.create_shadow_session(self, session, user, read_only, screen_size, options)

	def xauth_enabled(self, session):
		return	True				#always use xauth with NX: used for cookie authentication

	def start_session_object(self, session, user, is_preload):
		if not is_preload:
			self.set_xauth(session, user)		#needed for NX! authentication fails otherwise!
		return ServerUtilBase.start_session_object(self, session, user, is_preload)

	def start_display(self, session, user, is_preload):
		"""
		Start a display.
		NX cannot start a session without having the client attach to it within 30 seconds...
		So we can't start a real prelaunch session here.
		We just reserve the ports (and bind to them)
		which will still allow the clients to pre-load the SSH tunnels in advance.
		"""
		self.sdebug(None, session, user, is_preload)
		if is_preload:
			session.port_reserve = ReservePort(session.host, session.port)
			return	True

		nx_options = self.get_nxagent_options(session, user)
		DISPLAY = nx_display_value(nx_options, session.user, session.display, agent=True, optionsfile_dir=self.get_session_dir(session))
		name = "%s:%s" % (session.user, session.name.replace(" ", "_"))

		#ie: nxagent -R -name 'user: gedit' :8
		args_list = [self.config.nxagent_command]
		if not session.full_desktop:
			args_list.append("-R")			#rootless
		elif session.shared_desktop:
			args_list.append("-S")			#shadow
		else:
			args_list.append("-D")			#desktop

		if session.full_desktop:
			args_list += self.get_X_geometry_args(session.screen_size)

		#somehow this is different from the composite=1 we set in the options file?
		args_list += ["+extension","Composite"]
		args_list += ["-dpi", "96"]
		#args_list += ["-nocomposite"]
		args_list += ["-name", name, session.display]
		env = session.get_env();
		env["DISPLAY"] = DISPLAY
		self.slog("starting %s with env %s" % (args_list, env), session, user, is_preload)
		return	self.start_daemon(session, args_list, env)

	def process_log_line(self, session, line):
		status = None
		for cre,new_status,callback in self.state_map:
			match = cre.match(line)
			#self.sdebug("match(%s)=%s for status=%s" % (re, match, new_status), line)
			if not match:
				continue
			if callback and not session.reloading:
				try :
					status = callback(session, match)
				except Exception, e:
					self.serr(None, e)
			else:
				status = new_status
			return	status
		return	None
