Merge pull request #507 from DrDroidLab/global_variable_condition

adds global variable conditions

Author: jayeshsadhwani99

executor/playbook_result_conditional_evaluators/global_variable_evaluators/__init__.py

@@ -0,0 +1,57 @@
+from executor.playbook_result_conditional_evaluators.global_variable_evaluators.global_variable_evaluator import \
+    GlobalVariableEvaluator
+from protos.base_pb2 import Operator
+from protos.playbooks.playbook_pb2 import PlaybookTaskExecutionLog
+from protos.playbooks.playbook_global_variable_evaluator_pb2 import GlobalVariableResultRule, CompareGlobalVariable
+from utils.proto_utils import proto_to_dict
+from utils.dict_utils import get_nested_value
+
+class CompareGlobalVariableEvaluator(GlobalVariableEvaluator):
+
+    def evaluate(self, rule: GlobalVariableResultRule, playbook_task_execution_log: [PlaybookTaskExecutionLog]) -> bool:
+        if rule.type != GlobalVariableResultRule.Type.COMPARE_GLOBAL_VARIABLE:
+            raise ValueError(f'Rule type {GlobalVariableResultRule.Type.Name(rule.type)} not supported')
+
+        compare_global_variable: CompareGlobalVariable = rule.compare_global_variable
+        operator = compare_global_variable.operator
+        variable_name = compare_global_variable.variable_name.value
+        threshold = compare_global_variable.threshold.value
+
+        global_variable_set = next(
+                    (tr.execution_global_variable_set for tr in playbook_task_execution_log), None)
+        global_variable_set_dict = proto_to_dict(global_variable_set) if global_variable_set else {}
+        value = get_nested_value(global_variable_set_dict, variable_name)
+
+        # compare current time with first member of cron schedules
+        if operator == Operator.EQUAL_O:
+            return value == threshold
+        elif operator == Operator.GREATER_THAN_O:
+            try:
+                value = int(value)
+                threshold = int(threshold)
+            except ValueError:
+                return False
+            return value > threshold
+        elif operator == Operator.GREATER_THAN_EQUAL_O:
+            try:
+                value = int(value)
+                threshold = int(threshold)
+            except ValueError:
+                return False
+            return value >= threshold
+        elif operator == Operator.LESS_THAN_O:
+            try:
+                value = int(value)
+                threshold = int(threshold)
+            except ValueError:
+                return False
+            return value < threshold
+        elif operator == Operator.LESS_THAN_EQUAL_O:
+            try:
+                value = int(value)
+                threshold = int(threshold)
+            except ValueError:
+                return False
+            return value <= threshold
+        else:
+            raise ValueError(f'Operator {Operator.Name(operator)} not supported')

executor/playbook_result_conditional_evaluators/global_variable_evaluators/compare_global_variable_evaluator.py

@@ -0,0 +1,8 @@
+from protos.playbooks.playbook_pb2 import PlaybookTaskExecutionLog
+from protos.playbooks.playbook_global_variable_evaluator_pb2 import GlobalVariableResultRule
+
+
+class GlobalVariableEvaluator:
+
+    def evaluate(self, rule: GlobalVariableResultRule, playbook_task_execution_log: [PlaybookTaskExecutionLog]) -> bool:
+        pass

executor/playbook_result_conditional_evaluators/global_variable_evaluators/global_variable_evaluator.py

@@ -12,18 +12,22 @@
     TimeseriesResultEvaluator
 from executor.playbook_result_conditional_evaluators.task_result_evalutors.bash_command_result_evaluator import \
     BashCommandOutputResultEvaluator
+from executor.playbook_result_conditional_evaluators.global_variable_evaluators.compare_global_variable_evaluator import \
+    CompareGlobalVariableEvaluator
 from protos.base_pb2 import LogicalOperator
 from protos.playbooks.playbook_commons_pb2 import PlaybookTaskResultType
 from protos.playbooks.playbook_pb2 import PlaybookStepResultCondition, PlaybookTaskResultRule, \
     PlaybookTaskExecutionLog
 from protos.playbooks.playbook_step_result_evaluator_pb2 import PlaybookStepResultRule
+from protos.playbooks.playbook_global_variable_evaluator_pb2 import GlobalVariableResultRule
 
 
 class StepConditionEvaluator:
 
     def __init__(self):
         self._task_rule_map = {}
         self._step_rule_map = {}
+        self._variable_rule_map = {}
 
     def register_task_result_evaluator(self, result_type: PlaybookTaskResultType,
                                        task_result_evaluator: TaskResultEvaluator):
@@ -33,6 +37,10 @@ def register_step_result_evaluator(self, step_rule_type: PlaybookStepResultRule.
                                        step_result_evaluator: StepResultEvaluator):
         self._step_rule_map[step_rule_type] = step_result_evaluator
 
+    def register_variable_evaluator(self, variable_rule_type: GlobalVariableResultRule.Type,
+                                       variable_evaluator: TaskResultEvaluator):
+        self._variable_rule_map[variable_rule_type] = variable_evaluator
+
     def evaluate(self, condition: PlaybookStepResultCondition,
                  playbook_task_execution_log: [PlaybookTaskExecutionLog]) -> (bool, Dict):
         if not condition.rule_sets:
@@ -63,6 +71,15 @@ def evaluate(self, condition: PlaybookStepResultCondition,
                     raise ValueError(f"Step result type {PlaybookStepResultRule.Type.Name(sr.type)} not supported")
                 evaluation = step_result_evaluator.evaluate(sr, playbook_task_execution_log)
                 all_evaluations.append(evaluation)
+
+            variable_rules: [PlaybookTaskResultRule] = rs.variable_rules
+            for vr in variable_rules:
+                variable_evaluator = self._variable_rule_map.get(vr.type)
+                if not variable_evaluator:
+                    raise ValueError(f"Task result type {task_result.type} not supported")
+                evaluation = variable_evaluator.evaluate(vr, playbook_task_execution_log)
+                all_evaluations.append(evaluation)
+            
             logical_operator = rs.logical_operator
             if logical_operator == LogicalOperator.AND_LO:
                 all_rs_evaluations.append(all(all_evaluations))
@@ -99,3 +116,5 @@ def evaluate(self, condition: PlaybookStepResultCondition,
 
 step_condition_evaluator.register_step_result_evaluator(PlaybookStepResultRule.Type.COMPARE_TIME_WITH_CRON,
                                                         CompareTimeWithCronEvaluator())
+step_condition_evaluator.register_variable_evaluator(GlobalVariableResultRule.Type.COMPARE_GLOBAL_VARIABLE,
+                                                        CompareGlobalVariableEvaluator())

executor/playbook_result_conditional_evaluators/step_condition_evaluator.py

@@ -77,6 +77,10 @@ def execute_playbook_step_impl(tr: TimeRange, account: Account, step: PlaybookSt
         execution_global_variable_set = Struct()
         if global_variable_set and global_variable_set.items():
             execution_global_variable_set.CopyFrom(global_variable_set)
+
+        task_execution_global_variable_set = Struct()
+        if global_variable_set and global_variable_set.items():
+            task_execution_global_variable_set.CopyFrom(global_variable_set)
         for task_proto in tasks:
             if not execution_global_variable_set or not execution_global_variable_set.items():
                 if task_proto.global_variable_set and task_proto.global_variable_set.items():
@@ -123,9 +127,9 @@ def execute_playbook_step_impl(tr: TimeRange, account: Account, step: PlaybookSt
                     continue
             for bev in bulk_execution_var_values:
                 if is_bulk_execution and bulk_task_var:
-                    execution_global_variable_set[bulk_task_var] = bev
+                    task_execution_global_variable_set[bulk_task_var] = bev
                 try:
-                    task_result = playbook_source_facade.execute_task(account.id, tr, execution_global_variable_set,
+                    task_result = playbook_source_facade.execute_task(account.id, tr, task_execution_global_variable_set,
                                                                       task_proto)
                     task_interpretation: InterpretationProto = task_result_interpret(interpreter_type, task_proto,
                                                                                      task_result)

executor/tasks.py

@@ -9,6 +9,7 @@ import "protos/playbooks/intelligence_layer/interpreter.proto";
 
 import "protos/playbooks/playbook_task_result_evaluator.proto";
 import "protos/playbooks/playbook_step_result_evaluator.proto";
+import "protos/playbooks/playbook_global_variable_evaluator.proto";
 
 import "protos/playbooks/source_task_definitions/cloudwatch_task.proto";
 import "protos/playbooks/source_task_definitions/grafana_task.proto";
@@ -118,6 +119,7 @@ message PlaybookStepResultCondition {
     LogicalOperator logical_operator = 1;
     repeated PlaybookTaskResultRule rules = 2;
     repeated PlaybookStepResultRule step_rules = 3;
+    repeated GlobalVariableResultRule variable_rules = 4;
   }
   LogicalOperator logical_operator = 1;
   repeated RuleSet rule_sets = 2;

protos/playbooks/playbook.proto

@@ -22,6 +22,7 @@ enum PlaybookTaskResultType {
   BASH_COMMAND_OUTPUT = 4;
   TEXT = 5;
   LOGS = 6;
+  GLOBAL_VARIABLE = 7;
 }
 
 message ExternalLink {