#!/usr/bin/python
# PSEXEC functionality over COM port
#
# Author:
#  kgermanov (@axcient.com)

import os
import cmd
import sys
import time
import fcntl
import string
import select
import signal
import logging
import argparse
import threading
import subprocess

STR_FILTER  = "MONITOR:"
ACCESS_DENIED_ERROR = 13

#--------------------------------------------
class StdOutWatcher:
    def __init__(self, filename, shell, stdredirect, timeout_start, timeout_execution, str_filter, timeout_step = 1):
        self.__timeout_step = int(timeout_step)
        self.__filename = filename
        self.__shell = shell
        self.__stdredirect = stdredirect
        self.__timeout_start = int(timeout_start)
        self.__timeout_execution = int(timeout_execution)
        self.__str_filter = str_filter
        self.__buffer = ''
        self.__exit_code = int(1)
        self.__total_time = int(0)
        self.__execution_time = int(0)
        self.__is_canceled = False
        self.__is_ready = threading.Event()

    def timeout_step(self):
        return self.__timeout_step

    def is_canceled(self):
        return self.__is_canceled

    def cancel(self):
        self.__is_canceled = True

    def filter_data(self, data):
        #logging.debug("input FILTER {%s} " %(data))

        output = self.__buffer + data
        self.__buffer = ''
        res = ''
        for line in output.splitlines( 1 ) :
            #logging.debug("FILTER1 {%s} {%s} %i" % (line, line[: len(self.__str_filter)], line.endswith('\n')))
            if line.endswith('\n'):
                if line.startswith(self.__str_filter):
                    self.handle_ctrl_codes(line[len(self.__str_filter):].rstrip())
                else:
                    res += line
            else:
                #logging.debug("FILTER2 {%s} {%s}" % (line, line[: len(self.__str_filter)]))
                if line.startswith(self.__str_filter[ : len(line)]):
                    self.__buffer = line
                else:
                    res += line

        return res

    def handle_ctrl_codes(self, ctrl):
        logging.debug("Have ctrl code:%s" %(ctrl))
        if ctrl.startswith("Result=") and self.__is_ready.is_set():
            param, value = ctrl.split("=", 1)
            self.set_exit_code(value)

        elif ctrl.startswith("Ready"):
            self.__execution_time = 0
            self.__is_ready.set()

    def start_check_health(self):
        check_health_th = threading.Thread(target=self.check_health)
        check_health_th.setDaemon(True)
        check_health_th.start()

    def check_health(self):
        while True:
            if self.__is_canceled:
                break

            #logging.debug("Check health, timers: %i %i %i" % (self.__total_time, self.__execution_time, os.path.exists(self.__filename)));

            self.__total_time += self.__timeout_step
            if not self.__is_ready.is_set() and 0 < self.__timeout_start and self.__timeout_start < self.__total_time:
                logging.error("Timeout of wait service");
                self.set_exit_code(123)
                break

            self.__execution_time += self.__timeout_step
            if self.__is_ready.is_set() and 0 < self.__timeout_execution and self.__timeout_execution < self.__execution_time:
                logging.error("Total time of validation greater then %i hours" % (self.__timeout_execution/60/60));
                self.set_exit_code(124)
                break

            if not os.path.exists(self.__filename):
                logging.error("Unexpected reboot detected");
                self.set_exit_code(125)
                break

            time.sleep(self.__timeout_step)

        # This piece of code normally shouldn't be executed,
        # but sometimes main thread is not terminated from the first attempt.
        # We are using infinite loop for sending SIGINT signal to avoid main thread hanging.
        # This thread (loop) will be terminated after terminating main thread.
        # See RB-4099 for details
        while True:
            time.sleep(5)
            logging.info("Terminating main shell")
            self.__shell.stop()

    def print_error(self, err_code):
        logging.debug("Handle return code:%i" % (err_code))
        if self.__exit_code == ACCESS_DENIED_ERROR:
            logging.error("Access denied. Wrong credentionals?");

    def set_exit_code(self, code):
        self.__exit_code = int(code)
        self.stop()
        self.print_error(self.__exit_code)

    def exit_code(self):
        return self.__exit_code

    def stop(self):
        self.cancel()
        self.__stdredirect.stop()
        self.__shell.stop()

    def wait_service_ready(self):
        logging.info("Wait until service is not ready")
        while not self.__is_ready.wait(1): # and self.__is_starting:
             pass

#--------------------------------------------
class RemoteStdRedirectPipe(threading.Thread):
    def __init__(self, filename, stream, encode='cp437'):
        threading.Thread.__init__(self)
        self.__stream = stream
        self.__watcher= None
        self.__encode = encode
        self.__is_stopped = threading.Event()
        self.__filehandle = self.open_non_blocked(filename)

    def open_non_blocked(self, filename):
        filehandle = open(filename, mode='rb', buffering=0)

        fd = filehandle.fileno()
        flag = fcntl.fcntl(fd, fcntl.F_GETFL)
        fcntl.fcntl(fd, fcntl.F_SETFL, flag | os.O_NONBLOCK)
        try:
           fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except Exception:
           raise Exception("Cannot lock %s. May be already attached?" % (filename))

        return filehandle

    def set_watcher(self, watcher):
        self.__watcher= watcher

    def stop(self):
        self.__is_stopped.set()

    def run(self):
        while not self.__is_stopped.is_set():
            ready_to_read, _, _ = select.select([self.__filehandle], [], [], self.__watcher.timeout_step())
            if not ready_to_read:
                continue

            ans = self.__filehandle.read()
            ans = self.__watcher.filter_data(ans)

            self.__stream.write(ans.decode(self.__encode))
            self.__stream.flush()

#--------------------------------------------
class RemoteShell(cmd.Cmd):
    def __init__(self, filename, encode='cp437'):
        cmd.Cmd.__init__(self)
        self.__encode = encode
        self.prompt = ''
        self.__filehandle = open(filename, mode='w')
        self.__stdin_encode = sys.stdin.encoding
        if not self.__stdin_encode:
            self.__stdin_encode = 'UTF-8'

    def emptyline(self):
        self.send_data('\r\n')
        return

    def default(self, line):
        try:
            self.send_data(line.decode(self.__stdin_encode).encode(self.__encode) + '\r\n')
        except Exception, err:
            logging.error( err )
            logging.warning("EXCEPTION: %s" % str(err), exc_info=True)


    def send_data(self, data):
        self.__filehandle.write(data)
        self.__filehandle.flush()

    def stop(self):
        os.kill(os.getpid(), signal.SIGINT)

#--------------------------------------------
class ComExec:
    def __init__(self, id, timeout, timeout_execution, force_handshake):
        self.__timeout = timeout
        self.__timeout_execution = timeout_execution
        self.__pty = self.get_pty(id)
        self.__force_handshake = force_handshake
        logging.info("PTY file: '%s'" % (self.__pty) )

    def run(self, cmd):
        stdout_pipe = RemoteStdRedirectPipe(self.__pty, sys.stdout)
        shell = RemoteShell(self.__pty)

        watcher = StdOutWatcher(self.__pty, shell, stdout_pipe, self.__timeout, self.__timeout_execution, STR_FILTER)
        stdout_pipe.set_watcher(watcher)

        res = 0
        try:
            watcher.start_check_health()

            stdout_pipe.start()

            if self.__force_handshake is True:
                shell.emptyline()
            watcher.wait_service_ready()

            logging.info("Try execute: '%s'" % (cmd) )
            shell.default(cmd)
            shell.cmdloop()

        except KeyboardInterrupt:
            logging.info("Interrupted")
            shell.default('exit')
            res = watcher.exit_code()

        finally:
            logging.info("End shell. Wait pipe...")
            watcher.cancel()
            if stdout_pipe.is_alive():
                stdout_pipe.stop()
                stdout_pipe.join(1)

            sys.stdout.flush()

        logging.info("Result: %s" % (res))
        return res

    def get_pty(self, id):
        if id.startswith("/dev/"):
             return id

        if self.is_xen():
             return self.get_xen_pty(id)

        return self.get_kvm_pty(id)

    def is_xen(self):
       return os.path.isfile('/sys/hypervisor/type')

    def get_xen_pty(self, id):
        if not id.isdigit():
            id = self.execute_shell("xl list | grep " + id + " | awk '{print $2;}'")

        if not id:
            raise Exception("Cannot find any VM with this id")

        return self.execute_shell("xenstore-read /local/domain/" + id + "/serial/0/tty")

    def get_kvm_pty(self, id):
        if not id.isdigit():
            id = self.execute_shell("virsh -c qemu:///system list | grep " + id + " | awk '{print $1;}'")

        if not id:
            raise Exception("Cannot find any VM with this id")

        return self.execute_shell("virsh -c qemu:///system ttyconsole " + id)


    def execute_shell(self, cmd):
        logging.debug("Try execute: %s" % (cmd))

        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
        (out, err) = proc.communicate()

        if proc.returncode != 0:
            raise Exception("Cannot execute command:%s. Return code:%i" % (cmd, proc.returncode) )

        return out.rstrip()

#--------------------------------------------
if __name__ == '__main__':
    parser = argparse.ArgumentParser(add_help = True, description = "PSEXEC like functionality example using ComSvc.")

    parser.add_argument('id', default = '1', help='domain id VM (int or as 1_4 or /dev/...)')
    parser.add_argument('commands', nargs='*', help='command to execute at the target (w/o path) - (default:cmd.exe)')
    parser.add_argument('--debug', action='store_true', help='Turn DEBUG output ON')
    parser.add_argument('--cmd-file', action="store", help='command file.')
    parser.add_argument('--timeout', default=1800, action='store', help='Set timeout to wait response service')
    parser.add_argument('--timeout-execution', default=0, action='store', help='Set timeout to execute command')
    parser.add_argument('--runas', metavar='{DOMAIN\\}USERNAME%PASSWORD', action='store', help='Run as user')
    parser.add_argument('--force-handshake', action='store_true', help='Send command to server without wait signal from server. Required if server start without "-t" parameter')

    if len(sys.argv)==1:
        parser.print_help()
        sys.exit(1)

    options = parser.parse_args()

    # Init the example's logger theme
    rootLogger = logging.getLogger()
    if options.debug is True:
        rootLogger.setLevel(logging.DEBUG)
    else:
        rootLogger.setLevel(logging.ERROR)

    logDir="/mnt/ax/scratch"
    if (not os.path.exists(logDir)) or (options.debug is True):
        handler = logging.StreamHandler(sys.stderr)
    else:
        handler = logging.FileHandler(logDir + "/comexe.log")

        handler_err = logging.StreamHandler(sys.stderr)
        handler_err.setLevel(logging.ERROR)
        rootLogger.addHandler(handler_err)

    handler.setFormatter(logging.Formatter('%(asctime)s [%(process)d] [%(levelname)s] %(message)s'))
    rootLogger.addHandler(handler)
    rootLogger.setLevel(logging.WARNING)

    try:
        command = 'cmd /Q'
        if options.commands:
            command = " ".join(options.commands)

        if options.cmd_file:
            with open(options.cmd_file, "r") as command_file:
                command = command_file.read()

        if options.runas:
            command = 'set runas ' + options.runas + '\r\n' + command

        executor = ComExec(options.id, options.timeout, options.timeout_execution, options.force_handshake)
        sys.exit(executor.run(command))

    except Exception, err:
        sys.stdout.flush()
        logging.error( err )
        logging.warning("MAIN EXCEPTION: %s" % str(err), exc_info=True)
        sys.exit(1)
