Merge branch 'main' of github.com:DrDroidLab/PlayBooks into sandbox

connectors/views.py

@@ -1,5 +1,7 @@
 import logging
 from typing import Union
+
+from google.protobuf.struct_pb2 import Struct
 from google.protobuf.wrappers_pb2 import UInt64Value, StringValue, BoolValue
 
 from django.http import HttpResponse, JsonResponse
@@ -22,7 +24,7 @@
 from protos.base_pb2 import Message, Meta, Page, TimeRange, SourceKeyType
 from protos.connectors.api_pb2 import CreateConnectorRequest, CreateConnectorResponse, GetConnectorsListRequest, \
     GetConnectorsListResponse, GetSlackAlertTriggerOptionsRequest, GetSlackAlertTriggerOptionsResponse, \
-    GetSlackAlertsRequest, GetSlackAlertsResponse, GetSlackAppManifestRequest, GetSlackAppManifestResponse, \
+        GetSlackAlertsRequest, GetSlackAlertsResponse, GetSlackAppManifestRequest, GetSlackAppManifestResponse, \
     UpdateConnectorRequest, UpdateConnectorResponse, GetConnectorKeysOptionsRequest, \
     GetConnectorKeysOptionsResponse, GetConnectorKeysRequest, GetConnectorKeysResponse, \
     GetConnectorPlaybookSourceOptionsRequest, GetConnectorPlaybookSourceOptionsResponse, GetConnectedPlaybooksRequest, \
@@ -33,6 +35,7 @@
     SlackAlert as SlackAlertProto
 from protos.base_pb2 import Source, SourceModelType
 from protos.connectors.connector_pb2 import Connector as ConnectorProto
+from utils.proto_utils import dict_to_proto
 
 logger = logging.getLogger(__name__)
 
@@ -290,11 +293,12 @@ def slack_alerts_search(request_message: GetSlackAlertsRequest) -> \
     total_count = qs.count()
     qs = filter_page(qs, page)
     qs = qs.values('id', 'alert_type', 'title', 'text', 'data_timestamp', 'slack_channel_metadata_model__id',
-                   'channel_id',
+                   'channel_id', 'data',
                    'slack_channel_metadata_model__metadata__channel_name')
 
     slack_alerts = []
     for a in qs:
+        data_struct = dict_to_proto(a['data'], Struct)
         slack_alerts.append(SlackAlertProto(id=UInt64Value(value=a['id']),
                                             alert_type=StringValue(value=a['alert_type']),
                                             alert_title=StringValue(value=a['title']),
@@ -304,7 +308,8 @@ def slack_alerts_search(request_message: GetSlackAlertsRequest) -> \
                                                 channel_id=StringValue(value=a['channel_id']),
                                                 channel_name=StringValue(
                                                     value=a['slack_channel_metadata_model__metadata__channel_name'])),
-                                            alert_timestamp=int(a['data_timestamp'].timestamp())))
+                                            alert_timestamp=int(a['data_timestamp'].timestamp()),
+                                            alert_json=data_struct))
     return GetSlackAlertsResponse(meta=get_meta(page=page, total_count=total_count),
                                   slack_alerts=slack_alerts)
 

executor/crud/playbook_execution_crud.py

@@ -49,10 +49,10 @@ def create_playbook_execution(account: Account, time_range: TimeRange, playbook_
         raise e
 
 
-def update_db_account_playbook_execution_status(account: Account, playbook_run_id: int,
+def update_db_account_playbook_execution_status(account: Account, playbook_execution_id: int,
                                                 status: PlaybookExecutionStatusType):
     try:
-        playbook_execution = account.playbookexecution_set.get(id=playbook_run_id)
+        playbook_execution = account.playbookexecution_set.get(id=playbook_execution_id)
         playbook_execution.status = status
         update_fields = ['status']
         if status == PlaybookExecutionStatusType.RUNNING:
@@ -65,10 +65,10 @@ def update_db_account_playbook_execution_status(account: Account, playbook_run_i
         return True
     except PlayBookExecution.DoesNotExist:
         logger.error(
-            f"Failed to get playbook execution for account_id: {account.id}, playbook_run_id: {playbook_run_id}")
+            f"Failed to get playbook execution for account_id: {account.id}, playbook_run_id: {playbook_execution_id}")
     except Exception as e:
         logger.error(
-            f"Failed to get playbook execution for account_id: {account.id}, playbook_run_id: {playbook_run_id}, error: {e}")
+            f"Failed to get playbook execution for account_id: {account.id}, playbook_run_id: {playbook_execution_id}, error: {e}")
     return False
 
 
@@ -115,7 +115,7 @@ def bulk_create_playbook_execution_log(account, playbook, playbook_execution, al
             raise e
 
         for result in all_task_executions:
-            playbook_execution_log = PlayBookTaskExecutionLog(
+            playbook_task_execution_log = PlayBookTaskExecutionLog(
                 account=account,
                 playbook=playbook,
                 playbook_execution=playbook_execution,
@@ -125,9 +125,10 @@ def bulk_create_playbook_execution_log(account, playbook, playbook_execution, al
                 playbook_task_result=result['task_result'],
                 interpretation=result['task_interpretation'],
                 created_by=user,
-                time_range=time_range
+                time_range=time_range,
+                execution_global_variable_set=result['execution_global_variable_set']
             )
-            all_db_playbook_execution_logs.append(playbook_execution_log)
+            all_db_playbook_execution_logs.append(playbook_task_execution_log)
         for step_relation_execution_log in all_step_relation_executions:
             playbook_step_relation_execution_log = PlayBookStepRelationExecutionLog(
                 account=account,

executor/migrations/0042_playbooktaskexecutionlog_execution_global_variable_set.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.1.13 on 2024-07-20 12:02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('executor', '0041_migrate_logs_task_executions_to_new_def'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='playbooktaskexecutionlog',
+            name='execution_global_variable_set',
+            field=models.JSONField(blank=True, null=True),
+        ),
+    ]

executor/models.py

@@ -394,6 +394,10 @@ class Meta:
     def proto_partial(self) -> PlaybookExecutionProto:
         time_range_proto = dict_to_proto(self.time_range, TimeRange) if self.time_range else TimeRange()
 
+        execution_global_variable_set_proto = Struct()
+        if self.execution_global_variable_set:
+            execution_global_variable_set_proto.update(self.execution_global_variable_set)
+
         return PlaybookExecutionProto(
             id=UInt64Value(value=self.id),
             playbook_run_id=StringValue(value=self.playbook_run_id),
@@ -404,10 +408,15 @@ def proto_partial(self) -> PlaybookExecutionProto:
             created_at=int(self.created_at.replace(tzinfo=timezone.utc).timestamp()),
             created_by=StringValue(value=self.created_by) if self.created_by else None,
             time_range=time_range_proto,
+            execution_global_variable_set=execution_global_variable_set_proto
         )
 
     @property
     def proto(self) -> PlaybookExecutionProto:
+        execution_global_variable_set_proto = Struct()
+        if self.execution_global_variable_set:
+            execution_global_variable_set_proto.update(self.execution_global_variable_set)
+
         playbook_step_execution_logs = self.playbookstepexecutionlog_set.all()
         if not playbook_step_execution_logs:
             playbook_execution_logs = self.playbooktaskexecutionlog_set.all()
@@ -442,7 +451,8 @@ def proto(self) -> PlaybookExecutionProto:
             created_at=int(self.created_at.replace(tzinfo=timezone.utc).timestamp()),
             created_by=StringValue(value=self.created_by) if self.created_by else None,
             time_range=time_range_proto,
-            step_execution_logs=step_execution_logs
+            step_execution_logs=step_execution_logs,
+            execution_global_variable_set=execution_global_variable_set_proto
         )
 
     @property
@@ -547,8 +557,14 @@ class PlayBookTaskExecutionLog(models.Model):
     created_by = models.TextField(null=True, blank=True)
     time_range = models.JSONField(null=True, blank=True)
 
+    execution_global_variable_set = models.JSONField(null=True, blank=True)
+
     @property
     def proto(self) -> PlaybookTaskExecutionLogProto:
+        execution_global_variable_set_proto = Struct()
+        if self.execution_global_variable_set:
+            execution_global_variable_set_proto.update(self.execution_global_variable_set)
+
         task = self.playbook_task_definition.proto_with_connector_source(self.playbook.id, self.playbook_step.id)
         time_range_proto = dict_to_proto(self.time_range, TimeRange) if self.time_range else TimeRange()
         return PlaybookTaskExecutionLogProto(
@@ -559,7 +575,8 @@ def proto(self) -> PlaybookTaskExecutionLogProto:
             interpretation=dict_to_proto(self.interpretation,
                                          InterpretationProto) if self.interpretation else InterpretationProto(),
             created_by=StringValue(value=self.created_by) if self.created_by else None,
-            time_range=time_range_proto
+            time_range=time_range_proto,
+            execution_global_variable_set=execution_global_variable_set_proto
         )
 
     @property
@@ -640,6 +657,6 @@ def proto(self) -> PlaybookStepRelationExecutionLogProto:
             evaluation_result=BoolValue(value=self.evaluation_result),
             evaluation_output=evaluation_output_proto,
             step_relation_interpretation=dict_to_proto(self.interpretation,
-                                        InterpretationProto) if self.interpretation else InterpretationProto(),
+                                                       InterpretationProto) if self.interpretation else InterpretationProto(),
 
         )

executor/source_processors/kubectl_api_processor.py

@@ -59,22 +59,32 @@ def execute_command(self, command):
         command = command.strip()
         if 'kubectl' in command:
             command = command.replace('kubectl', '')
+        if '|' in command:
+            commands = [cmd.strip() for cmd in command.split('|')]
+        else:
+            commands = [command]
         if self.__ca_cert:
             kubectl_command = [
                                   "kubectl",
                                   f"--server={self.__api_server}",
                                   f"--token={self.__token}",
                                   f"--certificate-authority={self.__ca_cert}"
-                              ] + command.split()
+                              ] + commands[0].split()
         else:
             kubectl_command = [
                                   "kubectl",
                                   f"--server={self.__api_server}",
-                                  f"--token={self.__token}"
-                              ] + command.split()
+                                  f"--token={self.__token}",
+                                  f"--insecure-skip-tls-verify=true"
+                              ] + commands[0].split()
         try:
             process = subprocess.Popen(kubectl_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
             stdout, stderr = process.communicate()
+            if len(commands) > 1:
+                for cmd in commands[1:]:
+                    process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                                               stderr=subprocess.PIPE, text=True, shell=True)
+                    stdout, stderr = process.communicate(input=stdout)
             if process.returncode == 0:
                 print("Command Output:", stdout)
                 return stdout

executor/source_task_executors/bash_task_executor.py

@@ -69,7 +69,7 @@ def execute_command(self, time_range: TimeRange, global_variable_set: Dict, bash
                 for key, value in global_variable_set.items():
                     updated_commands = []
                     for command in commands:
-                        command = command.replace(f"{{{key}}}", value)
+                        command = command.replace(key, str(value))
                         updated_commands.append(command)
                     commands = updated_commands
             try:

executor/source_task_executors/eks_task_executor.py

@@ -329,7 +329,7 @@ def execute_kubectl_command(self, time_range: TimeRange, global_variable_set: Di
                 for key, value in global_variable_set.items():
                     updated_commands = []
                     for command in commands:
-                        command = command.replace(f"{{{key}}}", value)
+                        command = command.replace(key, str(value))
                         updated_commands.append(command)
                     commands = updated_commands
             try:

executor/source_task_executors/gke_task_executor.py

@@ -313,7 +313,7 @@ def execute_kubectl_command(self, time_range: TimeRange, global_variable_set: Di
                 for key, value in global_variable_set.items():
                     updated_commands = []
                     for command in commands:
-                        command = command.replace(f"{{{key}}}", value)
+                        command = command.replace(key, str(value))
                         updated_commands.append(command)
                     commands = updated_commands
             try:

executor/source_task_executors/kubernetes_task_executor.py

@@ -39,7 +39,7 @@ def execute_command(self, time_range: TimeRange, global_variable_set: Dict,
                 for key, value in global_variable_set.items():
                     updated_commands = []
                     for command in commands:
-                        command = command.replace(f"{{{key}}}", value)
+                        command = command.replace(key, str(value))
                         updated_commands.append(command)
                     commands = updated_commands
             try: