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

executor/crud/playbooks_crud.py

@@ -217,13 +217,13 @@ def get_or_create_db_task(account: Account, created_by, task: PlaybookTaskProto)
     task_dict.pop('reference_id', None)
     task_md5 = md5(str(task_dict).encode('utf-8')).hexdigest()
     try:
-        db_task, _ = PlayBookTask.objects.get_or_create(account=account,
-                                                        name=task.name.value,
-                                                        task_md5=task_md5,
-                                                        created_by=created_by,
-                                                        defaults={'task': task_dict,
-                                                                  'description': task.description.value,
-                                                                  'notes': task.notes.value})
+        db_task, _ = PlayBookTask.objects.update_or_create(account=account,
+                                                           name=task.name.value,
+                                                           task_md5=task_md5,
+                                                           created_by=created_by,
+                                                           defaults={'task': task_dict,
+                                                                     'description': task.description.value,
+                                                                     'notes': task.notes.value})
         return db_task, None
     except IntegrityError:
         db_task = PlayBookTask.objects.get(account=account, name=task.name.value, task_md5=task_md5,

executor/playbook_source_facade.py

@@ -69,6 +69,7 @@ def get_source_options(self, account_id) -> [PlaybookSourceOptions]:
                 model_types = task_info['model_types']
                 model_options: [PlaybookSourceOptions.TaskTypeOption.SourceModelTypeMap] = []
                 task_result_type = task_info.get('result_type', PlaybookTaskResultType.UNKNOWN)
+                form_fields = task_info.get('form_fields', [])
                 for m in model_types:
                     model_type_display_name = model_type_display_name_maps.get(m, SourceModelType.Name(m))
                     model_options.append(
@@ -79,7 +80,8 @@ def get_source_options(self, account_id) -> [PlaybookSourceOptions]:
                                                          task_type=StringValue(value=task_type_name),
                                                          category=StringValue(value=task_type_category),
                                                          supported_model_types=model_options,
-                                                         result_type=task_result_type))
+                                                         result_type=task_result_type,
+                                                         form_fields=form_fields))
             display_name = integrations_connector_type_display_name_map.get(source, Source.Name(source))
             source_options.append(PlaybookSourceOptions(source=source,
                                                         display_name=StringValue(value=display_name),

executor/playbook_source_manager.py

@@ -6,11 +6,35 @@
 from executor.source_processors.processor import Processor
 from protos.base_pb2 import TimeRange, Source
 from protos.connectors.connector_pb2 import Connector as ConnectorProto
+from protos.literal_pb2 import LiteralType
 from protos.playbooks.playbook_commons_pb2 import PlaybookTaskResult
 from protos.playbooks.playbook_pb2 import PlaybookTask
+from protos.ui_definition_pb2 import FormField
 from utils.proto_utils import proto_to_dict, dict_to_proto
 
 
+def resolve_global_variables(global_variable_set: Dict, form_fields: [FormField], source_type_task_def: Dict) -> Dict:
+    all_string_fields = [ff.key_name.value for ff in form_fields if ff.data_type == LiteralType.STRING]
+    all_string_array_fields = [ff.key_name.value for ff in form_fields if ff.data_type == LiteralType.STRING_ARRAY]
+    all_composite_fields = {}
+    for ff in form_fields:
+        if ff.is_composite:
+            all_composite_fields[ff.key_name.value] = ff.composite_fields
+    for gk, gv in global_variable_set.items():
+        for tk, tv in source_type_task_def.items():
+            if tk in all_string_fields:
+                source_type_task_def[tk] = tv.replace(gk, gv)
+            elif tk in all_string_array_fields:
+                source_type_task_def[tk] = [item.replace(gk, gv) for item in tv]
+            elif tk in all_composite_fields:
+                composite_fields = all_composite_fields[tk]
+                for item in source_type_task_def[tk]:
+                    for cf in composite_fields:
+                        if cf.data_type == LiteralType.STRING:
+                            item[cf.key_name.value] = item[cf.key_name.value].replace(gk, gv)
+    return source_type_task_def
+
+
 class PlaybookSourceManager:
     source: Source = Source.UNKNOWN
     task_proto = None
@@ -73,8 +97,16 @@ def execute_task(self, account_id, time_range: TimeRange, global_variable_set: D
             task_type = source_task_proto.type
             if task_type in self.task_type_callable_map:
                 try:
-                    return self.task_type_callable_map[task_type]['executor'](time_range, global_variable_set,
-                                                                              source_task_proto, source_connector_proto)
+                    task_type_name = self.task_proto.TaskType.Name(task_type).lower()
+                    source_type_task_def = source_task.get(task_type_name, {})
+                    form_fields = self.task_type_callable_map[task_type]['form_fields']
+                    resolved_source_type_task_def = resolve_global_variables(global_variable_set, form_fields,
+                                                                             source_type_task_def)
+                    source_task[task_type_name] = resolved_source_type_task_def
+                    resolved_task_def_proto = dict_to_proto(source_task, self.task_proto)
+                    return self.task_type_callable_map[task_type]['executor'](time_range,
+                                                                              resolved_task_def_proto,
+                                                                              source_connector_proto)
                 except Exception as e:
                     raise Exception(f"Error while executing task for source: {source_str} with error: {e}")
             else:

executor/source_task_executors/api_task_executor.py

@@ -1,15 +1,16 @@
 import json
-from typing import Dict
 
 import requests
 from google.protobuf.struct_pb2 import Struct
 
-from google.protobuf.wrappers_pb2 import StringValue, UInt64Value
+from google.protobuf.wrappers_pb2 import StringValue, UInt64Value, Int64Value
 from executor.playbook_source_manager import PlaybookSourceManager
 from protos.base_pb2 import Source, TimeRange
 from protos.connectors.connector_pb2 import Connector as ConnectorProto
+from protos.literal_pb2 import LiteralType, Literal
 from protos.playbooks.playbook_commons_pb2 import PlaybookTaskResult, ApiResponseResult, PlaybookTaskResultType
 from protos.playbooks.source_task_definitions.api_task_pb2 import Api
+from protos.ui_definition_pb2 import FormField
 
 method_proto_string_mapping = {
     Api.HttpRequest.Method.GET: "GET",
@@ -31,19 +32,46 @@ def __init__(self):
                 'model_types': [],
                 'result_type': PlaybookTaskResultType.API_RESPONSE,
                 'display_name': 'Trigger an API',
-                'category': 'Actions'
+                'category': 'Actions',
+                'form_fields': [
+                    FormField(key_name=StringValue(value="method"),
+                              display_name=StringValue(value="Method"),
+                              description=StringValue(value='Select Method'),
+                              data_type=LiteralType.STRING,
+                              valid_values=[
+                                  Literal(type=LiteralType.STRING, string=StringValue(value="GET")),
+                                  Literal(type=LiteralType.STRING, string=StringValue(value="POST")),
+                                  Literal(type=LiteralType.STRING, string=StringValue(value="PUT")),
+                                  Literal(type=LiteralType.STRING, string=StringValue(value="PATCH")),
+                                  Literal(type=LiteralType.STRING, string=StringValue(value="DELETE"))
+                              ]),
+                    FormField(key_name=StringValue(value="url"),
+                              display_name=StringValue(value="URL"),
+                              description=StringValue(value='Enter URL'),
+                              data_type=LiteralType.STRING),
+                    FormField(key_name=StringValue(value="headers"),
+                              display_name=StringValue(value="Headers (Enter JSON)"),
+                              data_type=LiteralType.STRING,
+                              is_optional=True),
+                    FormField(key_name=StringValue(value="payload"),
+                              display_name=StringValue(value="Payload/Body (Enter JSON)"),
+                              data_type=LiteralType.STRING,
+                              is_optional=True),
+                    FormField(key_name=StringValue(value="timeout"),
+                              display_name=StringValue(value="Timeout (in seconds)"),
+                              description=StringValue(value='Enter Timeout (in seconds)'),
+                              data_type=LiteralType.LONG,
+                              default_value=Literal(type=LiteralType.LONG, long=Int64Value(value=120))),
+                ]
             },
         }
 
-    def execute_http_request(self, time_range: TimeRange, global_variable_set: Dict,
-                             api_task: Api, api_connector_proto: ConnectorProto) -> PlaybookTaskResult:
+    def execute_http_request(self, time_range: TimeRange, api_task: Api,
+                             api_connector_proto: ConnectorProto) -> PlaybookTaskResult:
         try:
             http_request = api_task.http_request
             method = http_request.method
             url = http_request.url.value
-            if global_variable_set:
-                for key, value in global_variable_set.items():
-                    url = url.replace(f"{{{key}}}", value)
             headers = http_request.headers.value
             headers = json.loads(headers) if headers else {}
             if 'Content-Type' not in headers:

executor/source_task_executors/azure_task_executor.py

@@ -1,6 +1,5 @@
 import logging
 from datetime import timedelta
-from typing import Dict
 
 from google.protobuf.wrappers_pb2 import StringValue, UInt64Value
 
@@ -10,8 +9,10 @@
 
 from protos.base_pb2 import TimeRange, Source, SourceModelType
 from protos.connectors.connector_pb2 import Connector as ConnectorProto
+from protos.literal_pb2 import LiteralType, Literal
 from protos.playbooks.playbook_commons_pb2 import PlaybookTaskResult, PlaybookTaskResultType, TableResult
 from protos.playbooks.source_task_definitions.azure_task_pb2 import Azure
+from protos.ui_definition_pb2 import FormField
 
 logger = logging.getLogger(__name__)
 
@@ -27,16 +28,30 @@ def __init__(self):
                 'model_types': [SourceModelType.AZURE_WORKSPACE],
                 'result_type': PlaybookTaskResultType.LOGS,
                 'display_name': 'Fetch logs from Azure Log Analytics Workspace',
-                'category': 'Logs'
+                'category': 'Logs',
+                'form_fields': [
+                    FormField(key_name=StringValue(value="workspace_id"),
+                              display_name=StringValue(value="Azure Workspace ID"),
+                              description=StringValue(value='Select Workspace ID'),
+                              data_type=LiteralType.STRING),
+                    FormField(key_name=StringValue(value="filter_query"),
+                              display_name=StringValue(value="Log Filter Query"),
+                              data_type=LiteralType.STRING),
+                    FormField(key_name=StringValue(value="timespan"),
+                              display_name=StringValue(value="Timespan (hours)"),
+                              description=StringValue(value='Enter Timespan (hours)'),
+                              data_type=LiteralType.STRING,
+                              default_value=Literal(type=LiteralType.STRING, string=StringValue(value="1")))
+                ]
             },
         }
 
     def get_connector_processor(self, azure_connector, **kwargs):
         generated_credentials = generate_credentials_dict(azure_connector.type, azure_connector.keys)
         return AzureApiProcessor(**generated_credentials)
 
-    def filter_log_events(self, time_range: TimeRange, global_variable_set: Dict,
-                          azure_task: Azure, azure_connector: ConnectorProto) -> PlaybookTaskResult:
+    def filter_log_events(self, time_range: TimeRange, azure_task: Azure,
+                          azure_connector: ConnectorProto) -> PlaybookTaskResult:
         try:
             tr_end_time = time_range.time_lt
             end_time = int(tr_end_time * 1000)
@@ -49,9 +64,6 @@ def filter_log_events(self, time_range: TimeRange, global_variable_set: Dict,
             timespan = timedelta(hours=int(timespan_delta)) if timespan_delta else timedelta(
                 seconds=end_time - start_time)
             query_pattern = task.filter_query.value
-            if global_variable_set:
-                for key, value in global_variable_set.items():
-                    query_pattern = query_pattern.replace(key, str(value))
 
             azure_api_processor = self.get_connector_processor(azure_connector)
 
@@ -75,7 +87,8 @@ def filter_log_events(self, time_range: TimeRange, global_variable_set: Dict,
                     table_rows.append(table_row)
 
             result = TableResult(
-                raw_query=StringValue(value=f'Execute {query_pattern} on Azure Log Analytics workspace: {workspace_id}'),
+                raw_query=StringValue(
+                    value=f'Execute {query_pattern} on Azure Log Analytics workspace: {workspace_id}'),
                 rows=table_rows,
                 total_count=UInt64Value(value=len(table_rows)),
             )

executor/source_task_executors/bash_task_executor.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 from google.protobuf.wrappers_pb2 import StringValue
 
 from connectors.crud.connector_asset_model_crud import get_db_connector_metadata_models
@@ -9,8 +7,10 @@
 from executor.source_processors.remote_server_processor import RemoteServerProcessor
 from protos.base_pb2 import TimeRange, Source, SourceModelType
 from protos.connectors.connector_pb2 import Connector as ConnectorProto
+from protos.literal_pb2 import LiteralType
 from protos.playbooks.playbook_commons_pb2 import PlaybookTaskResult, BashCommandOutputResult, PlaybookTaskResultType
 from protos.playbooks.source_task_definitions.bash_task_pb2 import Bash
+from protos.ui_definition_pb2 import FormField
 
 
 class BashSourceManager(PlaybookSourceManager):
@@ -24,7 +24,16 @@ def __init__(self):
                 'model_types': [SourceModelType.SSH_SERVER],
                 'result_type': PlaybookTaskResultType.BASH_COMMAND_OUTPUT,
                 'display_name': 'Execute a BASH Command',
-                'category': 'Actions'
+                'category': 'Actions',
+                'form_fields': [
+                    FormField(key_name=StringValue(value="remote_server"),
+                              display_name=StringValue(value="Remote Server"),
+                              description=StringValue(value='Select Remote Server'),
+                              data_type=LiteralType.STRING),
+                    FormField(key_name=StringValue(value="command"),
+                              display_name=StringValue(value="Command"),
+                              data_type=LiteralType.STRING),
+                ]
             },
         }
 
@@ -48,7 +57,7 @@ def get_connector_processor(self, remote_server_connector, **kwargs):
             generated_credentials['remote_host'] = remote_server_str
         return RemoteServerProcessor(**generated_credentials)
 
-    def execute_command(self, time_range: TimeRange, global_variable_set: Dict, bash_task: Bash,
+    def execute_command(self, time_range: TimeRange, bash_task: Bash,
                         remote_server_connector: ConnectorProto) -> PlaybookTaskResult:
         try:
             bash_command: Bash.Command = bash_task.command
@@ -65,13 +74,6 @@ def execute_command(self, time_range: TimeRange, global_variable_set: Dict, bash
 
             command_str = bash_command.command.value
             commands = command_str.split('\n')
-            if global_variable_set:
-                for key, value in global_variable_set.items():
-                    updated_commands = []
-                    for command in commands:
-                        command = command.replace(key, str(value))
-                        updated_commands.append(command)
-                    commands = updated_commands
             try:
                 outputs = {}
                 ssh_client = self.get_connector_processor(remote_server_connector, remote_server_str=remote_server_str)