Skip to content

Commit 2a81450

Browse files
authored
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
1 parent c2869da commit 2a81450

File tree

2 files changed

+17
-1
lines changed

2 files changed

+17
-1
lines changed

msticpy/context/tiproviders/kql_base.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import abc
1818
import contextlib
19+
import logging
1920
import warnings
2021
from collections import defaultdict
2122
from functools import lru_cache
@@ -41,6 +42,8 @@
4142
__version__ = VERSION
4243
__author__ = "Ian Hellen"
4344

45+
logger = logging.getLogger(__name__)
46+
4447

4548
@export
4649
class KqlTIProvider(TIProvider):
@@ -162,9 +165,12 @@ def lookup_iocs(
162165
"""
163166
if not self._connected:
164167
self._connect()
165-
if any(
168+
if self._query_provider.schema and any(
166169
table not in self._query_provider.schema for table in self._REQUIRED_TABLES
167170
):
171+
logger.error(
172+
"Required tables not found in schema: %s", self._REQUIRED_TABLES
173+
)
168174
return pd.DataFrame()
169175

170176
# We need to partition the IoC types to invoke separate queries
@@ -175,6 +181,9 @@ def lookup_iocs(
175181
result = self._check_ioc_type(ioc, ioc_type, query_type)
176182

177183
if result["Status"] != LookupStatus.NOT_SUPPORTED.value:
184+
logger.info(
185+
"Check ioc type for %s (%s): %s", ioc, ioc_type, result["Status"]
186+
)
178187
ioc_groups[result["IocType"]].add(result["Ioc"])
179188

180189
all_results: list[pd.DataFrame] = []
@@ -187,13 +196,15 @@ def lookup_iocs(
187196
query_type=query_type,
188197
)
189198
if not query_obj:
199+
logger.info("No query found for %s", ioc_type)
190200
warnings.warn(
191201
f"Could not find query for {ioc_type}, {query_type}",
192202
stacklevel=1,
193203
)
194204
continue
195205

196206
# run the query
207+
logger.info("Running query for %s with params %s", ioc_type, query_params)
197208
data_result: pd.DataFrame = query_obj(**query_params)
198209

199210
src_ioc_frame: pd.DataFrame = pd.DataFrame(obs_set, columns=["Ioc"])
@@ -226,7 +237,9 @@ def lookup_iocs(
226237
all_results.append(combined_results_df)
227238

228239
if all_results:
240+
logger.info("Combining results from %d queries", len(all_results))
229241
return pd.concat(all_results, ignore_index=True, sort=False, axis=0)
242+
logger.info("No results found in data for any iocs.")
230243
return pd.DataFrame()
231244

232245
@staticmethod
@@ -318,12 +331,14 @@ def _create_query_provider(self: Self, **kwargs: str) -> tuple[QueryProvider, st
318331
WORKSPACE_ID=workspace_id,
319332
)
320333
query_provider: QueryProvider = QueryProvider("LogAnalytics")
334+
logging.info("Connection string: %s", connect_str)
321335
return query_provider, connect_str
322336

323337
def _connect(self: Self) -> None:
324338
"""Connect to query provider."""
325339
print("MS Sentinel TI query provider needs authenticated connection.")
326340
self._query_provider.connect(self._connect_str)
341+
logging.info("Connected to Sentinel. (%s)", self._connect_str)
327342

328343
@staticmethod
329344
def _get_spelled_variants(name: str, **kwargs: str) -> str | None:

tests/test_pkg_imports.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
"pyperclip",
3030
"autogen",
3131
"importlib_resources",
32+
"notebookutils",
3233
}
3334
CONDA_PKG_EXCEPTIONS = {
3435
"vt-py",

0 commit comments

Comments
 (0)