Added additional DataProvider for Advanced Hunting via Graph (#725)

* Added additional DataProvider for advanced hunting via Graph

* Renamed Enum and set to new value

---------

Co-authored-by: Joey Dreijer <joey.dreijer@hema.nl>
Co-authored-by: Ian Hellen <ianhelle@microsoft.com>

msticpy/data/core/query_defns.py

@@ -96,6 +96,7 @@ class DataEnvironment(Enum):
     Kusto_New = 2  # alias of Kusto
     MSGraph = 4
     SecurityGraph = 4
+
     MDE = 5
     MDATP = 5  # alias of MDE
     LocalData = 6
@@ -115,6 +116,7 @@ class DataEnvironment(Enum):
     Kusto_KQLM = 17
     VelociraptorLogs = 18
     Velociraptor = 18
+    M365DGraph = 20
 
     @classmethod
     def parse(cls, value: Union[str, int]) -> "DataEnvironment":

msticpy/data/drivers/__init__.py

@@ -39,6 +39,7 @@
     ),
     DataEnvironment.MSSentinel_Legacy: ("kql_driver", "KqlDriver"),
     DataEnvironment.Kusto_Legacy: ("kusto_driver", "KustoDriver"),
+    DataEnvironment.M365DGraph: ("mdatp_driver", "MDATPDriver"),
 }
 
 CUSTOM_PROVIDERS: Dict[str, type] = {}

msticpy/data/drivers/mdatp_driver.py

@@ -9,6 +9,7 @@
 import pandas as pd
 
 from ..._version import VERSION
+from ...auth.azure_auth_core import AzureCloudConfig
 from ...auth.cloud_mappings import (
     get_defender_endpoint,
     get_m365d_endpoint,
@@ -48,14 +49,17 @@ def __init__(
         cs_dict = _get_driver_settings(
             self.CONFIG_NAME, self._ALT_CONFIG_NAMES, instance
         )
+
         self.cloud = cs_dict.pop("cloud", "global")
         if "cloud" in kwargs and kwargs["cloud"]:
             self.cloud = kwargs["cloud"]
 
         api_uri, oauth_uri, api_suffix = _select_api_uris(
             self.data_environment, self.cloud
         )
-        self.add_query_filter("data_environments", ("MDE", "M365D", "MDATP"))
+        self.add_query_filter(
+            "data_environments", ("MDE", "M365D", "MDATP", "GraphHunting")
+        )
 
         self.req_body = {
             "client_id": None,
@@ -69,6 +73,15 @@ def __init__(
         self.api_suffix = api_suffix
         if self.data_environment == DataEnvironment.M365D:
             self.scopes = [f"{api_uri}/AdvancedHunting.Read"]
+        elif self.data_environment == DataEnvironment.M365DGraph:
+            self.api_ver = kwargs.get("api_ver", "v1.0")
+            self.req_body = {
+                "client_id": None,
+                "client_secret": None,
+                "grant_type": "client_credentials",
+                "scope": f"{self.api_root}.default",
+            }
+            self.scopes = [f"{api_uri}/ThreatHunting.Read.All"]
         else:
             self.scopes = [f"{api_uri}/AdvancedQuery.Read"]
 
@@ -102,13 +115,21 @@ def query(
         )
         if isinstance(data, pd.DataFrame):
             # If we got a schema we should convert the DateTimes to pandas datetimes
-            if "Schema" not in response:
+            if ("Schema" or "schema") not in response:
                 return data
-            date_fields = [
-                field["Name"]
-                for field in response["Schema"]
-                if field["Type"] == "DateTime"
-            ]
+
+            if self.data_environment == DataEnvironment.M365DGraph:
+                date_fields = [
+                    field["name"]
+                    for field in response["schema"]
+                    if field["type"] == "DateTime"
+                ]
+            else:
+                date_fields = [
+                    field["Name"]
+                    for field in response["Schema"]
+                    if field["Type"] == "DateTime"
+                ]
             data = ensure_df_datetimes(data, columns=date_fields)
             return data
         return response
@@ -123,6 +144,15 @@ def _select_api_uris(data_environment, cloud):
             f"{login_uri}{{tenantId}}/oauth2/token",
             "/advancedhunting/run",
         )
+    if data_environment == DataEnvironment.M365DGraph:
+        az_cloud_config = AzureCloudConfig(cloud=cloud)
+        api_uri = az_cloud_config.endpoints.get("microsoftGraphResourceId")
+        graph_login = az_cloud_config.authority_uri
+        return (
+            api_uri,
+            f"{graph_login}{{tenantId}}/oauth2/v2.0/token",
+            "/security/runHuntingQuery",
+        )
     return (
         get_defender_endpoint(cloud),
         f"{login_uri}{{tenantId}}/oauth2/token",

msticpy/data/drivers/odata_driver.py

@@ -277,7 +277,8 @@ def query_with_results(self, query: str, **kwargs) -> Tuple[pd.DataFrame, Any]:
             )
             return None, json_response  # type: ignore
 
-        result = json_response.get("Results", json_response)
+        results_key = "Results" if "Results" in json_response else "results"
+        result = json_response.get(results_key, json_response)
 
         if not result:
             print("Warning - query did not return any results.")