#!/usr/bin/python3
# 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
import xml.etree.ElementTree as ElementTree


STR_FILTER = b'MONITOR:'
CMD_FILTER = STR_FILTER + b'cmd:'
CMD_FILTER_STR = CMD_FILTER.decode('utf-8')
ACCESS_DENIED_ERROR = 13
GENERAL_ERROR = 121
BREAKED_EXEC_ERROR = 122
TIMEOUT_WAITING_ERROR = 123
TIMEOUT_VALIDATION_ERROR = 124
UNEXPECTED_REBOOT_ERROR = 125
VM_NOT_FOUND_ERROR = 126


# --------------------------------------------
class StdOutWatcher:
    def __init__(self, filename, timeout_start, timeout_execution, str_filter, timeout_step=1):
        self.__timeout_step = int(timeout_step)
        self.__filename = filename
        self.__shell = None
        self.__stdredirect = None
        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()
        self.__is_cmd_starting = threading.Event()
        self.__details = []

    def set_stdredirect(self, stdredirect):
        self.__stdredirect = stdredirect

    def timeout_step(self):
        return self.__timeout_step

    def is_canceled(self):
        return self.__is_canceled

    def is_ready(self):
        return self.__is_ready.is_set()

    def cancel(self):
        self.__is_canceled = True

    def filter_data(self, data, pts_name, pts_shell_handle):
        # logging.debug("input FILTER {} {} ".format(self.__buffer.get(pts_name, b''), data))

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

        return res

    def handle_ctrl_codes(self, ctrl, pts_shell_handle):
        logging.debug("Have ctrl code: {}".format(ctrl))
        if ctrl.startswith(b'Result=') and self.__is_ready.is_set():
            param, value = ctrl.split(b'=', 1)
            logging.info("Got response: {}".format(ctrl))
            self.set_exit_code(value)
            self.__is_cmd_starting.clear()

        elif ctrl.startswith(b'Start'):
            self.__is_cmd_starting.set()
            logging.info("Cmd starting...")

        elif ctrl.startswith(b'Ready'):
            if self.__is_cmd_starting.is_set():
                self.set_exit_code(BREAKED_EXEC_ERROR, "Breaked execution process detected.")
                return

            if self.__shell is None:
                self.__shell = RemoteShell(pts_shell_handle)
                logging.info("Service started")

            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: {} {} {}".format
            #                (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:
                self.set_exit_code(TIMEOUT_WAITING_ERROR, "Timeout of waiting service.")
                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):
                self.set_exit_code(TIMEOUT_VALIDATION_ERROR, "Total time of validation greater than {} hours."
                                                             .format(self.__timeout_execution // 60 // 60))
                break

            if not os.path.exists(self.__filename):
                self.set_exit_code(UNEXPECTED_REBOOT_ERROR, "Unexpected reboot detected.")
                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 self.__shell is not None:
            time.sleep(5)
            logging.info("Terminating main shell")
            self.__shell.stop()

    def handle_error(self, err_code):
        logging.debug("Handling return code: {}".format(err_code))
        if self.__exit_code == ACCESS_DENIED_ERROR:
            self.__details.append("Access denied. Wrong credentials?")
            return
        elif self.__exit_code == TIMEOUT_VALIDATION_ERROR:
            return
        self.__details.append("Elapsed: {}s. Elapsed since start {}s".format(self.__execution_time, self.__total_time))

    def set_exit_code(self, code, details=""):
        self.__exit_code = int(code)
        self.stop()
        if details:
            self.__details.append(details)
        self.handle_error(self.__exit_code)

    def exit_code(self):
        return self.__exit_code

    def details(self):
        return " ".join(self.__details)

    def stop(self):
        self.cancel()
        if self.__stdredirect is not None:
            self.__stdredirect.stop()
        if self.__shell is not None:
            self.__shell.stop()

    def __wait_service_ready(self):
        logging.info("Waiting until service is not ready")
        while not self.__is_ready.wait(1):
            if self.__is_canceled:
                raise KeyboardInterrupt

    def __wait_cmd_starting(self, cmd):
        logging.info("Waiting until cmd is not starting")
        self.__shell.send_cmd(cmd)
        while not self.__is_cmd_starting.wait(1):
            if self.__is_canceled:
                raise KeyboardInterrupt

            self.__shell.send_cmd(cmd)

    def cmdloop(self, cmd):
        self.__wait_service_ready()
        logging.info("Trying to execute command: '{}'".format(cmd))
        if self.__shell is None:
            raise Exception("Shell is not setup. Is it ready?")

        try:
            self.__shell.send_cmd(cmd)
            self.__wait_cmd_starting(cmd)
            self.__shell.cmdloop()
        except KeyboardInterrupt:
            logging.info("Execution was interrupted by SIGINT from RemoteShell")
            self.__shell.send_exit()
            raise


# --------------------------------------------
class RemoteStdRedirectPipe(threading.Thread):
    def __init__(self, watcher, filenames, stream, force_handshake=False, no_filter=False, encode='cp437'):
        threading.Thread.__init__(self)
        self.__stream = stream
        self.__watcher = watcher
        self.__encode = encode
        self.__no_filter = no_filter
        self.__is_stopped = threading.Event()
        self.__filenames = filenames
        self.__filehandles = list(map(lambda fn: self.open_non_blocked(fn), filenames))
        self.__filenames = dict(zip(self.__filehandles, filenames))

        shell_handles = list(map(lambda fn: self.open_for_write(fn, force_handshake), filenames))
        self.__shell_handles = dict(zip(self.__filehandles, shell_handles))

        self.__watcher.set_stdredirect(self)

    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 {}. May be already attached?".format(filename))

        # Cleanup input buffer from garbage
        ready_to_read, _, _ = select.select([filehandle], [], [], 0)

        if ready_to_read:
            filehandle.read()

        return filehandle

    def open_for_write(self, filename, force_handshake):
        filehandle = open(filename, mode='wb')

        if force_handshake:
            filehandle.write(CMD_FILTER + b'\r\n')
            filehandle.write(b'exit' + b'\r\n')
            filehandle.flush()

        return filehandle

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

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

            for handle in ready_to_read:
                ans_no_filtered = handle.read()
                ans = self.__watcher.filter_data(ans_no_filtered, self.__filenames[handle],
                                                 self.__shell_handles[handle])
                if self.__no_filter:
                    ans = ans_no_filtered

                if self.__watcher.is_ready() or self.__no_filter:
                    try:
                        ans_decoded = ans.decode(self.__encode)
                        self.__stream.write(ans_decoded)
                    except:
                        self.__stream.write(ans)

                    self.__stream.flush()

                    if not is_ready and self.__watcher.is_ready():
                        is_ready = True
                        self.lock_to_handle(handle)
                        break

    def lock_to_handle(self, found_handle):
        filename = self.__filenames[found_handle]
        shell_handle = self.__shell_handles[found_handle]

        logging.info("Answered on {}".format(filename))

        self.__filehandles = [found_handle]
        self.__filenames = {found_handle: filename}
        self.__shell_handles = {found_handle: shell_handle}


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

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

    def send_cmd(self, line):
        self.default(CMD_FILTER_STR + line)

    def send_exit(self):
        self.default('exit')

    def default(self, line):
        try:
            if sys.version_info[0] > 2:
                self.send_data(line.encode(self.__encode) + b'\r\n')
            else:
                self.send_data(line.decode(self.__stdin_encode).encode(self.__encode) + b'\r\n')
        except Exception as err:
            logging.error(err)
            logging.warning("EXCEPTION: {}".format(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 VMNotFoundException(Exception):
    def __init__(self, id):
        super(VMNotFoundException, self).__init__("Cannot find any VM with id: " + id)


# --------------------------------------------
class ComExec:
    def __init__(self, id, timeout, timeout_execution, force_handshake, no_filter):
        self.__timeout = timeout
        self.__timeout_execution = timeout_execution
        self.__ptys = self.get_ptys(id)
        self.__force_handshake = force_handshake
        self.__no_filter = no_filter
        logging.debug("Use PTY files: '{}'".format(self.__ptys))

    def run(self, cmd):
        watcher = StdOutWatcher(self.__ptys[0], self.__timeout, self.__timeout_execution, STR_FILTER)
        stdout_pipe = RemoteStdRedirectPipe(watcher, self.__ptys, sys.stdout, self.__force_handshake, self.__no_filter)

        res = 0
        try:
            watcher.start_check_health()
            stdout_pipe.start()

            watcher.cmdloop(cmd)

        except KeyboardInterrupt:
            logging.info("Execution was interrupted by SIGINT from RemoteShell")
            res = watcher.exit_code()
        finally:
            logging.info("Ending shell. Waiting pipe...")
            watcher.cancel()
            if stdout_pipe.is_alive():
                stdout_pipe.stop()
                stdout_pipe.join(1)

            sys.stdout.flush()

        details = watcher.details()
        if res != 0:
            logging.error("Command execution failed with code {}. {}".format(watcher.exit_code(), details))
        else:
            logging.info("Command successfully executed. {}".format(details))

        return res

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

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

        return self.get_kvm_ptys(id)

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

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

        if not id:
            raise VMNotFoundException(input_id)

        serials = self.execute_shell("xenstore-list /local/domain/" + id + "/serial")
        serials = [x.strip() for x in serials if x.strip()]
        return list(map(lambda el: self.execute_shell("xenstore-read /local/domain/" + id + "/serial/" +
                                                      el + "/tty"), serials))

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

        if not id:
            raise VMNotFoundException(input_id)

        xml = self.execute_shell("virsh -c qemu:///system dumpxml " + id)
        serials_elements = ElementTree.fromstring(xml).findall(".//serial[target]/source/[@path]")
        return list(map(lambda el: el.get('path'), serials_elements))

    def execute_shell(self, cmd):
        logging.debug("Try to get PTY. Execute: {}".format(cmd))

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

        if proc.returncode != 0:
            raise Exception("Cannot execute command:{}. Return code:{}".format(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('--verbose', action='store_true', help='Turn INFO output ON')
    parser.add_argument('--cmd-file', action="store", help='command file.')
    parser.add_argument('--log-path', default='/mnt/ax/scratch', action='store',
                        help='Log file location - (default:/mnt/ax/scratch/)')
    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('--no-force-handshake', action='store_true', help='No send handshake to server on start')
    parser.add_argument('--no-filtered', action='store_true', help='Show no filtered data for debug')

    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)
    elif options.verbose is True:
        rootLogger.setLevel(logging.INFO)
    else:
        rootLogger.setLevel(logging.ERROR)

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

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

    handler.setFormatter(logging.Formatter('%(asctime)s [%(process)5s] [%(levelname)7s] %(message)s'))
    rootLogger.addHandler(handler)

    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' + CMD_FILTER_STR + command

        logging.info("Starting comexe for VM {}".format(options.id))

        executor = ComExec(options.id, options.timeout, options.timeout_execution,
                           not options.no_force_handshake, options.no_filtered)
        sys.exit(executor.run(command))
    except VMNotFoundException as err:
        sys.stdout.flush()
        logging.error(err)
        sys.exit(VM_NOT_FOUND_ERROR)
    except Exception as err:
        sys.stdout.flush()
        logging.error(err)
        logging.warning("MAIN EXCEPTION: {}".format(str(err)), exc_info=True)
        sys.exit(GENERAL_ERROR)