From f351fe3a0892aff8fdbb17a55f19de95fd2a1ff4 Mon Sep 17 00:00:00 2001
From: Ariel <ariel.ropek@panther.com>
Date: Wed, 27 Nov 2024 11:24:08 -0700
Subject: [PATCH] PantherFlow Investigator Helper

---
 global_helpers/panther_aws_helpers.py  |  2 ++
 global_helpers/panther_base_helpers.py | 24 ++++++++++++++++++++++++
 global_helpers/panther_okta_helpers.py |  4 ++++
 3 files changed, 30 insertions(+)

diff --git a/global_helpers/panther_aws_helpers.py b/global_helpers/panther_aws_helpers.py
index 738db8cb3..1b3927838 100644
--- a/global_helpers/panther_aws_helpers.py
+++ b/global_helpers/panther_aws_helpers.py
@@ -4,6 +4,7 @@
 from typing import Any, Dict, List
 
 import boto3
+from panther_base_helpers import pantherflow_investigation
 from panther_config import config
 
 
@@ -38,6 +39,7 @@ def aws_rule_context(event):
         "sourceIPAddress": event.get("sourceIPAddress", "<MISSING_SOURCE_IP>"),
         "userAgent": event.get("userAgent", "<MISSING_USER_AGENT>"),
         "userIdentity": event.get("userIdentity", "<MISSING_USER_IDENTITY>"),
+        "PantherFlow Investigation": pantherflow_investigation(event),
     }
 
 
diff --git a/global_helpers/panther_base_helpers.py b/global_helpers/panther_base_helpers.py
index 6a0970d3b..ba8c2cd9b 100644
--- a/global_helpers/panther_base_helpers.py
+++ b/global_helpers/panther_base_helpers.py
@@ -327,3 +327,27 @@ def add_parse_delay(event, context: dict) -> dict:
     parsing_delay = time_delta(event.get("p_event_time"), event.get("p_parse_time"))
     context["parseDelay"] = f"{parsing_delay}"
     return context
+
+
+# generate a PantherFlow investigation from an event
+def pantherflow_investigation(event, interval="30m"):
+    logtype = event.get("p_log_type", "").lower().replace(".", "_")
+    timestamp = event.get("p_event_time", "")
+
+    query = f"""
+union panther_signals.public.correlation_signals
+    , panther_logs.public.{logtype}
+| where p_event_time between datetime('{timestamp}') - time.parse_timespan('{interval}') .. datetime('{timestamp}') + time.parse_timespan('{interval}')
+"""
+
+    first = True
+    for key, value in event.items():
+        if key.startswith("p_any_") and key != "p_any_aws_account_ids":
+            if first:
+                query += f"| where arrays.overlap({key}, {value})\n"
+                first = False
+            else:
+                query += f"     or arrays.overlap({key}, {value})\n"
+    query += "| sort p_event_time\n"
+
+    return query
diff --git a/global_helpers/panther_okta_helpers.py b/global_helpers/panther_okta_helpers.py
index ffbfb8af8..8715fc777 100644
--- a/global_helpers/panther_okta_helpers.py
+++ b/global_helpers/panther_okta_helpers.py
@@ -1,3 +1,6 @@
+from panther_base_helpers import pantherflow_investigation
+
+
 def okta_alert_context(event):
     """Returns common context for automation of Okta alerts"""
     return {
@@ -12,4 +15,5 @@ def okta_alert_context(event):
         "authentication_context": event.get("authenticationcontext", {}),
         "security_context": event.get("securitycontext", {}),
         "ips": event.get("p_any_ip_addresses", []),
+        "PantherFlow Investigation": pantherflow_investigation(event),
     }