Ahenk/opt/ahenk/base/command/commander.py

import configparser
import datetime
import json
import os
import queue as Queue
import threading

from base.Scope import Scope
from base.command.fifo import Fifo
from base.model.enum.ContentType import ContentType
from base.model.enum.MessageCode import MessageCode
from base.model.enum.MessageType import MessageType
from base.system.system import System
from base.util.util import Util


class Commander(object):
    def __init__(self):
        pass

    def set_event(self, *args):
        params = args[0]
        data = {}

        if System.Ahenk.is_running() is True:

            if params[1] == 'clean':
                print('Ahenk stopping')
                System.Process.kill_by_pid(int(System.Ahenk.get_pid_number()))
                self.clean()
                return False

            elif params[1] == 'login' and len(params) == 5:
                print('{1} {0}ing'.format(str(params[1]), str(params[2])))
                data['event'] = params[1]
                data['username'] = params[2]
                data['desktop'] = params[3]
                data['display'] = params[4]


            elif params[1] == 'logout' and len(params) == 3:
                print('{1} {0}ing'.format(str(params[1]), str(params[2])))
                data['event'] = params[1]
                data['username'] = params[2]

            elif params[1] == 'send' and len(params) > 5:
                data['event'] = params[1]
                response = {}
                response['timestamp'] = str(datetime.datetime.now().strftime("%d-%m-%Y %I:%M"))
                response['responseMessage'] = 'This content was sent via ahenk terminal command'

                if params[2] == '-t':
                    response['responseCode'] = MessageCode.TASK_PROCESSED.value
                    response['type'] = MessageType.TASK_STATUS.value
                    response['taskId'] = params[3]
                    if params[4] == '-m':
                        response['contentType'] = ContentType.TEXT_PLAIN.value
                        response['responseData'] = params[5]
                    elif params[4] == '-f':
                        if os.path.exists(str(params[5])):
                            response['contentType'] = self.get_relevant_type(str(params[5]))
                            response['responseData'] = Util.read_file(str(params[5]), 'rb')
                    else:
                        print('Wrong or missing parameter. Usage: send -t <task_id> -m|-f <message_content>|<file_path>')
                        return None

                    if len(params) > 6:
                        if params[6] == '-e':
                            response['responseCode'] = MessageCode.TASK_ERROR.value
                        elif params[6] == '-w':
                            response['responseCode'] = MessageCode.TASK_WARNING.value
                        elif params[6] == '-s':
                            response['responseCode'] = MessageCode.TASK_PROCESSED.value
                        else:
                            print('Wrong or missing parameter.(-e|-s|-w parameters are optional) Usage: send -t <task_id> -m|-f <message_content>|<file_path> -e|-s|-w')
                            return None

                elif params[2] == '-p' and len(params) > 7:
                    response['responseCode'] = MessageCode.POLICY_PROCESSED.value
                    response['type'] = MessageType.POLICY_STATUS.value
                    response['policyVersion'] = params[3]

                    if params[4] == '-c':
                        response['commandExecutionId'] = params[5]

                        if params[6] == '-m':
                            response['contentType'] = ContentType.TEXT_PLAIN.value
                            response['responseData'] = params[7]
                        elif params[6] == '-f':
                            if os.path.exists(str(params[7])):
                                response['contentType'] = self.get_relevant_type(str(params[7]))
                                response['responseData'] = Util.read_file(str(params[7]), 'rb')
                        else:
                            print('Wrong or missing parameter. Usage: send -p <policy_version> -c <command_execution_id> -m|-f <message_content>|<file_path>')
                            return None

                        if len(params) > 8:
                            if params[8] == '-e':
                                response['responseCode'] = MessageCode.POLICY_ERROR.value
                            elif params[8] == '-w':
                                response['responseCode'] = MessageCode.POLICY_WARNING.value
                            elif params[8] == '-s':
                                response['responseCode'] = MessageCode.POLICY_PROCESSED.value
                            else:
                                print('Wrong or missing parameter.(-e|-s|-w parameters are optional) Usage: send -p <policy_version> -c <command_execution_id> -m|-f <message_content>|<file_path> -e|-s|-w')
                                return None

                    else:
                        print('Wrong or missing parameter. Usage: send -p <policy_version> -c <command_execution_id> -m|-f <message_content>|<file_path> -e|-s|-w')
                        return None
                print('RESPONSE=' + str(response).replace("'", '"'))
                data['message'] = json.loads(str(response).replace("'", '"'))

            else:
                print('Wrong or missing parameter. Usage : %s start|stop|restart|status|clean|send')
                return None

        else:

            if params[1] == 'clean':
                self.clean()

            else:
                print('Ahenk not running!')
                return None

        if len(data) > 0:
            fifo = Fifo()
            thread = threading.Thread(target=fifo.push(str(json.dumps(data)) + '\n'))
            thread.start()

        return True

    def get_relevant_type(self, extension):

        extension = extension.lower()
        if extension == 'json':
            return ContentType.APPLICATION_JSON
        elif extension == 'txt':
            return ContentType.TEXT_PLAIN
        elif extension == 'dec':
            return ContentType.APPLICATION_MS_WORD
        elif extension == 'pdf':
            return ContentType.APPLICATION_PDF
        elif extension == 'xls':
            return ContentType.APPLICATION_VND_MS_EXCEL
        elif extension == 'jpeg' or extension == 'jpg':
            return ContentType.IMAGE_JPEG
        elif extension == 'png':
            return ContentType.IMAGE_PNG
        elif extension == 'html' or extension == 'htm':
            return ContentType.TEXT_HTML
        else:
            return ContentType.TEXT_PLAIN

    def get_event(self):
        fifo = Fifo()
        queue = Queue.Queue()
        thread = threading.Thread(target=fifo.pull(queue))
        thread.start()
        thread.join()
        result = queue.get()
        if result is not None:
            return result
        else:
            return None

    def clean(self):
        print('Ahenk cleaning..')
        try:
            config = configparser.ConfigParser()
            config._interpolation = configparser.ExtendedInterpolation()
            config.read(System.Ahenk.config_path())
            db_path = config.get('BASE', 'dbPath')

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

            config.set('CONNECTION', 'uid', '')
            config.set('CONNECTION', 'password', '')

            with open(System.Ahenk.config_path(), 'w') as file:
                config.write(file)
            file.close()
            print('Ahenk cleaned.')
        except Exception as e:
            print('Error while running clean command. Error Message {}'.format(str(e)))

    def status(self):
        ahenk_state = False

        if System.Ahenk.is_running() is True:
            ahenk_state = True
        return "Ahenk Active:{0}\nInstalled Plugins:{1}".format(ahenk_state, str(System.Ahenk.installed_plugins()))

    def force_clean(self):
        # TODO
        pass