#!/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.

""" A (not so) simple command line wrapper.
"""

"""
This tool expect a command line, run it, output stdout and stderr in a
demuxable-way, and give-up with the command line exit code
For exemple, "ls /mnt /nothing" whould give you:
***************************************
ls: /nothing: No such file or directory
/mnt:
nfs
***************************************
and "pulse2_output_wrapper 0 ls /mnt /nothing" will give you:
***************************************
1199893195.42 C: ls·/mnt·/nothing
1199893195.46 E: ls: /nothing: No such file or directory
1199893195.46 E:
1199893195.46 O: /mnt:
1199893195.46 O: nfs
1199893195.46 O:
1199893195.46 X: 1
***************************************
Output lines are formatted like this:
<timestamp><space><tag><column><space><output line>
With <tag>:
  W: the wrapper line, using PULSE2_WRAPPER_ARG_SEPARATOR_STR as separator
  C: the command line, using PULSE2_WRAPPER_ARG_SEPARATOR_STR as separator
  O: something on stdout
  E: something on stderr
  X: the exit code

Lines are always parsed to be fully utf-8 compliant
"""
import twisted.internet.protocol
import twisted.internet.reactor
import twisted.python.usage

import time
import sys
import os
import signal
import re

from pulse2.consts import PULSE2_WRAPPER_ERROR_FAILURE, PULSE2_WRAPPER_ERROR_SIGNAL_BASE, PULSE2_WRAPPER_ERROR_PRECHECK, PULSE2_WRAPPER_ERROR_PRECHECK_BASE, PULSE2_WRAPPER_ERROR_PRECHECK_TABLE
from pulse2.consts import PULSE2_WRAPPER_ARG_SEPARATOR_STR

class ExitCode(object): # the only way I found to exit with the right exit code
    value = PULSE2_WRAPPER_ERROR_SIGNAL_BASE + signal.SIGKILL # prepare to be killed with SIGKILL

    def __new__(type):
        if not '_the_instance' in type.__dict__:
            type._the_instance = object.__new__(type)
        return type._the_instance
    def set(self, value):
        self.value = value
    def get(self):
        return self.value

class Options(twisted.python.usage.Options):

    exec_string = []
    thru_string = []
    check_command = []
    do_check = False

    optParameters = [
        ["exec", None, None, "The command to be executed (on the client or on the server, see below), arg separator is '%s'" % PULSE2_WRAPPER_ARG_SEPARATOR_STR],
        ["thru", None, None, "The command to wrap the command to execute (if executed on client, see below) and to wrap the client-check, arg separator is '%s', must respect the SSH syntax" % PULSE2_WRAPPER_ARG_SEPARATOR_STR],
        ["check-path", None, '/usr/bin/pulse2-client-check.sh', "The 'checker' path (on the client)"],
        ["action", None, None, "Tell to the client we are going to perform a particular action, f.e. 'INVENTORY', 'REBOOT'"],
        ["check-server-side", None, None, "Ask the client to show us some information for further check, f.e. IP=1.2.3.4 will ask the client to display IP, which will be compared against 1.2.3.4 by the wrapper"],
        ["check-client-side", None, None, "Ask the client to check for some informations, f.e. IP=1.2.3.4 will ask the client if ots IP is 1.2.3.4"],
        ["max-log-size", None, None, "Keeps the only first bytes of the output"],
        ["max-exec-time", None, None, "Kill command above this number of seconds"],
    ]

    optFlags = [
        ["examples", None, "show me some examples"],
        ["exec-server-side", None, "--exec arg is executed on server"],
        ["exec-client-side", None, "--exec arg is executed on client using --thru to handle the connection"],
        ["no-wrap", None, "Do not wrap output"],
        ["only-stdout", None, "Only write stdout on the output"],
        ["remove-empty-lines", None, "Do not write empty lines"],
    ]

    def __init__(self):
        twisted.python.usage.Options.__init__(self)

    def opt_version(self):
        rev = "$Rev$";
        print rev;

    def opt_examples(self):
        print 'here are some useful examples:'
        print
        print ' + a simple "ls /home" client side:'
        print '   pulse2_output_wrapper --exec ls·/home --thru ssh·root@client'
        print ' + same, cheking IP client-side:'
        print '   pulse2_output_wrapper --exec ls·/home --thru ssh·root@client --check-client-side IP=1.2.3.4'
        print ' + a simple "rsync /home" server side:'
        print '   pulse2_output_wrapper --exec rsync·/home·root@client:/home'
        print ' + same, cheking IP client-side:'
        print '   pulse2_output_wrapper --exec rsync·/home·root@client:/home --exec-server-side --thru ssh·root@client --check-client-side IP=1.2.3.4'
        sys.exit(0)

    def postOptions(self):
        if not self['exec']:
            print '%s: missing string to execute: --help for usage details.' % (sys.argv[0])
            sys.exit(2)
        self.exec_string = self['exec'].split(PULSE2_WRAPPER_ARG_SEPARATOR_STR)
        if self['thru']:
            self.thru_string = self['thru'].split(PULSE2_WRAPPER_ARG_SEPARATOR_STR)
        else:
            self.thru_string = None
        self.check_command = [self['check-path']]
        if self['check-server-side']:
            self.check_command += ['--show', ','.join(map(lambda x: x.split('=')[0], self['check-server-side'].split(',')))]
            self.do_check = True
        if self['action']:
            self.check_command += ['--action', self['action']]
            self.do_check = True
        if self['check-client-side']:
            self.check_command += ['--verify', self['check-client-side']]
            self.do_check = True
        if self.do_check and not self.thru_string:
            print '%s: check asked but missing string to pass thru: %s --help for usage details.' % (sys.argv[0], sys.argv[0])
            sys.exit(3)

class shOutputWrapper(twisted.internet.protocol.ProcessProtocol):
    """ A twisted.internet.protocol.ProcessProtocol override to handle on-the-fly encoding transcription
    """

    def __init__(self, config):
        self.exitcode = 0
        self.max_log_size = None
        self.config = config
        if config['max-log-size']:
            self.max_log_size = int(config['max-log-size'])
        self.size = 0
        self.output_exhausted = False
        self.end = False

    def truncate_to_max_log_size(self, line):
        if not self.max_log_size: # no trunc
            self.size += len(line)
            return line
        spaceleft = self.max_log_size-self.size
        self.size += len(line)
        return line[0:spaceleft]

    def outReceived(self, data):
        """ Format stdout
        """
        if self.end:
            return True
        for line in data.splitlines():
            if not self.output_exhausted:
                if self.size >= self.max_log_size: # limit exhausted
                    self.output_exhausted = True
                    format("Max buffer size reach, remaining output suppressed", 'W')
                else:
                    format(self.truncate_to_max_log_size(line), 'O')
        return True

    def errReceived(self, data):
        """ Format stderr
        """
        if self.end:
            return True
        for line in data.splitlines():
            if not self.output_exhausted:
                if self.size >= self.max_log_size: # limit exhausted
                    self.output_exhausted = True
                    format("Max buffer size reach, remaining output suppressed", 'W')
                else:
                    format(self.truncate_to_max_log_size(line), 'E')
        return True

    def processEnded(self, reason):
        """ Fired when process has finished
        """
        if reason.value.exitCode != None:
            # keeps the given exit code
            exitcode = int(reason.value.exitCode)
        else:
            # killed on signal, shift by PULSE2_WRAPPER_ERROR_SIGNAL_BASE
            exitcode = int(reason.value.signal) + PULSE2_WRAPPER_ERROR_SIGNAL_BASE

        format(str(exitcode), 'X')
        ExitCode().set(exitcode)
        if twisted.internet.reactor.running:
            twisted.internet.reactor.stop()

class shCheckCommand(twisted.internet.protocol.ProcessProtocol):
    """ Override some stuff from shOutputWrapper
    """
    def __init__(self, config):
        self.exitcode = 0
        self.config = config
        self.stdout = ""
        self.stderr = ""

    def processEnded(self, reason):
        """ Fired when process has finished
        """
        # compute good exitcode
        if reason.value.exitCode != None:
            # keeps the given exit code
            exitcode = int(reason.value.exitCode)
        else:
            # killed on signal, shift by PULSE2_WRAPPER_ERROR_SIGNAL_BASE
            exitcode = int(reason.value.signal) + PULSE2_WRAPPER_ERROR_SIGNAL_BASE

        if exitcode == 0: # continue in flow if everything went good
            if self.config['check-server-side']: # if something has to be checked server-side, do it
                for verif in self.config['check-server-side'].split(','): # iterate over checks to do
                    if not re.search('^%s$' % verif, self.stdout, re.M): # if not found, tell it and exit
                        if config.thru_string:
                            format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.thru_string), 'T')
                        format("PRE-COMMAND FAILED SERVER SIDE: expected result '%s' was not found !" % (verif), 'P')
                        format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.check_command), 'P')
                        for line in self.stdout.splitlines():
                            format(line, 'P')
                        format(str(self._getPrecommandExitCode(verif)), 'X')
                        ExitCode().set(self._getPrecommandExitCode(verif))
                        if twisted.internet.reactor.running:
                            twisted.internet.reactor.stop()
                        return
            runCommand(self.config)
        else: # something goes bad, print pre-command output then exits
            if config.thru_string:
                format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.thru_string), 'T')
            format("PRE-COMMAND FAILED CLIENT SIDE: exitcode = %s !" % (exitcode), 'P')
            format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.check_command), 'P')
            stdout_split = self.stdout.splitlines()
            stderr_split = self.stderr.splitlines()
            for line in stdout_split:
                format(line, 'P')
            for line in stderr_split:
                format(line, 'P')
            if exitcode == PULSE2_WRAPPER_ERROR_PRECHECK:
                # A failure in pre-command should return PULSE2_WRAPPER_ERROR_PRECHECK, else something bad
                # happened and we don't touch the exit code

                if len(stdout_split) > 0:
                    exitcode = self._getPrecommandExitCode(stdout_split.pop())
                elif len(stderr_split) > 0:
                    exitcode = self._getPrecommandExitCode(stderr_split.pop())
            format(str(exitcode), 'X')
            ExitCode().set(exitcode)
            if twisted.internet.reactor.running:
                twisted.internet.reactor.stop()
            return

    def errReceived(self, data):
        self.stderr += data

    def outReceived(self, data):
        self.stdout += data

    def _getPrecommandExitCode(self, check):
        """
        Return special known exit code for pre-command
        """
        exitcode = PULSE2_WRAPPER_ERROR_PRECHECK_BASE # to handle unknown conditions
        for label in PULSE2_WRAPPER_ERROR_PRECHECK_TABLE:
            if check.startswith(label):
                exitcode = PULSE2_WRAPPER_ERROR_PRECHECK_TABLE[label]
        return exitcode


def format(line, tag):
    """ Format a line as expected (see above)
    * attempt to guess encoding (currently try 'windows-1252', 'utf-8',
     'cp850', 'latin-1' and fallback to ascii)
    * convert to utf-8
    * format line
    """
    if config["only-stdout"] and not tag == 'O':
        return
    if config["remove-empty-lines"] and not len(line):
        return
    for codec in ['utf-8', 'cp850', 'latin-1', 'windows-1252']:
        if type(line) == type(unicode()): # input line can already be unicode
            uline = line
            break
        try:
            uline = unicode(line, codec)
            break
        except UnicodeDecodeError:
            continue
    else:
        uline = unicode(line, 'ascii', 'replace')
    try:
        line = uline.encode('utf-8')
    except:
        line = uline.encode('utf-8', 'replace')

    if not config["no-wrap"]:
        print "%.2f %s: %s" % (time.time(), tag, line)
    else:
        print line
        sys.stdout.flush()

def runCommand(config):
    if config.thru_string and not config['exec-server-side']: # handle command wrapping if needed
        format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.thru_string), 'T')
        format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.exec_string), 'C')
        headarg = config.thru_string[0]
        tailargs = config.thru_string + config.exec_string
    else:
        format(PULSE2_WRAPPER_ARG_SEPARATOR_STR.join(config.exec_string), 'C')
        headarg = config.exec_string[0]
        tailargs = config.exec_string

    # fork process
    if not os.access(headarg, os.X_OK):
        print '%s: %s can\'t be executed !' % (sys.argv[0], headarg)
        ExitCode().set(PULSE2_WRAPPER_ERROR_FAILURE)
        if twisted.internet.reactor.running:
            twisted.internet.reactor.stop()
        return
    handler = twisted.internet.reactor.spawnProcess(shOutputWrapper(config), headarg, tailargs, purged_env)

    # install termination handler to cleanly terminated sub-process in case of emergency
    twisted.internet.reactor.addSystemEventTrigger('before', 'shutdown', stopProcess, handler)

    # install timeout handler => process gets killed if above config['max-exec-time'] seconds
    if config['max-exec-time']:
        twisted.internet.reactor.callLater(int(config['max-exec-time']), stopProcess, handler) # max_exec_time handling

def runCheckCommand(config):
    if config.thru_string: # handle command wrapping if needed
        headarg = config.thru_string[0]
        tailargs = config.thru_string + config.check_command
    else:
        headarg = config.check_command[0]
        tailargs = config.check_command

    # fork process
    if not os.access(headarg, os.X_OK):
        print '%s: %s can\'t be executed !' % (sys.argv[0], headarg)
        ExitCode().set(PULSE2_WRAPPER_ERROR_FAILURE)
        if twisted.internet.reactor.running:
            twisted.internet.reactor.stop()
        return
    handler = twisted.internet.reactor.spawnProcess(shCheckCommand(config), headarg, tailargs, purged_env)

    # install termination handler to cleanly terminated sub-process in case of emergency
    twisted.internet.reactor.addSystemEventTrigger('before', 'shutdown', stopProcess, handler)

    # install timeout handler => process gets killed if above config['max-exec-time'] seconds
    if config['max-exec-time']:
        twisted.internet.reactor.callLater(int(config['max-exec-time']), stopProcess, handler)

def stopProcess(handler):
    if handler.pid:
        handler.signalProcess(signal.SIGKILL)
    # give the reactor a chance to terminate cleanly
    time.sleep(1)

# Parse command line
config = Options()
try:
    config.parseOptions()
except twisted.python.usage.UsageError, errortext:
    print '%s: %s' % (sys.argv[0], errortext)
    print '%s: Try --help for usage details.' % (sys.argv[0])
    sys.exit(PULSE2_WRAPPER_ERROR_FAILURE)

purged_env = os.environ # hold "secured" env vars

# catch ssh-agent auth socket if available for further usage
if os.getenv('SSH_AUTH_SOCK'):
    purged_env['SSH_AUTH_SOCK'] = os.getenv('SSH_AUTH_SOCK')

# delayed start to give the reactor a chance to install its signals handlers
if config.do_check: # if we have to execute a pre-command, do it, else jump to the real command
    twisted.internet.reactor.callWhenRunning(runCheckCommand, config)
else:
    twisted.internet.reactor.callWhenRunning(runCommand, config)

# fire the reactor
twisted.internet.reactor.run()
sys.exit(ExitCode().get())
