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

from winswitch.consts import WINDOWS_TYPE, CONSOLE_CONNECT, CONSOLE_DISCONNECT, SESSION_LOCK, SESSION_UNLOCK, DEFAULT_RDP_PORT, MAGIC_NO_PASSWORD, MAGIC_SERVER_PASSWORD
from winswitch.globals import OSX, WIN32
from winswitch.objects.session import Session
from winswitch.objects.server_session import ServerSession
from winswitch.objects.server_command import ServerCommand
from winswitch.util.file_io import get_client_session_rdp_config_file
from winswitch.util.common import csv_list, save_binary_file, parse_screensize, get_bool, generate_UUID, is_valid_dir, is_valid_file
from winswitch.virt.client_util_base import ClientUtilBase
from winswitch.net.net_util import wait_for_socket
from winswitch.virt.options_common import FULLSCREEN, KEYMAP

unique = 0
ENABLE_KEYMAP_OPTION = True
DEBUG_RDP = ("--debug-rdp" in sys.argv) or OSX
FREERDP_DISABLE_TLS = False

class	RDPClientBase(ClientUtilBase):
	"""
	Starts and manages RDP client connections.
	We use mstsc.exe on Windows and rdesktop on *nix.
	"""

	CONNECT_FAILED_RE = re.compile(r"ERROR:\s.*:\sunable to connect")
	RESOLUTION_CHANGED_RE = re.compile(r"WARNING:\s*Remote desktop changed from ([0-9]*x[0-9]*) to ([0-9]*)x([0-9]*).")

	def	__init__(self, update_session_status, notify_callback):
		self.keymaps = {}
		self.default_keymap = ""
		ClientUtilBase.__init__(self, WINDOWS_TYPE, update_session_status, notify_callback)
		if ENABLE_KEYMAP_OPTION:
			self.keymaps = self.find_all_keymaps()
			self.default_keymap = self.get_default_keymap()
		self.ignore_process_returncodes.append(2)		#closing the window causes this returncode!

	def find_all_keymaps(self):
		if WIN32 or OSX:
			return	{}
		keymaps = {}
		for d in [os.path.join(os.path.expanduser("~/.rdesktop"), "keymaps"),
					"/usr/share/rdesktop/keymaps",
					"/usr/local/share/rdesktop/keymaps",
					os.path.join(os.getcwd(), "keymaps")]:
			if not is_valid_dir(d):
				self.sdebug("invalid dir skipped: %s" % d)
				continue
			for f in os.listdir(d):
				if f in keymaps or f in ["common", "modifiers"]:
					continue
				filename = os.path.join(d, f)
				if is_valid_file(filename):
					keymaps[f] = filename
		self.slog("=%s" % keymaps.keys())
		return	keymaps

	def get_default_keymap(self):
		if WIN32 or OSX:
			return	None
		lang = os.environ.get("LANG")
		if not lang:
			return	None
		layout = None
		if self.settings.xkbmap_query:
			for line in self.settings.xkbmap_query.splitlines():
				if line.startswith("layout:"):
					layout = line[len("layout:"):].strip()
		l = lang
		if l.find("."):
			l = l[:l.find(".")]	#-> "en_US"
		keys = [l]				#en_US
		langs = l.split("_")	#-> ["en", "US"]
		if layout:
			keys.append("%s-%s" % (langs[0], layout))	#en-gb
		keys.append(langs[0])		#en
		self.sdebug("trying to find keymap for lang=%s and layout=%s using keys=%s" % (lang, layout, keys))
		for k in keys:
			if k in self.keymaps.keys():
				return	k
		return	None


	def handle_session_event(self, servers, event, name):
		local = self.get_local_serversession(servers)
		self.last_session_event = event
		self.slog("local=%s" % str(local), servers, event, name)
		if not local:
			return
		(server, session) = local
		from winswitch.util.main_loop import callLater, callFromThread
		def set_status(new_status):
			callFromThread(self.update_session_status, server, session, new_status)
		if event==CONSOLE_CONNECT or event==SESSION_UNLOCK:
			set_status(Session.STATUS_CONNECTED)
		elif event==CONSOLE_DISCONNECT:
			""" mark the session as disconnected, but only if we don't get a CONNECT shortly after: """
			def check_still_disconnected():
				if self.last_session_event==CONSOLE_DISCONNECT:
					set_status(Session.STATUS_AVAILABLE)
			callLater(0.2, check_still_disconnected)
		elif event==SESSION_LOCK:
			set_status(Session.STATUS_IDLE)

	def get_local_serversession(self, servers):
		""" Identifies which one (if any) of the servers owns the current session (main display) """
		self.sdebug(None, servers)
		for server in servers:
			if server.local:
				sessions = server.get_sessions_by_type(self.session_type)
				for session in sessions:
					self.sdebug("owner(%s)=%s, uuid=%s" % (session, session.owner, self.settings.uuid), servers)
					if session.owner==self.settings.uuid:
						return	(server, session)
		return None


	def	client_start_session(self, server, client, command, screen_size, opts):
		self.slog(None, server, client, command, screen_size, opts)
		session = self.initialize_new_session(server, command, screen_size, options=opts)
		seamless_command = None
		if command.type==ServerCommand.COMMAND:
			seamless_command = session.command
		server.add_session(session)
		self.start_RDP_client(server, session, seamless_command)

	def do_real_attach(self, server, session, host, port):
		""" Detect firewalls which may be blocking the port then connect """
		if session.port_tested:
			self.sdebug("port already tested, starting client", server, session, host, port)
			self.start_RDP_client(server, session, host, port, None)
			return
		#test the port:
		def session_errored_out():
			return	session.status in [Session.STATUS_CLOSED, Session.STATUS_UNAVAILABLE, Session.STATUS_SUSPENDED]
		def session_port_ok():
			self.sdebug()
			session.port_tested = True
			self.start_RDP_client(server, session, host, port, None)
		def session_port_failed():
			if session_errored_out():
				return
			if port==DEFAULT_RDP_PORT:
				port_descr = "the 'Remote Desktop' port"
			else:
				port_descr = "port %s" % port
			self.notify_error("RDP Session Failed", "Cannot connect to server %s, please ensure that the firewall is not blocking %s" % (server.get_display_name(), port_descr))
			self.update_session_status(server, session, Session.STATUS_AVAILABLE)
		wait_for_socket(host, port, max_wait=10, success_callback=session_port_ok, error_callback=session_port_failed, abort_test=session_errored_out)


	def	start_RDP_client(self, server, session, host, port, seamless_command):
		if WIN32:
			name = "%s on %s" % (session.name, server.get_display_name())
			rdp_config_file = get_client_session_rdp_config_file(session.ID, name)
			self.save_rdp_config_file(rdp_config_file, server, session, host, port)
			args = [self.settings.rdesktop_command, rdp_config_file]
			log_args = True
		else:
			args = self.rdp_command_args(server, session, host, port, seamless_command)
			log_args = False or DEBUG_RDP			#dont log password!
		self.schedule_connect_check(session, kill_it=True)			#kill it if it is still CONNECTING after 30s
		self.exec_client(server, session, args, onstart_status=Session.STATUS_CONNECTED, onexit_status=Session.STATUS_AVAILABLE, log_args=log_args)

	def rdp_command_args(self, server, session, host, port, seamless_command):
		""" constructs the rdesktop/xfreerdp command line for a session """
		args = [self.settings.rdesktop_command]
		is_freerdp = self.settings.rdesktop_command.lower().find("freerdp")>=0
		if seamless_command:
			if is_freerdp:
				self.serror("cannot use seamless mode with freerdp! defaulting to desktop mode", server, session, host, port, seamless_command)
			else:
				args.append("-A")
				args.append("-s")
				assert server.rdp_seamless_command
				args.append("%s %s" % (server.rdp_seamless_command, session.command))
				#ie: -s "C:\SeamlessRDP\seamlessrdpshell.exe notepad"
		if ENABLE_KEYMAP_OPTION and self.default_keymap:
			args.append("-k")
			args.append(self.default_keymap)
		args.append("-u")
		args.append(session.user)
		if session.password==MAGIC_SERVER_PASSWORD:
			if not server.password:
				self.notify_error("Password is missing", "The RDP session '%s' needs to use the server password,\nbut the password is not set!" % session.name)
			else:
				args.append("-p")
				args.append(server.password)
		elif session.password!=MAGIC_NO_PASSWORD:
			args.append("-p")
			args.append(session.password)
		fullscreen = session.options.get(FULLSCREEN)
		if fullscreen is not None and get_bool(fullscreen):
			args.append("-f")
		if is_freerdp:
			if FREERDP_DISABLE_TLS:
				args.append("--no-tls")
		else:
			""" only rdesktop supports setting the session name """
			args.append("-T")
			args.append("%s on %s" % (session.name, server.get_display_name()))
		if server.line_speed<56*1000:		#<56K
			args.append("-z")		#compression should probably be made an option
			perf = "m"				#modem
		elif server.line_speed>1000*1000:	#>1Mbit
			perf = "l"				#lan
		else:
			perf = "b"				#broadband
		args.append("-x")
		args.append(perf)
		ss = parse_screensize(session.screen_size)
		if ss:
			(width, height, depth) = ss
			args.append("-g")
			args.append("%sx%s" % (width, height))
			if depth in [8, 15, 16, 24]:
				args.append("-a %s" % depth)
		if session.port==DEFAULT_RDP_PORT:
			args.append("%s" % host)
		else:
			args.append("%s:%s" % (host, port))
		return args

	def initialize_new_session(self, server, command, screen_size):
		global unique
		unique += 1
		session = ServerSession(self.config)
		session.display = "RDP%s" % unique
		session.ID = generate_UUID()
		session.host = server.host
		session.port = server.rdp_port
		session.name = command.name
		session.command = command.command
		session.commands = [session.command]
		session.status = Session.STATUS_CONNECTED
		session.session_type = WINDOWS_TYPE
		session.start_time = time.time()
		session.actor = self.settings.uuid
		session.owner = self.settings.uuid
		session.set_default_icon_data(command.get_icon_data())
		session.screen_size = screen_size
		return	session

	def client_detach_session(self, server, client, session):
		"""
		Override so we force the server to set the session state to "AVAILABLE".
		RDP session state transitions between CONNECTED and AVAILABLE are handled by
		the do_real_attach() method above, but if we re-started the client without
		re-setting the state to AVAILABLE it may never reach it.
		"""
		ClientUtilBase.client_detach_session(self, server, client, session)
		self.update_session_status(server, session, Session.STATUS_AVAILABLE)

	def save_rdp_config_file(self, rdp_config_file, server, session, host, port):
		self.sdebug(None, rdp_config_file, server, session)
		ss = parse_screensize(session.screen_size)
		if not ss:
			ss = (1024, 768, 32)
		(width, height, depth) = ss
		if not depth:
			depth = 32
		#found some documentation here:
		#http://www.xpunlimited.com/faq/index.php?sid=131402&lang=en&action=artikel&cat=4&id=112&artlang=en
		#http://www.teachout.com/Blog/tabid/165/EntryId/170/RDP-File-Configuration-Options.aspx
		config = [
				#general attributes:
				"displayconnectionbar:i:1", "keyboardhook:i:2"
				"autoreconnection enabled:i:0",
				#performance related:
				"compression:i:1",  "disable wallpaper:i:1", "disable full window drag:i:1",
				"allow desktop composition:i:0", "allow font smoothing:i:0",
				"disable menu anims:i:1", "disable themes:i:0",
				"disable cursor setting:i:0", "bitmapcachepersistenable:i:1",
				"bitmapcachepersistenable:i:1",
				#connection details:
				"authentication level:i:0",
				"negotiate security layer:i:1",
				"promptcredentialonce:i:0",
				"prompt for credentials:i:0",
				"alternate shell:s:",
				"shell working directory:s:",
				"remoteapplicationmode:i:0",					#seamless mode?
				"remoteapplicationname:s:",
				#screen dimensions:
				"desktopwidth:i:%s" % width, "desktopheight:i:%s" % height, "session bpp:i:%s" % depth,
				"winposstr:s:0,3,0,0,800,600",
				#not using a gateway:
				"gatewayhostname:s:",
				"gatewayusagemethod:i:4",
				"gatewaycredentialssource:i:4",
				"gatewayprofileusagemethod:i:0",
				#we will deal with this via gst/pulse?
				"audiomode:i:0",
				#not sure about those..
				"redirectprinters:i:1",
				"redirectcomports:i:0",
				"redirectsmartcards:i:1",
				"redirectclipboard:i:1",
				"redirectposdevices:i:0",
				]
		fullscreen = session.options.get(FULLSCREEN)
		if fullscreen and get_bool(fullscreen):
			config.append("screen mode id:i:2")		#2==fullscreen
		else:
			config.append("screen mode id:i:1")		#1==normal??
		#connection info:
		config.append("full address:s:%s" % host)
		config.append("username:s:%s" % session.user)
		config.append("port:i:%s" % port)
		def add_password(password):
			#encrypt the password:
			#as seen here: http://stackoverflow.com/questions/2374331/python-win32crypt-cryptprotectdata-difference-between-2-5-and-3-1
			#and here: http://www.remkoweijnen.nl/blog/2007/11/05/encrypt-rdp-password-in-python/
			try:
				import win32crypt	#@UnresolvedImport
				import binascii
				u_pwd = unicode(password)
				hashed = win32crypt.CryptProtectData(u_pwd, u"", None, None, None, 0)
				hex_pwd = str(binascii.hexlify(hashed)).upper()
				config.append(u"password 51:b:%s" % hex_pwd)
				#self.slog("unicode(%s)=%s" % (session.password, u_pwd.decode('utf-8', "ignore")), rdp_config_file, server, session)
				#self.slog("hashed=%s" % hashed.decode('utf-8', "ignore"), rdp_config_file, server, session)
				self.sdebug("hex_pwd=%s" % hex_pwd, rdp_config_file, server, session)
			except Exception, e:
				self.notify_error("Password error",
								"Failed to hash the password: %s.\nYou will have to type it in, sorry." % e)
				self.serr("failed to hash the password", e)

		if session.password==MAGIC_SERVER_PASSWORD:
			if not server.password:
				self.notify_error("Password is missing",
								"The RDP session '%s' needs to use the server password,\nbut the password is not set!\nYou will have to type it in." % session.name)
			else:
				add_password(server.password)
		elif session.password==MAGIC_NO_PASSWORD:
			pass
		else:
			if session.password:
				add_password(session.password)
			else:
				self.slog("no password available! you will be prompted...", rdp_config_file, server, session, host, port)
		config_str = csv_list(config, quote=None, sep='\n', before="", after="\n")
		save_binary_file(rdp_config_file, config_str)


	def handle_line(self, server, session, line):
		m = RDPClientBase.RESOLUTION_CHANGED_RE.match(line)
		if m:
			old_res = m.group(1)
			new_w = int(m.group(2))
			new_h = int(m.group(3))
			self.slog("resolution changed from %s to %sx%s" % (old_res, new_w, new_h))
			screen_size = "%sx%s" % (new_w, new_h)
			from winswitch.util.main_loop import callFromThread
			callFromThread(self.update_session_status, server, session, None, screen_size)
		if RDPClientBase.CONNECT_FAILED_RE.match(line):
			#TODO: notify user?
			self.serror(None, server, session, line)
		else:
			self.sdebug(None, server, session, line)
		#"disconnect: Server initiated disconnect." == rdesktop wrong password message?
		#"ui_error: ERROR: send: Broken pipe" == xfreerdp error?

	def get_options_defaults(self):
		d = { FULLSCREEN: False }
		if ENABLE_KEYMAP_OPTION:
			d[KEYMAP] = self.get_default_keymap()
		return	d



def main():
	from winswitch.util.simple_logger import Logger
	logger = Logger("rdp_client_base")
	from winswitch.objects.global_settings import get_settings
	get_settings(True, True).xkbmap_query = """rules:      base
model:      pc105
layout:     gb
"""
	def update_session_status(*args):
		logger.slog(None, *args)
	def notify_callback(*args):
		logger.slog(None, *args)
	rdp_client = RDPClientBase(update_session_status, notify_callback)
	assert rdp_client

if __name__ == "__main__":
	main()
