Updates for msticpy compatibility (#46)

ianhelle · web-flow · commit e5a4f61b532e · 2024-09-20T12:58:49.000-07:00
* Updates for msticpy compatibility

Fixing data_providers to check for required parameter before creating class
ip_summary - handle either dataclass or tuple representation of whois result
ti.py - TI call using ioc_col instead of deprecated obs_col
minor updates to unit tests.

* Updating version to 1.2.2

* Fixing call to data providers and pandas compat
diff --git a/msticnb/_version.py b/msticnb/_version.py
@@ -1,2 +1,3 @@
 """Version file."""
-VERSION = "1.2.1"
+
+VERSION = "1.2.2"
diff --git a/msticnb/data_providers.py b/msticnb/data_providers.py
@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 """Data Providers class and init function."""
 import inspect
+import logging
 import sys
 from collections import namedtuple
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
@@ -351,7 +352,16 @@ def _query_prov(self, provider: str, provider_defn: ProviderDefn, **kwargs) -> A
             prov_kwargs_args = self._get_provider_kwargs(provider, **kwargs)
 
             # instantiate the provider class (sending all kwargs)
-            created_provider = provider_defn.prov_class(provider, **prov_kwargs_args)
+            provider_class = provider_defn.prov_class
+            if (
+                "data_environment"
+                in inspect.signature(provider_class.__init__).parameters
+            ):
+                prov_kwargs_args = {"data_environment": provider, **prov_kwargs_args}
+            logging.info(
+                "Creating provider %s with args %s", provider, prov_kwargs_args
+            )
+            created_provider = provider_defn.prov_class(**prov_kwargs_args)
             if created_provider.connected:
                 return created_provider
             # get the args required by connect function
diff --git a/msticnb/nb/azsent/host/win_host_events.py b/msticnb/nb/azsent/host/win_host_events.py
@@ -281,17 +281,39 @@ def _get_win_security_events(qry_prov, host_name, timespan):
 
 @set_text(docs=_CELL_DOCS, key="display_event_pivot")
 def _display_event_pivot(event_pivot):
-    display(
-        event_pivot.style.map(lambda x: "color: white" if x == 0 else "")
-        .map(
-            lambda x: "background-color: lightblue"
-            if not isinstance(x, str) and x > 0
-            else ""
+    if pd.__version__ < "2.1.0":
+        styled_data = (
+            event_pivot.style.applymap(lambda x: "color: white" if x == 0 else "")
+            .applymap(
+                lambda x: (
+                    "background-color: lightblue"
+                    if not isinstance(x, str) and x > 0
+                    else ""
+                )
+            )
+            .set_properties(
+                subset=["Activity"], **{"width": "400px", "text-align": "left"}
+            )
+            .highlight_max(axis=1)
+            .hide(axis="index")
         )
-        .set_properties(subset=["Activity"], **{"width": "400px", "text-align": "left"})
-        .highlight_max(axis=1)
-        .hide(axis="index")
-    )
+    else:
+        styled_data = (
+            event_pivot.style.map(lambda x: "color: white" if x == 0 else "")
+            .map(
+                lambda x: (
+                    "background-color: lightblue"
+                    if not isinstance(x, str) and x > 0
+                    else ""
+                )
+            )
+            .set_properties(
+                subset=["Activity"], **{"width": "400px", "text-align": "left"}
+            )
+            .highlight_max(axis=1)
+            .hide(axis="index")
+        )
+    display(styled_data)
 
 
 # %%
@@ -408,17 +430,39 @@ def _create_acct_event_pivot(account_event_data):
 
 @set_text(docs=_CELL_DOCS, key="display_acct_event_pivot")
 def _display_acct_event_pivot(event_pivot_df):
-    display(
-        event_pivot_df.style.map(lambda x: "color: white" if x == 0 else "")
-        .map(
-            lambda x: "background-color: lightblue"
-            if not isinstance(x, str) and x > 0
-            else ""
+    if pd.__version__ < "2.1.0":
+        styled_data = (
+            event_pivot_df.style.applymap(lambda x: "color: white" if x == 0 else "")
+            .applymap(
+                lambda x: (
+                    "background-color: lightblue"
+                    if not isinstance(x, str) and x > 0
+                    else ""
+                )
+            )
+            .set_properties(
+                subset=["Activity"], **{"width": "400px", "text-align": "left"}
+            )
+            .highlight_max(axis=1)
+            .hide(axis="index")
         )
-        .set_properties(subset=["Activity"], **{"width": "400px", "text-align": "left"})
-        .highlight_max(axis=1)
-        .hide(axis="index")
-    )
+    else:
+        styled_data = (
+            event_pivot_df.style.map(lambda x: "color: white" if x == 0 else "")
+            .map(
+                lambda x: (
+                    "background-color: lightblue"
+                    if not isinstance(x, str) and x > 0
+                    else ""
+                )
+            )
+            .set_properties(
+                subset=["Activity"], **{"width": "400px", "text-align": "left"}
+            )
+            .highlight_max(axis=1)
+            .hide(axis="index")
+        )
+    display(styled_data)
 
 
 @set_text(docs=_CELL_DOCS, key="display_acct_mgmt_timeline")
diff --git a/msticnb/nb/azsent/network/ip_summary.py b/msticnb/nb/azsent/network/ip_summary.py
@@ -963,7 +963,11 @@ def _populate_host_entity(result, geo_lookup=None):
 # Public IP functions
 def _get_whois(src_ip, result):
     """Get WhoIs data and split out networks."""
-    _, whois_dict = get_whois_info(src_ip)
+    whois_result = get_whois_info(src_ip)
+    if hasattr(whois_result, "properties"):
+        whois_dict = whois_result.properties
+    else:
+        whois_dict = whois_result[1]
     result.whois = pd.DataFrame(whois_dict)
     result.whois_nets = pd.DataFrame(whois_dict.get("nets", []))
     if df_has_data(result.whois):
diff --git a/msticnb/nblib/ti.py b/msticnb/nblib/ti.py
@@ -42,7 +42,7 @@ def get_ti_results(
     ti_merged_df = None
     iocs = data[col].dropna().unique()
     nb_markdown(f"Querying TI for {len(iocs)} indicators...")
-    ti_results = ti_lookup.lookup_iocs(data=data, obs_col=col)
+    ti_results = ti_lookup.lookup_iocs(data=data, ioc_col=col)
     if isinstance(ti_results, pd.DataFrame) and not ti_results.empty:
         ti_results = ti_results[ti_results["Severity"].isin(["warning", "high"])]
         ti_merged_df = data.merge(ti_results, how="inner", left_on=col, right_on="Ioc")
diff --git a/tests/nb/azsent/network/test_ip_summary.py b/tests/nb/azsent/network/test_ip_summary.py
@@ -130,7 +130,7 @@ def test_ip_summary_notebooklet(
     respx.get(re.compile(r"http://rdap\.arin\.net/.*")).respond(200, json=rdap_response)
     respx.get(
         re.compile(r"https://otx\.alienvault.*|https://www\.virustotal.*")
-    ).respond(200, json=_OTX_RESP)
+    ).respond(200, json=OTX_RESP)
     respx.get(re.compile(r"https://check\.torproject\.org.*")).respond(404)
     respx.get(re.compile(r"https://api\.greynoise\.io.*")).respond(404)
     respx.get(re.compile(r".*SecOps-Institute/Tor-IP-Addresses.*")).respond(
@@ -230,7 +230,7 @@ def test_ip_summary_notebooklet_all(
     respx.get(re.compile(r"http://rdap\.arin\.net/.*")).respond(200, json=rdap_response)
     respx.get(
         re.compile(r"https://otx\.alienvault.*|https://www\.virustotal.*")
-    ).respond(200, json=_OTX_RESP)
+    ).respond(200, json=OTX_RESP)
     respx.get(re.compile(r"https://check\.torproject\.org.*")).respond(404)
     respx.get(re.compile(r"https://api\.greynoise\.io.*")).respond(404)
     respx.get(re.compile(r".*SecOps-Institute/Tor-IP-Addresses.*")).respond(
@@ -296,7 +296,7 @@ def test_ip_summary_mde_data(
     respx.get(re.compile(r"http://rdap\.arin\.net/.*")).respond(200, json=rdap_response)
     respx.get(
         re.compile(r"https://otx\.alienvault.*|https://www\.virustotal.*")
-    ).respond(200, json=_OTX_RESP)
+    ).respond(200, json=OTX_RESP)
     respx.get(re.compile(r"https://check\.torproject\.org.*")).respond(404)
     respx.get(re.compile(r"https://api\.greynoise\.io.*")).respond(404)
     respx.get(re.compile(r".*SecOps-Institute/Tor-IP-Addresses.*")).respond(
@@ -331,7 +331,7 @@ def test_ip_summary_mde_data(
     check.is_instance(result.ti_results, pd.DataFrame)
 
 
-_OTX_RESP = {
+OTX_RESP = {
     "ioc_param": "url",
     "response": {
         "response": "Found stuff",
diff --git a/tests/nb/azsent/network/test_network_flow_summary.py b/tests/nb/azsent/network/test_network_flow_summary.py
@@ -27,6 +27,7 @@
     TILookupMock,
     get_test_data_path,
 )
+from .test_ip_summary import OTX_RESP
 
 # pylint: disable=no-member
 
@@ -89,6 +90,15 @@ def test_network_flow_summary_notebooklet(
     #     LocalData_query_paths=[test_data],
     # )
     respx.get(re.compile(r"http://rdap\.arin\.net/.*")).respond(200, json=rdap_response)
+    respx.get(
+        re.compile(r"https://otx\.alienvault.*|https://www\.virustotal.*")
+    ).respond(200, json=OTX_RESP)
+    respx.get(re.compile(r"https://check\.torproject\.org.*")).respond(404)
+    respx.get(re.compile(r"https://api\.greynoise\.io.*")).respond(404)
+    respx.get(re.compile(r".*SecOps-Institute/Tor-IP-Addresses.*")).respond(
+        200, content=b"12.34.56.78\n12.34.56.78\n12.34.56.78"
+    )
+    respx.get(re.compile(r"https://api\.greynoise\.io/.*")).respond(404)
 
     test_nb = nblts.azsent.network.NetworkFlowSummary()
     tspan = TimeSpan(period="1D")
diff --git a/tests/unit_test_lib.py b/tests/unit_test_lib.py
@@ -141,24 +141,24 @@ def lookup_ioc(
         for i in range(3):
             hit = random.randint(1, 10) > 5
 
-            result_args = dict(
-                Provider=f"TIProv-{i}",
-                Ioc=observable,
-                IocType=ioc_type,
-                QuerySubtype="mock",
-                Result=True,
-                Severity=2 if hit else 0,
-                Details=f"Details for {observable}",
-                RawResult=f"Raw details for {observable}",
-            )
+            result_args = {
+                "Provider": f"TIProv-{i}",
+                "Ioc": observable,
+                "IocType": ioc_type,
+                "QuerySubtype": "mock",
+                "Result": True,
+                "Severity": 2 if hit else 0,
+                "Details": f"Details for {observable}",
+                "RawResult": f"Raw details for {observable}",
+            }
             if check_mp_version("2.0"):
                 result_args["sanitized_value"] = observable
             else:
                 result_args["SafeIoC"] = observable
             result_list.append(result_args)
         return pd.DataFrame(result_list)
 
-    def lookup_iocs(self, data, obs_col: Optional[str] = None, **kwargs):
+    def lookup_iocs(self, data, ioc_col: Optional[str] = None, **kwargs):
         """Lookup fake TI."""
         del kwargs
         item_result: List[pd.DataFrame] = []
@@ -169,7 +169,7 @@ def lookup_iocs(self, data, obs_col: Optional[str] = None, **kwargs):
             )
         elif isinstance(data, pd.DataFrame):
             item_result.extend(
-                self.lookup_ioc(observable=row[obs_col]) for row in data.itertuples()
+                self.lookup_ioc(observable=row[ioc_col]) for row in data.itertuples()
             )
         elif isinstance(data, list):
             item_result.extend(self.lookup_ioc(observable=obs) for obs in data)
