# -*- coding: utf-8 -*-

# Copyright (C) 2010-2012 by Mike Gabriel <mike.gabriel@das-netzwerkteam.de>
#
# PyHoca CLI is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# PyHoca CLI is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the
# Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.
#
# Contributors to the code of this programme:
#    2010 Dick Kniep <dick.kniep@lindix.nl>
#    2010 Jörg Sawatzki <joerg.sawatzki@web.de>

import x2go
import sys
import os
import copy
import time
import getpass
from types import *
from gevent import socket

# for debugging
import pprint

# a list of available X2Go print actions
from x2go.defaults import X2GO_PRINT_ACTIONS

# use current_home as user home dir
current_home = os.path.expanduser("~")

from x2go.utils import touch_file as _touch_file

# define and create known_hosts file (if not there)
ssh_known_hosts_filename = os.path.join(x2go.LOCAL_HOME, x2go.X2GO_SSH_ROOTDIR, 'known_hosts')
if not os.path.isfile(ssh_known_hosts_filename):
    try: _touch_file(ssh_known_hosts_filename)
    except OSError: pass
# define and create ssh_config file (if not there)
ssh_config_filename = os.path.join(x2go.LOCAL_HOME, x2go.X2GO_SSH_ROOTDIR, 'config')
if not os.path.isfile(ssh_config_filename):
    try: _touch_file(ssh_config_filename)
    except OSError: pass


# sometimes we have to fail...
def runtime_error(m, parser=None, exitcode=-1):
    """\
    STILL UNDOCUMENTED
    """
    if parser is not None:
        parser.print_usage()
    sys.stderr.write ("%s: error: %s\n" % (os.path.basename(sys.argv[0]), m))
    sys.exit(exitcode)


class PyHocaCLI(x2go.X2GoClient):
    """\
    L{PyHocaCLI} is a command-line X2Go Client wrapping around Python X2Go (C{X2GoClient} class).

    """
    x2go_session_hash = None

    def _runtime_error(self, m, exitcode=-1):
        runtime_error(m, exitcode=exitcode)

    def list_sessions(self, s_hash):
        """\
        List up server-side available sessions for the logged in user.

        The list of available session is printed to stdout.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # retrieve a session list
        print
        print "Available runing/suspended X2Go sessions"
        print "========================================"
        _peername = self._X2GoClient__get_session_server_peername(s_hash)
        print "Host: %s - [%s]:%s" % (self._X2GoClient__get_session_server_hostname(s_hash), _peername[0], _peername[1])
        print "Username: %s" % self._X2GoClient__get_session_username(s_hash)
        print 
        session_infos = self._X2GoClient__list_sessions(s_hash)
        if session_infos:
            for session_info in session_infos.values():
                print "Session Name: %s" % session_info
                print "-------------"
                print "cookie: %s" % session_info.cookie
                print "agent PID: %s" % session_info.agent_pid
                print "display: %s" % session_info.display
                print "status: %s" % session_info.status
                print "graphic port: %s" % session_info.graphics_port
                print "snd port: %s" % session_info.snd_port
                print "sshfs port: %s" % session_info.sshfs_port
                print "username: %s" % session_info.username
                print "hostname: %s" % session_info.hostname
                # TODO: turn into datetime object
                print "create date: %s" % session_info.date_created
                # TODO: turn into datetime object
                print "suspended since: %s" % session_info.date_suspended
                print
        else:
            print "No running/suspended sessions found on X2Go server."
            print

    def list_desktops(self, s_hash):
        """\
        List up server-side available desktop sessions that are
        available for sharing to the logged in user.

        The list of sharable desktop sessions is printed to stdout.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # retrieve a desktop list
        print
        print "X2Go desktops available for sharing"
        print "==================================="
        _peername = self._X2GoClient__get_session_server_peername(s_hash)
        print "Host: %s - [%s]:%s" % (self._X2GoClient__get_session_server_hostname(s_hash), _peername[0], _peername[1])
        print "Username: %s" % self._X2GoClient__get_session_username(s_hash)
        print 
        desktop_list = self._X2GoClient__list_desktops(s_hash)
        if desktop_list:
            for desktop_name in desktop_list:
                print desktop_name
        else:
            print "No X2Go desktop sessions found that are available for desktop sharing."
            print

    def list_profiles(self):
        """\
        List up client-side configured session profile configurations.

        The list of session profiles is printed to stdout.

        """
        # retrieve a session list
        print
        print "Available X2Go session profiles"
        print "==============================="
        _profiles = self._X2GoClient__get_profiles()
        for _profile_id in _profiles.profile_ids:
            _profile_config = _profiles.get_profile_config(_profile_id)
            _session_params = _profiles.to_session_params(_profile_id)
            print _profile_config
            print _session_params
            print 

    def clean_sessions(self, s_hash):
        """\
        Clean all running/suspended sessions on the X2Go server that the session with UUID C{s_hash}
        is connected to.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        _profiles = self._X2GoClient__get_profiles()
        if self.args.session_profile and _profiles.has_profile(self.args.session_profile):
            _server = _profiles.get_session_param(self.args.session_profile, 'server')
        else:
            _server = self.args.server
        # clean all sessions from X2Go server
        self.logger('cleaning up all running sessions from X2Go server: %s' % _server, loglevel=x2go.loglevel_NOTICE, )
        self._X2GoClient__clean_sessions(s_hash)

    def new_session(self, s_hash):
        """\
        Launch a new X2Go session via the C{X2GoSession} object with session UUID C{s_hash}.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # start a new session and run a command
        self.logger('starting a new X2Go session', loglevel=x2go.loglevel_INFO, )
        if self.args.session_profile is None:
            self.logger('command for new session is: %s' % self.args.command, loglevel=x2go.loglevel_DEBUG, )
        else:
            self.logger('command from session profile to run is: %s' % self.session_registry(self.x2go_session_hash).get_session_cmd(), loglevel=x2go.loglevel_DEBUG, )
        self._X2GoClient__start_session(s_hash)

    def _auto_resume_newest(self, s_hash):
        session_infos = self._X2GoClient__list_sessions(s_hash)
        if session_infos:
            newest_session_name = x2go.utils.session_names_by_timestamp(session_infos)[-1]
            self._pyhoca_logger("Resuming newest X2Go session %s..." % newest_session_name, loglevel=x2go.loglevel_NOTICE, )
            self._X2GoClient__resume_session(s_hash, newest_session_name)
            return True
        return False

    def _auto_resume_oldest(self, s_hash):
        session_infos = self._X2GoClient__list_sessions(s_hash)
        if session_infos:
            oldest_session_name = x2go.utils.session_names_by_timestamp(session_infos)[0]
            self._pyhoca_logger("Resuming oldest X2Go session %s..." % oldest_session_name, loglevel=x2go.loglevel_NOTICE, )
            self._X2GoClient__resume_session(s_hash, oldest_session_name)
            return True
        return False

    def resume_session(self, s_hash):
        """\
        Use C{X2GoSession} object with session UUID C{s_hash} to resume a server-side suspended X2Go session.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # resume a running session
        self.logger('resuming X2Go session: %s' % self.args.resume, loglevel=x2go.loglevel_INFO, )
        available_sessions = self._X2GoClient__list_sessions(s_hash)
        if available_sessions and self.args.resume == 'OLDEST':
            self._auto_resume_oldest(s_hash)
        elif available_sessions and self.args.resume == 'NEWEST':
            self._auto_resume_newest(s_hash)
        elif self.args.resume in available_sessions.keys():
            self._X2GoClient__resume_session(s_hash, self.args.resume)
        else:
            self._runtime_error('requested session not available on X2Go server [%s]:%s' % (self.args.server, self.args.remote_ssh_port), exitcode=20)

    def share_desktop_session(self, s_hash):
        """\
        Use C{X2GoSession} object with session UUID C{s_hash} to share a server-side X2Go desktop session that is available for sharing.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # start a new session and run a command
        _desktop = self.args.share_desktop
        _share_mode = self.args.share_mode
        self.logger('sharing X2Go session: %s' % _desktop, loglevel=x2go.loglevel_INFO, )
        try:
            self._X2GoClient__share_desktop_session(s_hash, desktop=_desktop, share_mode=_share_mode)
        except x2go.X2GoDesktopSharingException, e:
            self._runtime_error('%s' % str(e), exitcode=54)


    def suspend_session(self, s_hash):
        """\
        Suspend X2Go Session object with session UUID C{s_hash}.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # send a suspend request to a session
        self.logger('requesting X2Go session suspend of session: %s' % self.args.suspend, loglevel=x2go.loglevel_INFO, )
        available_sessions = self._X2GoClient__list_sessions(s_hash)

        if self.args.suspend == "ALL":
            _session_names = available_sessions
        else:
            _session_names = [ self.args.suspend ]

        for _session_name in _session_names:
            if _session_name in available_sessions.keys():
                self._X2GoClient__suspend_session(s_hash, _session_name)
                self._pyhoca_logger("X2Go session %s has been suspended" % _session_name, loglevel=x2go.loglevel_NOTICE, )
            else:
                _server = self.args.server
                _remote_ssh_port = self.args.remote_ssh_port
                if not self.args.server and self.args.session_profile:
                    _server = self.session_profiles.get_value(self.session_profiles.to_profile_id(self.args.session_profile), 'host')
                    _remote_ssh_port = self.session_profiles.get_value(self.session_profiles.to_profile_id(self.args.session_profile), 'sshport')
                self._runtime_error('session %s not available on X2Go server [%s]:%s' % (_session_name, _server, _remote_ssh_port), exitcode=21)

    def terminate_session(self, s_hash):
        """\
        Terminate X2Go Session with session UUID C{s_hash}.

        @param s_hash: session UUID
        @type s_hash: C{str}

        """
        # send a terminate request to a session
        self.logger('requesting X2Go session terminate of session: %s' % self.args.terminate, loglevel=x2go.loglevel_INFO, )
        available_sessions = self._X2GoClient__list_sessions(s_hash)

        if self.args.terminate == "ALL":
            _session_names = available_sessions
        else:
            _session_names = [ self.args.terminate ]

        for _session_name in _session_names:
            if _session_name in available_sessions.keys():
                self._X2GoClient__terminate_session(s_hash, _session_name)
                self._pyhoca_logger("X2Go session %s has been terminated" % _session_name, loglevel=x2go.loglevel_NOTICE, )
            else:
                _server = self.args.server
                _remote_ssh_port = self.args.remote_ssh_port
                if not self.args.server and self.args.session_profile:
                    _server = self.session_profiles.get_value(self.session_profiles.to_profile_id(self.args.session_profile), 'host')
                    _remote_ssh_port = self.session_profiles.get_value(self.session_profiles.to_profile_id(self.args.session_profile), 'sshport')
                self._runtime_error('session %s not available on X2Go server [%s]:%s' % (_session_name, _server, _remote_ssh_port), exitcode=22)

    def __init__(self, args, logger=None, liblogger=None):
        """\
        L{PyHocaCLI} Class constructor.

        @param args: a class with properties representing the command-line options that are available to L{PyHocaCLI} instances.
        @type args: C{argparse.ArgumentParser} (or similar)
        @param logger: you can pass an L{X2GoLogger} object to the
            L{PyHocaCLI} constructor for logging application events
        @type logger: Python X2Go C{X2GoLogger} instance
        @param liblogger: you can pass an L{X2GoLogger} object to the
            L{PyHocaCLI} constructor for logging application events, this object is forwarded to the C{X2GoClient} class in Python X2Go 
        @type liblogger: Python X2Go C{X2GoLogger} instance

        """
        self.args = args
        if logger is None:
            self._pyhoca_logger = x2go.X2GoLogger(tag='PyHocaCLI')
        else:
            self._pyhoca_logger = copy.deepcopy(logger)
            self._pyhoca_logger.tag = 'PyHocaCLI'

        # initialize the X2GoClient context and start the connection to the X2Go server
        self._pyhoca_logger('preparing requested X2Go session', loglevel=x2go.loglevel_NOTICE, )

        x2go.X2GoClient.__init__(self, logger=liblogger)

        _profiles = self._X2GoClient__get_profiles()
        if self.args.session_profile and not _profiles.has_profile(self.args.session_profile):
            self._runtime_error('no such session profile of name: %s' % (self.args.session_profile), exitcode=31)

        self.auth_attempts = int(self.args.auth_attempts)

        if self.args.session_profile:

            _cmdlineopt_to_sessionopt = {
                'command': 'cmd',
                'kb_layout': 'kblayout',
                'kb_type': 'kbtype',
                'sound': 'snd_system',
                'ssh_privkey': 'key_filename',
                'server': 'hostname',
                'remote_ssh_port': 'port',
            }

            # override session profile options by option values from the arg parser
            kwargs={}
            if hasattr(self.args, 'parser'):
                for a, v in self.args._get_kwargs():
                    if v != self.args.parser.get_default(a):
                        try:
                            kwargs[_cmdlineopt_to_sessionopt[a]] = v
                        except KeyError:
                            kwargs[a] = v

            # setup up the session profile based X2Go session
            self.x2go_session_hash = self._X2GoClient__register_session(profile_name=self.args.session_profile,
                                                                        known_hosts=ssh_known_hosts_filename,
                                                                        **kwargs
                                                                       )

        else:

            # setup up the manually configured X2Go session
            self.x2go_session_hash = self._X2GoClient__register_session(args.server, port=int(self.args.remote_ssh_port), 
                                                           known_hosts=ssh_known_hosts_filename,
                                                           username=self.args.username,
                                                           key_filename=self.args.ssh_privkey,
                                                           add_to_known_hosts=self.args.add_to_known_hosts,
                                                           profile_name = 'Pyhoca-Client_Session',
                                                           session_type=self.args.session_type,
                                                           link=self.args.link,
                                                           geometry=self.args.geometry,
                                                           pack=self.args.pack,
                                                           cache_type='unix-kde',
                                                           kblayout=self.args.kbd_layout,
                                                           kbtype=self.args.kbd_type,
                                                           snd_system=self.args.sound,
                                                           printing=self.args.printing,
                                                           print_action=self.args.print_action,
                                                           print_action_args=self.args.print_action_args,
                                                           share_local_folders=self.args.share_local_folders,
                                                           cmd=self.args.command)

    def authenticate(self):
        """\
        Authenticate this L{PyHocaCLI} instance with the remote X2Go server.

        """
        connected = False
        force_password_auth = False
        _username = self.args.username or self._X2GoClient__get_session_username(self.x2go_session_hash)
        try:

            _auth_count = self.auth_attempts +1
            while not connected and _auth_count:
                try:
                    self._X2GoClient__connect_session(self.x2go_session_hash, username=_username, password=self.args.password, force_password_auth=force_password_auth)
                    connected = True
                except x2go.AuthenticationException, e:
                    force_password_auth = True
                    self._pyhoca_logger('interactive login for user ,,%s\'\'' % _username, loglevel=x2go.loglevel_NOTICE, )
                    self.args.password = getpass.getpass()
                except x2go.PasswordRequiredException, e:
                    self._pyhoca_logger('interactive login for user ,,%s\'\'' % _username, loglevel=x2go.loglevel_NOTICE, )
                    self.args.password = getpass.getpass()
                except x2go.BadHostKeyException:
                    self._runtime_error('SSH host key verification for remote host [%s]:%s failed' % (self.args.server, self.args.remote_ssh_port), exitcode=-254)
                except x2go.SSHException, e:
                    if str(e) != 'not a valid DSA private key file':
                        self._runtime_error(str(e), exitcode=253)
                    force_password_auth = True

                # workaround for Python bug: http://bugs.python.org/issue11236
                try:
                    if self.args.password is not None and '\x03' in self.args.password:
                        raise KeyboardInterrupt()
                except KeyboardInterrupt:
                    self._runtime_error('Authentication cancelled by user by hitting Ctrl-C at password prompt', exitcode=-200)

                _auth_count -= 1

                if not connected and not _auth_count:
                    if self.auth_attempts >= 2:
                        self._runtime_error('authentication failed, too many failures during interactive login', exitcode=-201)
                    elif self.auth_attempts == 1:
                        self._runtime_error('interactive authentication failed', exitcode=-202)
                    else:
                        self._runtime_error('non-interactive authentication failed', exitcode=-203)

        except socket.error, e:
            self._runtime_error('a socket error occured while establishing the connection: %s' % str(e), exitcode=-245)

    def MainLoop(self):
        """\
        Start the main loop of this application.

        """
        try:

            if self.args.clean_sessions:
                self.clean_sessions(self.x2go_session_hash)

            # go through the possible X2Go client modes
            if self.args.list_sessions:
                # print a beautified session list for the user
                self.list_sessions(self.x2go_session_hash)
                sys.exit(0)

            if self.args.list_desktops:
                # print a beautified desktop list for the user
                self.list_desktops(self.x2go_session_hash)
                sys.exit(0)

            if self.args.list_profiles:
                # print a beautified profile list for the user
                self.list_profiles()
                sys.exit(0)

            if self.args.resume:
                self.resume_session(self.x2go_session_hash)

            if self.args.share_desktop:
                self.share_desktop_session(self.x2go_session_hash)

            elif self.args.suspend:
                self.suspend_session(self.x2go_session_hash)

            elif self.args.terminate:
                self.terminate_session(self.x2go_session_hash)

            elif self.args.new:
                self.new_session(self.x2go_session_hash)

            if not (self.args.new or self.args.resume or self.args.share_desktop or self.args.session_profile):
                sys.exit(0)

            if self.args.resume or self.args.new:
                self._pyhoca_logger("give the X2Go session some time to come up...", loglevel=x2go.loglevel_NOTICE, )

            i=0
            while 0 < self.get_session(self.x2go_session_hash).get_progress_status() < 100:
                time.sleep(1)
                i+=1

            if self._X2GoClient__session_ok(self.x2go_session_hash):

                profile_name = self._X2GoClient__get_session_profile_name(self.x2go_session_hash)
                session_name = self._X2GoClient__get_session_name(self.x2go_session_hash)
                self._pyhoca_logger("X2Go session is now running, the X2Go client's profile name is: %s" % profile_name, loglevel=x2go.loglevel_INFO, )
                self._pyhoca_logger("X2Go session name is: %s" % session_name, loglevel=x2go.loglevel_INFO, )

                if not self.args.from_stdin:

                    if self.args.share_desktop:
                        self._pyhoca_logger("Press CTRL+C to end desktop sharing for this session...", loglevel=x2go.loglevel_NOTICE, )
                    elif self.args.terminate_on_ctrl_c:
                        self._pyhoca_logger("Press CTRL+C to terminate the running session...", loglevel=x2go.loglevel_NOTICE, )
                    else:
                        self._pyhoca_logger("Press CTRL+C to suspend the running session...", loglevel=x2go.loglevel_NOTICE, )

                try:

                    session_duration = 0
                    mounted = False

                    while self._X2GoClient__session_ok(self.x2go_session_hash):
                        time.sleep(2)
                        session_duration +=2

                    # wait a little while before telling the user what has happened
                    time.sleep(2)

                    # refresh session status so we can be most accurate on what we report below
                    self._X2GoClient__list_sessions(self.x2go_session_hash)

                    # report about the session status once we get here...
                    if self._X2GoClient__has_session_terminated(self.x2go_session_hash):
                        self._pyhoca_logger("X2Go session %s has terminated" % session_name, loglevel=x2go.loglevel_NOTICE, )
                    elif self._X2GoClient__is_session_suspended(self.x2go_session_hash):
                        self._pyhoca_logger("X2Go session %s has been suspended" % session_name, loglevel=x2go.loglevel_NOTICE, )
                    elif self._X2GoClient__is_session_running(self.x2go_session_hash):
                        self._pyhoca_logger("X2Go session %s has been moved to a different screen" % session_name, loglevel=x2go.loglevel_NOTICE, )

                except KeyboardInterrupt:
                    if self.args.share_desktop:
                        self._pyhoca_logger("Terminating X2Go shared desktop session %s" % session_name, loglevel=x2go.loglevel_INFO, )
                        self._X2GoClient__terminate_session(self.x2go_session_hash)
                        time.sleep(2)
                        self._pyhoca_logger("X2Go session %s has been terminated" % session_name, loglevel=x2go.loglevel_NOTICE, )
                    elif self.args.terminate_on_ctrl_c:
                        self._pyhoca_logger("Terminating X2Go session %s" % session_name, loglevel=x2go.loglevel_INFO, )
                        self._X2GoClient__terminate_session(self.x2go_session_hash)
                        # giving nxproxy's SSH tunnel some time to settle
                        time.sleep(2)
                        self._pyhoca_logger("X2Go session %s has been terminated" % session_name, loglevel=x2go.loglevel_NOTICE, )
                    else:
                        self._pyhoca_logger("Suspending X2Go session %s" % session_name, loglevel=x2go.loglevel_INFO, )
                        self._X2GoClient__suspend_session(self.x2go_session_hash)
                        # giving nxproxy's SSH tunnel some time to settle
                        time.sleep(2)
                        self._pyhoca_logger("X2Go session %s has been suspended" % session_name, loglevel=x2go.loglevel_NOTICE, )

        except x2go.X2GoSessionException, e:
            self._pyhoca_logger("X2GoSessionException occured:", loglevel=x2go.loglevel_ERROR)
            self._pyhoca_logger("-> %s" % str(e), loglevel=x2go.loglevel_ERROR)
