#!/usr/bin/python

# This file is part of the  X2Go Project - http://www.x2go.org
# Copyright (C) 2012-2015 by Mike Gabriel <mike.gabriel@das-netzwerkteam.de>
#
# X2Go Session Broker 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.
#
# X2Go Session Broker 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.

import os
import sys
import types
import setproctitle
import argparse
import logging
from paramiko import AutoAddPolicy

# perform an authentication against the authentication mechanism configured for WSGI
try:
    import x2gobroker.defaults
except ImportError:
    sys.path.insert(0, os.path.join(os.getcwd(), '..'))
    import x2gobroker.defaults
import x2gobroker.loggers

import x2gobroker.agent
from x2gobroker.utils import drop_privileges

PROG_NAME = os.path.basename(sys.argv[0])
PROG_OPTIONS = sys.argv[1:]
try:
    _password_index = PROG_OPTIONS.index('--password')+1
    PROG_OPTIONS[_password_index] = "XXXXXXXX"
except ValueError:
    # ignore if --password option is not specified
    pass
setproctitle.setproctitle("%s %s" % (PROG_NAME, " ".join(PROG_OPTIONS)))

if __name__ == "__main__":

    agent_options = [
        {'args':['--list-tasks'], 'default': False, 'action': 'store_true', 'help': 'List available broker agent tasks', },
        {'args':['-u','--username', '--user'], 'default': None, 'metavar': 'USERNAME', 'help': 'When testing the broker agent, test on behalf of this user (default: none)', },
        {'args':['-t','--task'], 'default': 'ping', 'metavar': 'AGENT_TASK', 'help': 'Perform this task on the (remote) broker agent', },
        {'args':['-H','--host'], 'default': 'LOCAL', 'metavar': 'HOSTNAME', 'help': 'Test X2Go Session Broker Agent on this (remote) host (default: LOCAL)', },
        {'args':['-p','--port'], 'default': 22, 'metavar': 'PORT', 'help': 'For remote agent calls (via SSH) use this port as SSH port (default: 22)', },
    ]
    supplementary_options = [
        {'args':['-A', '--add-to-known-hosts'], 'default': False, 'action': 'store_true', 'help': 'For SSH connections to a remote broker agent: permanently add the remote system\'s host key to the list of known hosts', },
        {'args':['--session-id'], 'default': None, 'help': 'when testing the \'suspendsession\' or the \'terminatesession\' task, you have to additionally give a session ID to test those tasks on', },
        {'args':['--pubkey'], 'default': None, 'metavar': 'PUBKEY_AS_STRING', 'help': 'use this public key when testing the \'addauthkey\' and the \'delauthkey\' tasks', },
    ]
    misc_options = [
        {'args':['-C','--config-file'], 'default': None, 'metavar': 'CONFIG_FILE', 'help': 'Specify a special configuration file name, default is: {default}'.format(default=x2gobroker.defaults.X2GOBROKER_CONFIG), },
        {'args':['-d','--debug'], 'default': False, 'action': 'store_true', 'help': 'enable debugging code', },
    ]
    p = argparse.ArgumentParser(description='X2Go Session Broker (Agent Test Utility)',\
                                formatter_class=argparse.RawDescriptionHelpFormatter, \
                                add_help=True, argument_default=None)
    p_agent = p.add_argument_group('agent parameters')
    p_supplimentary = p.add_argument_group('supplimentary parameters')
    p_misc  = p.add_argument_group('miscellaneous parameters')

    for (p_group, opts) in ( (p_agent, agent_options), (p_supplimentary, supplementary_options), (p_misc, misc_options), ):
        for opt in opts:
            args = opt['args']
            del opt['args']
            p_group.add_argument(*args, **opt)

    cmdline_args = p.parse_args()

    if os.getuid() != 0:
        p.print_help()
        print
        print "*** The {progname} tool needs to be run with super-user privileges... ***".format(progname=PROG_NAME)
        print
        sys.exit(-1)

    if cmdline_args.username is None and not cmdline_args.list_tasks and cmdline_args.task not in ('ping', 'checkload'):
        p.print_help()
        print
        print "*** Cannot continue without username... ***"
        print
        sys.exit(-1)

    if cmdline_args.task in ('suspendsession', 'terminatesession') and not cmdline_args.session_id:
        p.print_help()
        print
        print "*** Cannot continue on this task without a given session ID... ***"
        print
        sys.exit(0);

    if cmdline_args.config_file is not None:
        x2gobroker.defaults.X2GOBROKER_CONFIG = cmdline_args.config_file

    if cmdline_args.debug:
        x2gobroker.defaults.X2GOBROKER_DEBUG = cmdline_args.debug
        # raise log level to DEBUG if requested...
        if x2gobroker.defaults.X2GOBROKER_DEBUG and not x2gobroker.defaults.X2GOBROKER_TESTSUITE:
            x2gobroker.loggers.logger_broker.setLevel(logging.DEBUG)
            x2gobroker.loggers.logger_error.setLevel(logging.DEBUG)

    username = cmdline_args.username
    hostname = cmdline_args.host
    port     = cmdline_args.port
    task     = cmdline_args.task

    list_tasks = cmdline_args.list_tasks


local_agent = (hostname == 'LOCAL')
query_mode = local_agent and 'LOCAL' or 'SSH'
if local_agent: remote_agent = None
else: remote_agent = {'hostaddr': hostname, 'port': port, }

if remote_agent and cmdline_args.add_to_known_hosts:
    remote_agent.update({
        'host_key_policy': AutoAddPolicy(),
    })


def call_agent(task, **kwargs):
    try:
        _result = agent_client_tasks[task](username=username, query_mode=query_mode, remote_agent=remote_agent, **kwargs)
        if _result == True:
            print "The broker agent could be reached but the task returned no printable output (which probably is fine)."
            print
        return _result
    except KeyError:
        print "No such task: '{task_name}'. Use --list-tasks to view information about".format(task_name=task)
        print "available tasks."
        print
        return False

if __name__ == "__main__":

    # drop root privileges and run as X2GOBROKER_DAEMON_USER
    drop_privileges(uid=x2gobroker.defaults.X2GOBROKER_DAEMON_USER, gid=x2gobroker.defaults.X2GOBROKER_DAEMON_GROUP)

    print
    print "X2Go Session Broker (Agent Test Utility)"
    print "----------------------------------------"

    agent_client_tasks = x2gobroker.agent.tasks
    if 'availabletasks' in agent_client_tasks:
        try:
            remote_agent_tasks = x2gobroker.agent.tasks_available(username=username, query_mode=query_mode, remote_agent=remote_agent)
        except x2gobroker.x2gobroker_exceptions.X2GoBrokerAgentException, e:
            print "{errmsg}.".format(errmsg=e)
            print
            sys.exit(0)

    if not local_agent and not x2gobroker.agent.has_remote_broker_agent_setup():

        print "This instance of X2Go Session Broker is not able to contact any remote"
        print "X2Go Session Broker Agent instances. Check this broker's SSH setup!!!"
        print
        print "Aborting any futher tests..."
        sys.exit(-1)

    if list_tasks:
        print "The queried broker agent supports these tasks / features:"
        print
        for task in remote_agent_tasks:
            try:
                print "   {task_name}: {task_function_obj}".format(task_name=task, task_function_obj=agent_client_tasks[task])
            except KeyError:
                print "   {task_name}: not supported by this broker version".format(task_name=task)
        print
        sys.exit(0);

    kwargs = {}
    pubkey = cmdline_args.pubkey
    if (task == 'addauthkey' or task == 'delauthkey') and not pubkey:
        pubkey, privkey = x2gobroker.agent.genkeypair(local_username=username, client_address="localhost")

    if pubkey:
        kwargs.update({
            'pubkey_hash': pubkey,
        })

    if task in ('suspendsession', 'terminatesession'):
        kwargs.update({
            'session_name': cmdline_args.session_id
        })

    result = call_agent(task, **kwargs)
    if type(result) is types.ListType:
        print "\n".join(result)
        print
    elif task.startswith('findbusyservers') and type(result) is types.DictType:
        if result:
            print "\n".join([ "{host} -- {usage}%".format(host=host, usage=usage) for host, usage in result.items() ])
        else:
            print "X2Go Server busy state: All servers are idle."
            # even an empty dict means, that we have been successful...
            result = True
        print

    if task == 'addauthkey' and result:
        on_host = " on {host}".format(host=cmdline_args.host) if cmdline_args.host != 'LOCAL' else ""
        print "NOTE: This test-run added the below SSH public key to X2Go's authorized_keys file"
        print "      for user '{username}{on_host}'.".format(username=username, on_host=on_host)
        print
        print "      The file location for this normally is $HOME/.x2go/authorized_keys."
        print "      MAKE SURE TO REMOVE THIS KEY MANUALLY (or use test the 'delauthkey' task)!!!"
        print
        print pubkey
        print

    if task == 'delauthkey' and result:
        on_host = " on {host}".format(host=cmdline_args.host) if cmdline_args.host != 'LOCAL' else ""
        print "NOTE: This test-run attempted to remove all occurences of the below SSH public key from"
        print "      X2Go's authorized_keys file for user '{username}{on_host}'.".format(username=username, on_host=on_host)
        print
        print "      The file location for this normally is $HOME/.x2go/authorized_keys."
        print "      PLEASE DOUBLE-CHECK THE USER'S authorized_keys file MANUALLY!!!"
        print
        print pubkey
        print

    where = "local" if query_mode == "LOCAL" else "remote"
    if result:
        print "The task '{task_name}' could be executed successfully on the {where} broker agent.".format(task_name=task, where=where)
    else:
        print "The task '{task_name}' failed to execute on the {where} broker agent.".format(task_name=task, where=where)
    print
