#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

# This file is part of the  X2Go Project - https://www.x2go.org
# Copyright (C) 2012-2020 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 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:
            (success, remote_agent_tasks) = x2gobroker.agent.tasks_available(username=username, query_mode=query_mode, remote_agent=remote_agent)
        except x2gobroker.x2gobroker_exceptions.X2GoBrokerAgentException as 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 dict:
        print("\n".join(result))
        print()
    elif task.startswith('findbusyservers') and type(result) is dict:
        if result:
            print("\n".join([ "{host} -- {usage}%".format(host=host, usage=usage) for host, usage in list(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()
