Fix sentinel TI provider (#797)

* Fix sentinel TI provider Provider was failiing trying to lookup table name in schema (which is not available) * Adding exception for notebookutils in test_pkg_imports
microsoft · Sep 12, 2024 · 2a81450 · 2a81450
1 parent c2869da
commit 2a81450
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/msticpy/context/tiproviders/kql_base.py b/msticpy/context/tiproviders/kql_base.py
@@ -16,6 +16,7 @@
 
 import abc
 import contextlib
+import logging
 import warnings
 from collections import defaultdict
 from functools import lru_cache
@@ -41,6 +42,8 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+logger = logging.getLogger(__name__)
+
 
 @export
 class KqlTIProvider(TIProvider):
@@ -162,9 +165,12 @@ def lookup_iocs(
         """
         if not self._connected:
             self._connect()
-        if any(
+        if self._query_provider.schema and any(
             table not in self._query_provider.schema for table in self._REQUIRED_TABLES
         ):
+            logger.error(
+                "Required tables not found in schema: %s", self._REQUIRED_TABLES
+            )
             return pd.DataFrame()
 
         # We need to partition the IoC types to invoke separate queries
@@ -175,6 +181,9 @@ def lookup_iocs(
             result = self._check_ioc_type(ioc, ioc_type, query_type)
 
             if result["Status"] != LookupStatus.NOT_SUPPORTED.value:
+                logger.info(
+                    "Check ioc type for %s (%s): %s", ioc, ioc_type, result["Status"]
+                )
                 ioc_groups[result["IocType"]].add(result["Ioc"])
 
         all_results: list[pd.DataFrame] = []
@@ -187,13 +196,15 @@ def lookup_iocs(
                     query_type=query_type,
                 )
             if not query_obj:
+                logger.info("No query found for %s", ioc_type)
                 warnings.warn(
                     f"Could not find query for {ioc_type}, {query_type}",
                     stacklevel=1,
                 )
                 continue
 
             # run the query
+            logger.info("Running query for %s with params %s", ioc_type, query_params)
             data_result: pd.DataFrame = query_obj(**query_params)
 
             src_ioc_frame: pd.DataFrame = pd.DataFrame(obs_set, columns=["Ioc"])
@@ -226,7 +237,9 @@ def lookup_iocs(
             all_results.append(combined_results_df)
 
         if all_results:
+            logger.info("Combining results from %d queries", len(all_results))
             return pd.concat(all_results, ignore_index=True, sort=False, axis=0)
+        logger.info("No results found in data for any iocs.")
         return pd.DataFrame()
 
     @staticmethod
@@ -318,12 +331,14 @@ def _create_query_provider(self: Self, **kwargs: str) -> tuple[QueryProvider, st
             WORKSPACE_ID=workspace_id,
         )
         query_provider: QueryProvider = QueryProvider("LogAnalytics")
+        logging.info("Connection string: %s", connect_str)
         return query_provider, connect_str
 
     def _connect(self: Self) -> None:
         """Connect to query provider."""
         print("MS Sentinel TI query provider needs authenticated connection.")
         self._query_provider.connect(self._connect_str)
+        logging.info("Connected to Sentinel. (%s)", self._connect_str)
 
     @staticmethod
     def _get_spelled_variants(name: str, **kwargs: str) -> str | None:

diff --git a/tests/test_pkg_imports.py b/tests/test_pkg_imports.py
@@ -29,6 +29,7 @@
     "pyperclip",
     "autogen",
     "importlib_resources",
+    "notebookutils",
 }
 CONDA_PKG_EXCEPTIONS = {
     "vt-py",