Compat fixes for Bokeh 3.7 (#840)

* Compat fixes for Bokeh 3.7
Also suppressing classes of mypy errors in bokeh modules since it is broadly incompatible with Bokeh.
Fixing flake8 error in common/utilty/types

* Fix for vt-py changing structure VTObject dictionary - causing conversion to dataframe to contain hundreds of columns.
Limiting json_normalize to max_levels=0

* Fix for using AzCli credential with Managed Identity

* Mypy fix for azure_auth_core.py

Author: ianhelle

msticpy/_version.py

@@ -1,3 +1,3 @@
 """Version file."""
 
-VERSION = "2.16.1"
+VERSION = "2.16.2"

msticpy/auth/azure_auth_core.py

@@ -4,6 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 """Azure KeyVault pre-authentication."""
+
 from __future__ import annotations
 
 import logging
@@ -16,6 +17,7 @@
 
 from azure.common.credentials import get_cli_profile
 from azure.core.credentials import TokenCredential
+from azure.core.exceptions import ClientAuthenticationError
 from azure.identity import (
     AzureCliCredential,
     AzurePowerShellCredential,
@@ -150,27 +152,62 @@ def _build_cli_client(
 ) -> AzureCliCredential:
     """Build a credential from Azure CLI."""
     del kwargs
-    if tenant_id:
-        return AzureCliCredential(tenant_id=tenant_id)
-    return AzureCliCredential()
+    if tenant_id is not None:
+        try:
+            logger.info("Creating Azure CLI credential with tenant_id")
+            cred = AzureCliCredential(tenant_id=tenant_id)
+            # Attempt to get a token immediately to validate the credential
+            cred.get_token("https://management.azure.com/.default")
+            return cred
+        except ClientAuthenticationError as ex:
+            logger.info("Azure CLI credential failed to authenticate: %s", str(ex))
+            # Check if the error is related to tenant ID
+            if "Tenant" not in str(ex).lower():
+                raise  # re-raise if it's a different error
+    logger.info("Creating Azure CLI credential without tenant_id")
+    cred = AzureCliCredential()
+    cred.get_token("https://management.azure.com/.default")
+    return cred
 
 
 def _build_msi_client(
     tenant_id: str | None = None,
-    aad_uri: str | None = None,
     client_id: str | None = None,
     **kwargs,
 ) -> ManagedIdentityCredential:
     """Build a credential from Managed Identity."""
     msi_kwargs: dict[str, Any] = kwargs.copy()
     client_id = client_id or os.environ.get(AzureCredEnvNames.AZURE_CLIENT_ID)
 
-    return ManagedIdentityCredential(
-        tenant_id=tenant_id,
-        authority=aad_uri,
-        client_id=client_id,
-        **msi_kwargs,
-    )
+    try:
+        cred = ManagedIdentityCredential(client_id=client_id)
+        cred.get_token("https://management.azure.com/.default")
+        return cred
+    except ClientAuthenticationError as ex:
+        logger.info(
+            (
+                "Managed Identity credential failed to authenticate: %s, retrying with args "
+                "tenant_id=%s, client_id=%s, kwargs=%s"
+            ),
+            str(ex),
+            tenant_id,
+            client_id,
+            msi_kwargs,
+        )
+
+        try:
+            # Retry passing previous parameter set
+            cred = ManagedIdentityCredential(
+                client_id=client_id, tenant_id=tenant_id, **msi_kwargs
+            )
+            cred.get_token("https://management.azure.com/.default")
+            return cred
+        except ClientAuthenticationError:
+            # If we fail again, just create with no params
+            logger.info(
+                "Managed Identity credential failed auth - retrying with no params"
+            )
+            return ManagedIdentityCredential()
 
 
 def _build_vscode_client(
@@ -380,7 +417,8 @@ def _az_connect_core(
     # Create the wrapped credential using the passed credential
     wrapped_credentials = CredentialWrapper(credential, resource_id=az_config.token_uri)
     return AzCredentials(
-        wrapped_credentials, ChainedTokenCredential(credential)  # type: ignore[arg-type]
+        wrapped_credentials,  # type: ignore[arg-type]
+        ChainedTokenCredential(credential),  # type: ignore[arg-type]
     )
 
 

msticpy/common/utility/types.py

@@ -280,7 +280,6 @@ def singleton(cls: type) -> Callable:
     instances: dict[type[object], object] = {}
 
     def get_instance(*args, **kwargs) -> object:
-        nonlocal instances
         if cls not in instances:
             instances[cls] = cls(*args, **kwargs)
         return instances[cls]

msticpy/context/vtlookupv3/vtlookupv3.py

@@ -191,7 +191,7 @@ def _parse_vt_object(
                 }
             else:
                 obj = attributes
-            vt_df: pd.DataFrame = pd.json_normalize(data=[obj])
+            vt_df: pd.DataFrame = pd.json_normalize(data=[obj], max_level=0)
             last_analysis_stats: dict[str, Any] | None = attributes.get(
                 VTObjectProperties.LAST_ANALYSIS_STATS.value,
             )
@@ -207,7 +207,7 @@ def _parse_vt_object(
             # Format dates for pandas
             vt_df = timestamps_to_utcdate(vt_df)
         elif obj_dict:
-            vt_df = pd.json_normalize([obj_dict])
+            vt_df = pd.json_normalize([obj_dict], max_level=0)
         else:
             vt_df = cls._item_not_found_df(
                 vt_type=vt_object.type,
@@ -885,7 +885,7 @@ def get_object(self: Self, vt_id: str, vt_type: str) -> pd.DataFrame:
                     "type": [response.type],
                 },
             )
-            attribs = pd.json_normalize(response.to_dict()["attributes"])
+            attribs = pd.json_normalize(response.to_dict()["attributes"], max_level=0)
             result_df = pd.concat([result_df, attribs], axis=1)
             result_df["context_attributes"] = response.to_dict().get(
                 "context_attributes",
@@ -1051,7 +1051,9 @@ def _extract_response(self: Self, response_list: list) -> pd.DataFrame:
         response_rows = []
         for response_item in response_list:
             # flatten nested dictionary and append id, type values
-            response_item_df = pd.json_normalize(response_item["attributes"])
+            response_item_df = pd.json_normalize(
+                response_item["attributes"], max_level=0
+            )
             response_item_df["id"] = response_item["id"]
             response_item_df["type"] = response_item["type"]
 

msticpy/vis/entity_graph_tools.py

@@ -33,6 +33,9 @@
 __version__ = VERSION
 __author__ = "Pete Bryan"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="arg-type"
+
 req_alert_cols = ["DisplayName", "Severity", "AlertType"]
 req_inc_cols = ["id", "name", "properties.severity"]
 
@@ -140,6 +143,7 @@ def _plot_with_timeline(self, hide: bool = False, **kwargs) -> LayoutDOM:
         """
         timeline = None
         tl_df = self.to_df()
+
         tl_type = "duration"
         # pylint: disable=unsubscriptable-object
         if len(tl_df["EndTime"].unique()) == 1 and not tl_df["EndTime"].unique()[0]:
@@ -150,22 +154,22 @@ def _plot_with_timeline(self, hide: bool = False, **kwargs) -> LayoutDOM:
             ):
                 print("No timestamps available to create timeline")
                 return self._plot_no_timeline(timeline=False, hide=hide, **kwargs)
-        # tl_df["TimeGenerated"] = pd.to_datetime(tl_df["TimeGenerated"], utc=True)
-        # tl_df["StartTime"] = pd.to_datetime(tl_df["StartTime"], utc=True)
-        # tl_df["EndTime"] = pd.to_datetime(tl_df["EndTime"], utc=True)
+
         graph = self._plot_no_timeline(hide=True, **kwargs)
         if tl_type == "duration":
+            # remove missing time values
             timeline = display_timeline_duration(
-                tl_df.dropna(subset=["TimeGenerated"]),
+                tl_df.dropna(subset=["StartTime", "EndTime"]),
                 group_by="Name",
                 title="Entity Timeline",
                 time_column="StartTime",
                 end_time_column="EndTime",
-                source_columns=["Name", "Description", "Type", "TimeGenerated"],
+                source_columns=["Name", "Description", "Type", "StartTime", "EndTime"],
                 hide=True,
                 width=800,
             )
         elif tl_type == "discreet":
+            tl_df = tl_df.dropna(subset=["TimeGenerated"])
             timeline = display_timeline(
                 tl_df.dropna(subset=["TimeGenerated"]),
                 group_by="Type",

msticpy/vis/matrix_plot.py

@@ -21,6 +21,9 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="call-arg, attr-defined"
+
 # wrap figure function to handle v2/v3 parameter renaming
 figure = bokeh_figure(figure)  # type: ignore[assignment, misc]
 

msticpy/vis/network_plot.py

@@ -33,6 +33,9 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="arg-type"
+
 _BOKEH_VERSION: Version = parse(version("bokeh"))
 
 # wrap figure function to handle v2/v3 parameter renaming

msticpy/vis/process_tree.py

@@ -79,6 +79,9 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="arg-type, call-arg, attr-defined"
+
 _DEFAULT_KWARGS = ["height", "title", "width", "hide_legend", "pid_fmt"]
 
 # wrap figure function to handle v2/v3 parameter renaming

msticpy/vis/timeline.py

@@ -43,6 +43,9 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="arg-type, call-arg"
+
 # wrap figure function to handle v2/v3 parameter renaming
 figure = bokeh_figure(figure)  # type: ignore[assignment, misc]
 

msticpy/vis/timeline_common.py

@@ -42,6 +42,9 @@
 __version__ = VERSION
 __author__ = "Ian Hellen"
 
+# mypy and Bokeh are not best friends
+# mypy: disable-error-code="arg-type, call-arg, attr-defined"
+
 TIMELINE_HELP = (
     "https://msticpy.readthedocs.io/en/latest/msticpy.vis.html"
     "#msticpy.vis.timeline.{plot_type}"
@@ -298,10 +301,11 @@ def create_range_tool(
                 y=y,
                 color=series_def["color"],
                 source=series_def["source"],
+                radius=1,
             )
     elif isinstance(data, pd.DataFrame):
         rng_select.circle(
-            x=time_column, y=y, color="blue", source=ColumnDataSource(data)
+            x=time_column, y=y, color="blue", source=ColumnDataSource(data), radius=1
         )
 
     range_tool = RangeTool(x_range=plot_range)