<filename>knowledge/StateAction.py
#!/usr/bin/python3
from pymongo import MongoClient
import time
import sys
import json
import os
from KnowledgeUtils import KnowledgeUtils

class StateAction:
    def __init__(self, name, run_keys, result_keys):
        """ input_list - list of all possible trials that could be made. The framework will ignore those that were run alrady and those that would lead to achieving something that is __achieved already """
        """ run_keys - input fields that determine if this combination has been __tried already. E.g. host, user, pass """
        """ result_keys - input fields that determine if the goal has been reached already. E.g. host, user (or just host)."""


        self.name = name
        self.utils = KnowledgeUtils()

        self.run_keys = run_keys
        self.result_keys = result_keys

    def run(self, retry=False, retry_successful=False, remove_previous=False):
        """ Call it to run the StateAction. """
        """ retry - if True, previous trials will be ignord. Does not run if goal is __achieved though. """
        """ retry_successful - if True, retries even when goal achieved.  """
        """ remove_previous - if True, removes previous successful result records for tried combinations """
        self.__scan(self.emit(), retry=retry, retry_successful=retry_successful, remove_previous=remove_previous)

    def emit(self):
        """ Overload this. Return a list of all possible inputs here. The framework will take care of not running on inputs that already run."""
        return []

    def work(self, args):
        """ Overload this. Work for single input happens here. Return dictionary with at least {'status': code} where code=0 for failure, 1 for sucess """
        # print("run: " + str(args))
        return {'status': 0}

    def strip_list(self, dic_list, keys):
        return [self.strip(d, keys) for d in dic_list]

    def strip(self, dic, keys):
        """ Takes a dictionary on input. Returns a flat dictionary comprising of dic stripped to the keys. Keys can follow dot notation, e.g. key.subkey.subsubkey """
        """ Returned dict will have key names of the deepest subkeys """
        ret = {}
        for k in keys:
            path=k.split('.')
            d=dic
            subkey=''
            while len(path)>0:
                subkey=path.pop(0)
                # print(subkey)
                arrsplit = subkey.split('[')
                subkey_main=arrsplit[0]
                d=d[subkey_main]
                if(len(arrsplit)>1):
                    index = int(arrsplit[1].split(']')[0])
                    d=d[index]
            ret[subkey_main]=d
        return ret

    def __scan(self, input_list, retry=False, retry_successful=False, remove_previous=False):
        """ The core engine of StateActions. Ensures work is done only once (according to the run_keys). Does not run upon given input element if the goal has been achievd alreay (as per result_keys). """
        """ input_list - list of all possible trials that could be made. The framework will ignore those that were run alrady and those that would lead to achieving something that is __achieved already """

        self.__run_count = 0
        total_count = len(input_list)
        inputs_count = 0

        goals = set()
        achievements = set()

        for m in input_list:
            if(remove_previous):
                self.__remove_result(self.strip(m, self.result_keys))

        for m in input_list:
            inputs_count += 1

            # counting goals encountered
            goals.add(str(self.strip(m, self.result_keys)))

            should_try = retry or not self.__tried(self.strip(m, self.run_keys))
            goal_achieved = not retry_successful and self.__achieved(self.strip(m, self.result_keys))

            # if(not goal_achieved):
            #     print("Should_try: {}  Goal_achieved: {}  user: {}  host: {}".format(should_try, goal_achieved, m['username'], m['host']))

            if (should_try and not goal_achieved):
                self.__run_count += 1
                print('\n--- Run {} ({} / {} = {}%)---'.format(self.__run_count, inputs_count, total_count, int(100*inputs_count/total_count) ))

                result = self.work(m)

                success = (result['status'] == 1)
                self.__record_trial(self.strip(m, self.run_keys), result, m)
                if (success):
                    self.__record_result(self.strip(m, self.result_keys), result, m)

            # it's just to count goals achieved
            if(self.__achieved(self.strip(m, self.result_keys))):
                achievements.add(str(self.strip(m, self.result_keys)))

        print("{} runs permformed (out of {} inputs emitted)".format(self.__run_count, len(input_list)))
        print("Goals achieved: {} / {} = {}%".format(len(achievements), len(goals), int(100*len(achievements)/len(goals)) if len(goals)>0 else 100 ))

    def __find(self, p, subtype):
        params = {'key.'+k: v for k,v in p.items()}
        params['name'] = self.name
        params['type'] = 'state_action'
        params['subtype'] = subtype
        f = self.utils.knowledge.find_one(params)
        # print("    "+ subtype + "? " + str(params) + " ---> " + str(f))
        return bool(f)

    def __tried(self, params):
        return self.__find(params, 'trial')

    def __achieved(self, params):
        return self.__find(params, 'result')

    def __find_params(self, p, subtype):
        # Mongo search keys: prependng with "key."
        params = {'key.'+k: v for k,v in p.items()}
        params['name'] = self.name
        params['type'] = 'state_action'
        params['subtype'] = subtype
        params['pentest_id'] = self.utils.pentest_id
        return params

    def __save(self, p, result, original_input, subtype):
        params = self.__find_params(p, subtype)
        doc = {}
        doc['name'] = self.name
        doc['type'] = 'state_action'
        doc['subtype'] = subtype
        doc['pentest_id'] = self.utils.pentest_id

        doc['result'] = result
        doc['key'] = p.copy()
        doc['input'] = original_input

        doc['updated_at'] = time.time()
        f = self.utils.knowledge.replace_one(params, doc, upsert=True)
        # print("record "+subtype+": " + str(doc))

    def __record_trial(self, params, result, original_input):
        self.__save(params, result, original_input, 'trial')

    def __record_result(self, params, result, original_input):
        print("--- SUCCESS --- \n{}".format(result))
        self.__save(params, result, original_input, 'result')

    def __remove_result(self, params):
        params = self.__find_params(params, 'result')
        self.utils.knowledge.remove(params)


    #
    # @classmethod
    # def test(cls):
    #     s = StateAction('foo',['a', 'b'], ['a'])
    #     API CHANGED. s.scan([{'a':11 ,'b': 20}, {'a':11 ,'b': 22}, {'a':2 ,'b': 30}], )
