#!/usr/bin/python
#
# -*- coding: utf-8; -*-
#
# (c) 2007-2008 Mandriva, http://www.mandriva.com/
#
# $Id$
#
# This file is part of Pulse 2, http://pulse2.mandriva.org
#
# Pulse 2 is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# Pulse 2 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Pulse 2; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

"""
    This is the Pulse2 scheduler manager
    The goal is to provide a helper to run scheduler:
    * start (fork)
    * stop (kill)
"""

# Big modules
import os
import os.path
import sys
import logging
import logging.config
import optparse
import signal
import time
import errno
import resource

time_to_wait_for_start = 2
time_to_wait_for_stop = 2
check_count_stop = 300

from mmc.site import mmcconfdir
from pulse2.scheduler.config import SchedulerConfig

def start_scheduler(config):
    """
    daemonize a scheduler
    """
    logger = logging.getLogger()
    name = config.name

    logger.info('daemonization started: scheduler %s' % name)

    pid_path = config.pid_path
    if not os.path.isdir(pid_path):
        os.makedirs(pid_path)

    pid = get_pidfile(config)
    if pid:
        if pid_is_alive(pid):
            logger.error('failed to start scheduler "%s": seems to be running under pid %d' % (name, pid))
            print 'failed to start scheduler "%s": seems to be running under pid %d' % (name, pid)
            return False
        else:
            clear_pidfile(config)
            logger.error('while starting scheduler "%s": clearing old pid file (pid was %d)' % (name, pid))

    # as a scheduler itself do not fork(), fork is done here
    # do the UNIX double-fork magic, see Stevens' "Advanced
    # Programming in the UNIX Environment" for details (ISBN 0201563177)
    try:
        pid = os.fork()
        if pid > 0:
            # first parent stays alive
            # let's check if it is realy alive
            # check is done in two stages:

            # let's do the first stage
            time.sleep(time_to_wait_for_start)
            pid = get_pidfile(config)
            if not pid_is_alive(pid): # os.WNOHANG make os.waitpid() to return (0, 0) if child is alive
                logger.error('daemonization part #2 failed: scheduler "%s" died in less than %d seconds' % (name, time_to_wait_for_start))
                clear_pidfile(config)
                return False

            # TODO: implement second stage here !
            return True
    except OSError, e:
        logger.error('daemonization part #1 failed: %d (%s)' % (e.errno, e.strerror))
        print >> sys.stderr, "daemonization part #1 failed: %d (%s)" % (e.errno, e.strerror)
        sys.exit(1)

    # decouple from parent environment
    os.close(sys.stdin.fileno())
    os.close(sys.stdout.fileno())
    os.close(sys.stderr.fileno())
    os.chdir("/")
    os.setsid()

    # do second fork
    try:
        pid = os.fork()
        if pid > 0:
            logger.info('daemonization succedeed: "%s", pid %d' % (name, pid) ) # exit from second parent, print eventual PID before
            set_pidfile(config, pid) # records the PID from here, as it is the only place where we know it
            sys.exit(0)
    except OSError, e:
        logger.error('daemonization part #2 failed: %d (%s)' % (e.errno, e.strerror))
        sys.exit(1)



    # drop privs
    os.seteuid(0)
    os.setegid(0)
    os.setgid(config.daemon_group)
    os.setuid(config.daemon_user)

    scheduler_path = config.scheduler_path
    logger.info('daemonization ended: switching to %s' % [scheduler_path])
    os.execv(scheduler_path, (scheduler_path,))

def stop_scheduler(config, wait_for_completion = False):
    name = config.name
    logger = logging.getLogger()
    os.seteuid(0)
    os.setegid(0)

    pid = get_pidfile(config)
    if pid:
        if pid_is_alive(pid):
            logger.info('attempt to kill scheduler "%s" (pid %d)' % (name, pid))
            kill_pid(pid, signal.SIGTERM) # TODO: check that kill was successful
            attempts_left = check_count_stop
            while True:
                time.sleep(time_to_wait_for_stop)
                if not pid_is_alive(pid): # os.WNOHANG make os.waitpid() to return (0, 0) if child is alive
                    logger.info('scheduler "%s" killed using SIGTERM' % name)
                    clear_pidfile(config)
                    return True
                elif wait_for_completion:
                    sys.stdout.softspace = 0 # prevent to print additionnal spaces
                    print ".",
                    sys.stdout.flush()
                    logger.warn('Still waiting for scheduler "%s" to stop' % name)
                    attempts_left -= 1
                    if attempts_left == 0:
                        logger.error('Giving up waiting for scheduler "%s" to stop' % name)
                        return False
                else:
                    logger.error('failed to kill scheduler "%s" using SIGTERM' % name)
                    return False
        else:
            logger.error('failed to stop scheduler "%s": was supposed to be running under pid %d' % (name, pid))
            print 'failed to stop scheduler "%s": was supposed to be running under pid %d' % (name, pid)
            clear_pidfile(config)
            return False
    else:
        logger.error('failed to stop scheduler "%s": was not running' % (name))
        print 'failed to stop scheduler "%s": was not running' % (name)
        clear_pidfile(config)
        return True

def clear_pidfile(config):
    os.setegid(0)
    os.seteuid(0)
    pid_file = os.path.join(config.pid_path, '%s.pid' % config.name)
    if os.path.isfile(pid_file):
        os.unlink(pid_file)
    os.setegid(config.daemon_group)
    os.seteuid(config.daemon_user)

def set_pidfile(config, pid):
    os.setegid(0)
    os.seteuid(0)
    pid_path = config.pid_path
    pid_file = os.path.join(config.pid_path, '%s.pid' % config.name)
    if not os.path.isdir(pid_path):
        os.makedirs(pid_path)
    os.system("echo %d > %s" % (pid, pid_file))
    os.setegid(config.daemon_group)
    os.seteuid(config.daemon_user)

def pid_is_alive(pid):
    return kill_pid(pid, 0)

def kill_pid(pid, signal):
    os.setegid(0)
    os.seteuid(0)
    try:
        os.killpg(os.getpgid(pid), signal)
    except OSError, err:
        if err.errno == errno.ESRCH:
            os.setegid(config.daemon_group)
            os.seteuid(config.daemon_user)
            return False
        elif err.errno == errno.EPERM: # WTF ? I *am* root, am I ?
            os.setegid(config.daemon_group)
            os.seteuid(config.daemon_user)
            return True
        else:
            os.setegid(config.daemon_group)
            os.seteuid(config.daemon_user)
            return False
    else:
        os.setegid(config.daemon_group)
        os.seteuid(config.daemon_user)
        return True

def get_pidfile(config):
    os.setegid(0)
    os.seteuid(0)
    name = config.name
    pid_file = os.path.join(config.pid_path, '%s.pid' % name)
    if not os.path.isfile(pid_file):
        os.setegid(config.daemon_group)
        os.seteuid(config.daemon_user)
        return False
    f = open(pid_file)
    for line in f:
        try:
            ret = int(line)
            f.close()
            os.setegid(config.daemon_group)
            os.seteuid(config.daemon_user)
            return ret
        except ValueError:
            pass
    os.setegid(config.daemon_group)
    os.seteuid(config.daemon_user)
    return False

def set_resource_usage(config):
    """
    Set scheduler process resources usage.
    We can only do it as root.
    """
    limits = config.setrlimit.split()
    while limits:
        try:
            res, soft, hard = limits[:3]
        except ValueError:
            print "Malformed resource limit expression, must be: limit soft hard (was '%s')" % config.setrlimit
            sys.exit(3)
        try:
            resource.setrlimit(eval("resource." + res), (int(soft), int(hard)))
        except Exception, e:
            print "Can't set resource limits for resource='%s', soft='%s', hard='%s'" % (res, soft, hard)
            print e
            sys.exit(3)
        limits = limits[3:]

def log_resource_usage(config):
    """
    Log resources usage limits that have been set.
    """
    limits = config.setrlimit.split()
    while limits:
        res, soft, hard = limits[:3]
        try:
            resource.getrlimit(eval("resource." + res))
            logging.getLogger().info("Resource '%s' set to soft='%s', hard='%s'" % (res, soft, hard))
        except Exception, e:
            print "Can't get resource limits for resource='%s'" % res
            print e
            sys.exit(4)
        limits = limits[3:]

if __name__ == "__main__":

    def __start():
        logger.info('Starting Pulse2 Scheduler')
        print('Starting Pulse2 Scheduler ...'),
        sys.stdout.flush()
        if start_scheduler(config):
            logger.info('Scheduler started')
            print('OK')
            return 0
        else:
            logger.warn('Failed to start scheduler')
            print('FAILED')
            return 1

    def __graceful_stop():
        logger.info('stopping Pulse2 Scheduler')
        print('Stopping Pulse2 Scheduler ...'),
        sys.stdout.flush()
        if stop_scheduler(config, wait_for_completion = False):
            logger.info('Scheduler stopped')
            print('OK')
            return 0
        else:
            logger.warn('Scheduler is still running, please wait a few minutes before restarting it')
            print('IN PROGRESS')
            return 1

    def __stop():
        logger.info('stopping Pulse2 Scheduler')
        print('Stopping Pulse2 Scheduler ...'),
        sys.stdout.flush()
        if stop_scheduler(config, wait_for_completion = True):
            logger.info('Scheduler stopped')
            print('OK')
            return 0
        else:
            logger.warn('Failed to stop scheduler')
            print('FAILED')
            return 1

    def __status():
        pid = get_pidfile(config)
        if not pid :
            print('Pulse2 Scheduler is stopped')
            return 3
        elif pid_is_alive(pid) :
            print('Pulse2 Scheduler (pid %s) is running...' % pid)
            return 0
        else:
            print('Pulse2 Scheduler dead but pid file exists')
            return 1

    parser = optparse.OptionParser()
    parser.add_option("-c", "--config-file", help='path to the config file',
                      default=mmcconfdir + '/pulse2/scheduler/scheduler.ini')
    parser.add_option("-o", "--id", help='id of this scheduler')
    (options, args) = parser.parse_args()

    if not os.path.exists(options.config_file):
        print "Config file '%s' does not exist." % options.config_file
        sys.exit(1)

    if len(args)!=1:
        print "command is missing"
        sys.exit(2)

    # we have to advance carefully here in order not to break the logging system
    # hypothesis: the logging folder already exists with proper permissions
    # instanciate object
    config = SchedulerConfig()

    # gather right user and password
    config.presetup(options.config_file)

    # Try to create the default log directory
    if not os.path.exists(config.logdir):
        try:
            os.mkdir(config.logdir)
        except:
            # Maybe we are not root, so it may not be a problem if the mkdir
            # failed
            pass

    if args[0] in ("start", "restart", "graceful-restart") :
        set_resource_usage(config)

    # set umask and effective UID and GID values RIGHT NOW
    os.umask(config.umask)
    os.setegid(config.daemon_group)
    os.seteuid(config.daemon_user)
    # we changed our identify a few lines above, log file will be created
    # with proper uid/gid/umask
    logging.config.fileConfig(options.config_file)

    if args[0] in ("start", "restart", "graceful-restart") :
        log_resource_usage(config)

    # last step: parse conf
    config.setup(options.config_file)

    logger = logging.getLogger()

    if args[0] == "start":
        sys.exit(__start())
    elif args[0] == "graceful-stop":
        sys.exit(__graceful_stop())
    elif args[0] == "stop":
        sys.exit(__stop())
    elif args[0] == "restart":
        ret = __stop()
        if ret != 0:
            sys.exit(ret)
        sys.exit(__start())
    elif args[0] == "graceful-restart":
        ret = __graceful_stop()
        if ret != 0:
            sys.exit(ret)
        sys.exit(__start())
    elif args[0] == "status":
        sys.exit(__status())
