Added additional DataProvider for advanced hunting via Graph

msticpy/data/core/query_defns.py

@@ -96,6 +96,7 @@ class DataEnvironment(Enum):
     Kusto_New = 2  # alias of Kusto
     MSGraph = 4
     SecurityGraph = 4
+    GraphHunting = 5
     MDE = 5
     MDATP = 5  # alias of MDE
     LocalData = 6

msticpy/data/drivers/__init__.py

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

msticpy/data/drivers/mdatp_driver.py

@@ -18,6 +18,7 @@
 from ...common.utility import export
 from ..core.query_defns import DataEnvironment
 from .odata_driver import OData, QuerySource, _get_driver_settings
+from ...auth.azure_auth_core import AzureCloudConfig
 
 __version__ = VERSION
 __author__ = "Pete Bryan"
@@ -48,14 +49,15 @@ 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 +71,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.GraphHunting:
+            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 +113,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.GraphHunting:
+                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,8 +142,18 @@ def _select_api_uris(data_environment, cloud):
             f"{login_uri}{{tenantId}}/oauth2/token",
             "/advancedhunting/run",
         )
-    return (
-        get_defender_endpoint(cloud),
-        f"{login_uri}{{tenantId}}/oauth2/token",
-        "/advancedqueries/run",
-    )
+    elif data_environment == DataEnvironment.GraphHunting:
+        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",
+        )
+    else:
+        return (
+            get_defender_endpoint(cloud),
+            f"{login_uri}{{tenantId}}/oauth2/token",
+            "/advancedqueries/run",
+        )

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.")