#!/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-2019 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
import asyncore
import socket
import getpass
import logging.config
import atexit
import configparser

if os.path.isdir('/run'):
    RUNDIR = '/run'
else:
    RUNDIR = '/var/run'

try:
    import daemon
    import lockfile
    CAN_DAEMONIZE = True
    pidfile = '{run}/x2gobroker/x2gobroker-loadchecker.pid'.format(run=RUNDIR)
    daemon_logdir = '/var/log/x2gobroker/'
except ImportError:
    CAN_DAEMONIZE = False

from pwd import getpwnam
from grp import getgrnam

PROG_NAME = os.path.basename(sys.argv[0])
PROG_OPTIONS = sys.argv[1:]
setproctitle.setproctitle("%s %s" % (PROG_NAME, " ".join(PROG_OPTIONS)))

from x2gobroker import __VERSION__
from x2gobroker import __AUTHOR__
import x2gobroker.loadchecker

global load_checker

class LoadCheckerServiceHandler(asyncore.dispatcher_with_send):

    def __init__(self, sock, logger=None):
        self.logger = logger
        asyncore.dispatcher_with_send.__init__(self, sock)
        self._buf = ''

    def handle_read(self):
        data = self._buf + self.recv(1024).decode()
        if not data:
            self.close()
            return
        reqs, data = data.rsplit('\n', 1)
        self._buf = data
        output = ""
        for req in reqs.split('\n'):
            backend, profile_id, hostname = req.split('\r')
            if self.logger: self.logger.debug('LoadCheckServiceHandler.handle_read(): received load check query: backend={backend}, profile_id={profile_id}, hostname={hostname}'.format(backend=backend, profile_id=profile_id, hostname=hostname))
            if hostname:
                load_factor = load_checker.get_server_load(backend, profile_id, hostname)
                if load_factor is not None:
                    output += "{lf}".format(lf=load_factor)
                    if self.logger: self.logger.info('LoadCheckServiceHandler.handle_read(): load check result for backend={backend}, profile_id={profile_id}, hostname={hostname}: {lf}'.format(backend=backend, profile_id=profile_id, hostname=hostname, lf=load_factor))
                else:
                    output += "LOAD-UNAVAILABLE"
                    if self.logger: self.logger.warning('LoadCheckServiceHandler.handle_read(): load check failure for backend={backend}, profile_id={profile_id}, hostname={hostname}: LOAD-UNAVAILABLE'.format(backend=backend, profile_id=profile_id, hostname=hostname))
            else:
                load_factors = load_checker.get_profile_load(backend, profile_id)
                if load_factors:
                    for h in list(load_factors.keys()):
                        if load_factors[h] is not None:
                            output +="{hostname}:{loadfactor}\n".format(hostname=h, loadfactor=load_factors[h])
                            if self.logger: self.logger.info('LoadCheckServiceHandler.handle_read(): load check result for backend={backend}, profile_id={profile_id}, hostname={hostname}: {lf}'.format(backend=backend, profile_id=profile_id, hostname=h, lf=load_factors[h]))
                        else:
                            output += "{hostname}:LOAD-UNAVAILABLE\n".format(hostname=h)
                            if self.logger: self.logger.warning('LoadCheckServiceHandler.handle_read(): load check failure for backend={backend}, profile_id={profile_id}, hostname={hostname}: LOAD-UNAVAILABLE'.format(backend=backend, profile_id=profile_id, hostname=h))
                else:
                    if self.logger: self.logger.warning('LoadCheckServiceHandler.handle_read(): load check failure for backend={backend}, profile_id={profile_id}: LOAD-UNAVAILABLE'.format(backend=backend, profile_id=profile_id))
                    output += "\n"
            self.send(output.encode())

    def handle_close(self):
        self.close()


class LoadCheckerService(asyncore.dispatcher_with_send):

    def __init__(self, socketfile, owner='root', group_owner='root', permissions='0o660', logger=None):
        self.logger = logger
        asyncore.dispatcher_with_send.__init__(self)
        self.create_socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.set_reuse_addr()
        self.bind(socketfile)
        try:
            os.chown(socketfile, getpwnam(owner).pw_uid, getgrnam(group_owner).gr_gid)
            os.chmod(socketfile, int(permissions, 8))
        except OSError:
            pass
        self.listen(1)

    def handle_accept(self):
        conn, _ = self.accept()
        LoadCheckerServiceHandler(conn, logger=self.logger)


def loop():
    ### the "loop" has two tasks...

    # 1. Do regular queries to remote X2Go Broker Agent instances to collect
    #    load average, CPU usage and type, memory usage, etc.
    load_checker.start()

    # 2. Provide a listening UNIX domain socket file that can be used for querying
    #    server states.
    asyncore.loop()


def cleanup_on_exit():
    os.remove(X2GOBROKER_LOADCHECKER_SOCKET)
    try: os.remove(pidfile)
    except: pass


# load the defaults.conf file, if present
iniconfig_loaded = None
iniconfig_section = '-'.join(PROG_NAME.split('-')[1:])
X2GOBROKER_DEFAULTS = "/etc/x2go/broker/defaults.conf"
if os.path.isfile(X2GOBROKER_DEFAULTS) and os.access(X2GOBROKER_DEFAULTS, os.R_OK):
    iniconfig = configparser.RawConfigParser()
    iniconfig.optionxform = str
    iniconfig_loaded = iniconfig.read(X2GOBROKER_DEFAULTS)

# normally this would go into defaults.py, however, we do not want to pull in defaults.py here as that will create
# unwanted logfiles (access.log, broker.log, error.log) when x2gobroker-loadchecker is installed as standalone service
if 'X2GOBROKER_DEBUG' in os.environ:
    X2GOBROKER_DEBUG = ( os.environ['X2GOBROKER_DEBUG'].lower() in ('1', 'on', 'true', 'yes', ) )
elif iniconfig_loaded and iniconfig.has_option(iniconfig_section, 'X2GOBROKER_DEBUG'):
    X2GOBROKER_DEBUG=iniconfig.get(iniconfig_section, 'X2GOBROKER_DEBUG')
elif iniconfig_loaded and iniconfig.has_option('common', 'X2GOBROKER_DEBUG'):
    X2GOBROKER_DEBUG=iniconfig.get('common', 'X2GOBROKER_DEBUG')
else:
    X2GOBROKER_DEBUG = False

if 'X2GOBROKER_DAEMON_USER' in os.environ:
    X2GOBROKER_DAEMON_USER=os.environ['X2GOBROKER_DAEMON_USER']
elif iniconfig_loaded and iniconfig.has_option(iniconfig_section, 'X2GOBROKER_DAEMON_USER'):
    X2GOBROKER_DAEMON_USER=iniconfig.get(iniconfig_section, 'X2GOBROKER_DAEMON_USER')
elif iniconfig_loaded and iniconfig.has_option('common', 'X2GOBROKER_DAEMON_USER'):
    X2GOBROKER_DAEMON_USER=iniconfig.get('common', 'X2GOBROKER_DAEMON_USER')
else:
    X2GOBROKER_DAEMON_USER="x2gobroker"

if 'X2GOBROKER_LOADCHECKER_LOGCONFIG' in os.environ:
    X2GOBROKER_LOADCHECKER_LOGCONFIG=os.environ['X2GOBROKER_LOADCHECKER_LOGCONFIG']
elif iniconfig_loaded and iniconfig.has_option(iniconfig_section, 'X2GOBROKER_LOADCHECKER_LOGCONFIG'):
    X2GOBROKER_LOADCHECKER_LOGCONFIG=iniconfig.get(iniconfig_section, 'X2GOBROKER_LOADCHECKER_LOGCONFIG')
elif iniconfig_loaded and iniconfig.has_option('common', 'X2GOBROKER_LOADCHECKER_LOGCONFIG'):
    X2GOBROKER_LOADCHECKER_LOGCONFIG=iniconfig.get('common', 'X2GOBROKER_LOADCHECKER_LOGCONFIG')
else:
    X2GOBROKER_LOADCHECKER_LOGCONFIG="/etc/x2go/broker/x2gobroker-loadchecker-logger.conf"

if 'X2GOBROKER_LOADCHECKER_SOCKET' in os.environ:
    X2GOBROKER_LOADCHECKER_SOCKET=os.environ['X2GOBROKER_LOADCHECKER_SOCKET']
elif iniconfig_loaded and iniconfig.has_option(iniconfig_section, 'X2GOBROKER_LOADCHECKER_SOCKET'):
    X2GOBROKER_LOADCHECKER_SOCKET=iniconfig.get(iniconfig_section, 'X2GOBROKER_LOADCHECKER_SOCKET')
elif iniconfig_loaded and iniconfig.has_option('common', 'X2GOBROKER_LOADCHECKER_SOCKET'):
    X2GOBROKER_LOADCHECKER_SOCKET=iniconfig.get('common', 'X2GOBROKER_LOADCHECKER_SOCKET')
else:
    X2GOBROKER_LOADCHECKER_SOCKET="{run}/x2gobroker/x2gobroker-loadchecker.socket".format(run=RUNDIR)


if __name__ == '__main__':

    common_options = [
        {'args':['-s','--socket-file'], 'default': X2GOBROKER_LOADCHECKER_SOCKET, 'metavar': 'LOADCHECKERSOCKET', 'help': 'socket file for LoadChecker communication', },
        {'args':['-o','--owner'], 'default': 'root', 'help': 'owner of the LoadChecker socket file', },
        {'args':['-g','--group'], 'default': 'root', 'help': 'group ownership of the LoadChecker socket file', },
        {'args':['-p','--permissions'], 'default': '0660', 'help': 'set these file permissions for the LoadChecker socket file', },
        {'args':['-d','--debug'], 'default': False, 'action': 'store_true', 'help': 'enable debugging code', },
        {'args':['-i','--debug-interactively'], 'default': False, 'action': 'store_true', 'help': 'force output of log message to the stderr (rather than to the log files)', },

    ]
    if CAN_DAEMONIZE:
        common_options.extend([
            {'args':['-D', '--daemonize'], 'default': False, 'action': 'store_true', 'help': 'Detach the X2Go Broker process from the current terminal and fork to background', },
            {'args':['-P', '--pidfile'], 'default': pidfile, 'help': 'Alternative file path for the daemon\'s PID file', },
            {'args':['-L', '--logdir'], 'default': daemon_logdir, 'help': 'Directory where log files for the process\'s stdout and stderr can be created', },
        ])
    p = argparse.ArgumentParser(description='X2Go Session Broker (Load Checker Service)',\
                                formatter_class=argparse.RawDescriptionHelpFormatter, \
                                add_help=True, argument_default=None)
    p_common = p.add_argument_group('common parameters')

    for (p_group, opts) in ( (p_common, common_options), ):
        for opt in opts:
            args = opt['args']
            del opt['args']
            p_group.add_argument(*args, **opt)

    cmdline_args = p.parse_args()

    # standalone daemon mode (x2gobroker-loadchecker as daemon) or interactive mode (called from the cmdline)?
    if getpass.getuser() in (X2GOBROKER_DAEMON_USER, 'root') and not cmdline_args.debug_interactively:

        # we run in standalone daemon mode, so let's use the system configuration for logging
        logging.config.fileConfig(X2GOBROKER_LOADCHECKER_LOGCONFIG)

        # create loadchecker logger
        logger_loadchecker = logging.getLogger('loadchecker')

    else:
        logger_root = logging.getLogger()
        stderr_handler = logging.StreamHandler(sys.stderr)
        stderr_handler.setFormatter(logging.Formatter(fmt='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt=''))

        # all loggers stream to stderr...
        logger_root.addHandler(stderr_handler)

        logger_loadchecker = logging.getLogger('loadchecker')
        logger_loadchecker.addHandler(stderr_handler)
        logger_loadchecker.propagate = 0

    if cmdline_args.debug_interactively:
        cmdline_args.debug = True

    # raise log level to DEBUG if requested...
    if cmdline_args.debug or X2GOBROKER_DEBUG:
        X2GOBROKER_DEBUG = True
        logger_loadchecker.setLevel(logging.DEBUG)

    logger_loadchecker.info('X2Go Session Broker ({version}), written by {author}'.format(version=__VERSION__, author=__AUTHOR__))
    logger_loadchecker.info('Setting up the Load Checker service\'s environment...')
    logger_loadchecker.info('  X2GOBROKER_DEBUG: {value}'.format(value=X2GOBROKER_DEBUG))
    logger_loadchecker.info('  X2GOBROKER_LOADCHECKER_SOCKET: {value}'.format(value=X2GOBROKER_LOADCHECKER_SOCKET))

    load_checker = x2gobroker.loadchecker.LoadChecker(logger=logger_loadchecker)

    if CAN_DAEMONIZE and cmdline_args.daemonize:

        # create directory for the PID file
        pidfile = os.path.expanduser(cmdline_args.pidfile)
        if not os.path.isdir(os.path.dirname(pidfile)):
            try:
                os.makedirs(os.path.dirname(pidfile))
            except:
                pass
        if not os.access(os.path.dirname(pidfile), os.W_OK) or (os.path.exists(pidfile) and not os.access(pidfile, os.W_OK)):
            print("")
            p.print_usage()
            print("Insufficent privileges. Cannot create PID file {pidfile} path".format(pidfile=pidfile))
            print("")
            sys.exit(-3)

        # create directory for logging
        daemon_logdir = os.path.expanduser(cmdline_args.logdir)
        if not os.path.isdir(daemon_logdir):
            try:
                os.makedirs(daemon_logdir)
            except:
                pass
        if not os.access(daemon_logdir, os.W_OK):
            print("")
            p.print_usage()
            print("Insufficent privileges. Cannot create directory for stdout/stderr log files: {logdir}".format(logdir=daemon_logdir))
            print("")
            sys.exit(-3)
        else:
            if not daemon_logdir.endswith('/'):
                daemon_logdir += '/'

    socket_file = cmdline_args.socket_file

    if os.path.exists(socket_file):
        os.remove(socket_file)

    if not os.path.exists(os.path.dirname(socket_file)):
        os.makedirs(os.path.dirname(socket_file))

    runtimedir_permissions = int(cmdline_args.permissions, 8)
    if runtimedir_permissions & 0o400: runtimedir_permissions = runtimedir_permissions | 0o100
    if runtimedir_permissions & 0o040: runtimedir_permissions = runtimedir_permissions | 0o010
    if runtimedir_permissions & 0o004: runtimedir_permissions = runtimedir_permissions | 0o001
    try:
        os.chown(os.path.dirname(socket_file), getpwnam(cmdline_args.owner).pw_uid, getpwnam(cmdline_args.group).pw_gid)
        os.chmod(os.path.dirname(socket_file), runtimedir_permissions)
    except OSError:
        pass

    LoadCheckerService(socket_file, owner=cmdline_args.owner, group_owner=cmdline_args.group, permissions=cmdline_args.permissions, logger=logger_loadchecker)
    atexit.register(cleanup_on_exit)
    try:
        if CAN_DAEMONIZE and cmdline_args.daemonize:
            keep_fds = [int(fd) for fd in os.listdir('/proc/self/fd') if fd not in (0,1,2) ]
            daemon_stdout = open(daemon_logdir+'x2gobroker-loadchecker.stdout', 'w+')
            daemon_stderr = open(daemon_logdir+'x2gobroker-loadchecker.stderr', 'w+')
            with daemon.DaemonContext(stdout=daemon_stdout, stderr=daemon_stderr, files_preserve=keep_fds, umask=0o027, pidfile=lockfile.FileLock(pidfile), detach_process=True):
                open(pidfile, 'w+').write(str(os.getpid())+"\n")
                loop()
        else:
            loop()
    except KeyboardInterrupt:
        pass
