From 2680ce2f7f0609fd9de753c3d4bef3d18c6e13f4 Mon Sep 17 00:00:00 2001 From: Ian Hellen Date: Fri, 12 Nov 2021 14:41:41 -0800 Subject: [PATCH] Removing azure_data requirement from host_summary notebooklet for test (#22) * Removing azure_data requirement from host_summary notebooklet for test * Updating azure-pipelines to python 3.8 type hints in data_providers * Linting errors * Another failing test in test_metadata.py * Addressing McCabe and Prospector warnings * Fixing test breaks in ti_enrich.py and account_summary.py Adding additional McCabe suppressions to deal with diff versions (sometimes McCabe IDs the start of decorated function as the decorator line, in newer versions, it uses the def line) * Fixing error caused by msticpy bug in ti_enrich * Reverting change to calling SelectAlert since it fails on 1.4.5 and earlier --- azure-pipelines.yml | 2 +- msticnb/class_doc.py | 4 +- msticnb/data_providers.py | 21 ++++++---- msticnb/nb/azsent/account/account_summary.py | 1 + msticnb/nb/azsent/alert/ti_enrich.py | 2 +- msticnb/nb/azsent/host/host_logons_summary.py | 4 +- msticnb/nb/azsent/host/host_summary.py | 42 ++++++++++++------- msticnb/nb/azsent/host/host_summary.yaml | 5 ++- msticnb/nb/azsent/network/ip_summary.py | 2 +- .../nb/azsent/network/network_flow_summary.py | 2 +- msticnb/nb_metadata.py | 2 +- msticnb/nblib/azsent/host.py | 2 +- msticnb/notebooklet.py | 6 ++- setup.py | 6 +-- tests/test_metadata.py | 4 +- tests/test_nb_pivot.py | 4 +- 16 files changed, 66 insertions(+), 43 deletions(-) diff --git a/azure-pipelines.yml b/azure-pipelines.yml index ab12c7e..15946d9 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -24,7 +24,7 @@ stages: pool: vmImage: $(imageName) variables: - python.version: '3.6' + python.version: '3.8' steps: # Add an alias for Windows python=>python3 - script: alias python='python3' pip='pip3' diff --git a/msticnb/class_doc.py b/msticnb/class_doc.py index 0744613..8192d15 100644 --- a/msticnb/class_doc.py +++ b/msticnb/class_doc.py @@ -200,10 +200,10 @@ def _format_func_doc(func_name, func, full_doc=False, prop_set=None): if func_doc: if not full_doc: # Get the first line of the doc string - doc_lines.append(func_doc.split("\n")[0]) + doc_lines.append(func_doc.split("\n", maxsplit=1)[0]) return doc_lines - func_doc = inspect.cleandoc(func_doc).split("\n")[0] + func_doc = inspect.cleandoc(func_doc).split("\n", maxsplit=1)[0] if func_doc: refmt_headings = [] for doc_line in func_doc.split("\n"): diff --git a/msticnb/data_providers.py b/msticnb/data_providers.py index fea058c..8d3f4cb 100644 --- a/msticnb/data_providers.py +++ b/msticnb/data_providers.py @@ -7,7 +7,7 @@ import inspect import sys from collections import namedtuple -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union from msticpy.common.exceptions import MsticpyAzureConfigError from msticpy.common.wsconfig import WorkspaceConfig @@ -336,7 +336,7 @@ def get_def_providers(cls) -> List[str]: """ return cls._DEFAULT_PROVIDERS - def _query_prov(self, provider, provider_defn, **kwargs): + def _query_prov(self, provider: str, provider_defn: ProviderDefn, **kwargs) -> Any: try: # Get any keys with the provider prefix and initialize the provider prov_kwargs_args = self._get_provider_kwargs(provider, **kwargs) @@ -353,25 +353,30 @@ def _query_prov(self, provider, provider_defn, **kwargs): if not prov_connect_args and provider_defn.get_config: prov_connect_args = provider_defn.get_config() # call the connect function - created_provider.connect(**prov_connect_args) + try: + created_provider.connect(**prov_connect_args) + except Exception as err: # pylint: disable=broad-except + print(f"Connection attempt for {provider} failed.\n{err}") return created_provider except MsticpyAzureConfigError as mp_ex: if get_opt("verbose"): print("Warning:", mp_ex.args) return None - def _no_connect_prov(self, provider, provider_defn, **kwargs): + def _no_connect_prov( + self, provider: str, provider_defn: ProviderDefn, **kwargs + ) -> Any: # Get the args passed to __init__ for this provider prov_args = self._get_provider_kwargs(provider, **kwargs) # If there are none and there's a config function, call that. if not prov_args and provider_defn.get_config: prov_args = provider_defn.get_config() - # Instatiate the provider + # Instantiate the provider return provider_defn.prov_class(prov_args) # Helper methods @staticmethod - def _get_provider_kwargs(prefix, **kwargs): + def _get_provider_kwargs(prefix: str, **kwargs) -> Dict[str, str]: """Return the kwargs prefixed with "prefix_".""" if prefix == "LogAnalytics" and any( key for key in kwargs if key.startswith("AzureSentinel") @@ -389,7 +394,7 @@ def _get_provider_kwargs(prefix, **kwargs): } @staticmethod - def _get_connect_args(func, **kwargs): + def _get_connect_args(func: Callable, **kwargs) -> Dict[str, str]: """Get the arguments required by the `connect` function.""" connect_params = inspect.signature(func).parameters return { @@ -398,7 +403,7 @@ def _get_connect_args(func, **kwargs): # Provider get_config functions @staticmethod - def _azsent_get_config(**kwargs): + def _azsent_get_config(**kwargs) -> Dict[str, str]: if "workspace" in kwargs: ws_config = WorkspaceConfig(workspace=kwargs["workspace"]) elif "config_file" in kwargs: diff --git a/msticnb/nb/azsent/account/account_summary.py b/msticnb/nb/azsent/account/account_summary.py index 199f4a9..9c71368 100644 --- a/msticnb/nb/azsent/account/account_summary.py +++ b/msticnb/nb/azsent/account/account_summary.py @@ -1178,6 +1178,7 @@ def _create_ip_summary(data, ip_col, geoip): .pipe( (get_geoip_whois, "data"), geo_lookup=geoip, ip_col=ip_col ) # get geoip and whois + .drop(columns=["TimeGenerated", "Type"], errors="ignore") .merge(data, left_on="IpAddress", right_on=ip_col) ) for col in group_cols: diff --git a/msticnb/nb/azsent/alert/ti_enrich.py b/msticnb/nb/azsent/alert/ti_enrich.py index 324c156..62489d8 100644 --- a/msticnb/nb/azsent/alert/ti_enrich.py +++ b/msticnb/nb/azsent/alert/ti_enrich.py @@ -334,7 +334,7 @@ def _color_cells(val): color = color_chart[val.casefold()] else: color = "none" - return "background-color: %s" % color + return f"background-color: {color}" def _sev_score(sev): diff --git a/msticnb/nb/azsent/host/host_logons_summary.py b/msticnb/nb/azsent/host/host_logons_summary.py index 1723375..ee46377 100644 --- a/msticnb/nb/azsent/host/host_logons_summary.py +++ b/msticnb/nb/azsent/host/host_logons_summary.py @@ -107,9 +107,9 @@ class HostLogonsSummary(Notebooklet): # pylint: disable=too-few-public-methods metadata = _CLS_METADATA - @set_text(docs=_CELL_DOCS, key="run") # noqa:MC0001 + @set_text(docs=_CELL_DOCS, key="run") # noqa: MC0001 # pylint: disable=too-many-locals, too-many-branches - def run( + def run( # noqa:MC0001 self, value: Any = None, data: Optional[pd.DataFrame] = None, diff --git a/msticnb/nb/azsent/host/host_summary.py b/msticnb/nb/azsent/host/host_summary.py index 9e742e0..7a5d5ea 100644 --- a/msticnb/nb/azsent/host/host_summary.py +++ b/msticnb/nb/azsent/host/host_summary.py @@ -106,8 +106,8 @@ class HostSummary(Notebooklet): _cell_docs = _CELL_DOCS # pylint: disable=too-many-branches - @set_text(docs=_CELL_DOCS, key="run") # noqa MC0001 - def run( + @set_text(docs=_CELL_DOCS, key="run") # noqa: MC0001 + def run( # noqa:MC0001 self, value: Any = None, data: Optional[pd.DataFrame] = None, @@ -206,7 +206,7 @@ def run( ) if azure_api: host_entity.AzureDetails["ResourceDetails"] = azure_api[ - "resoure_details" + "resource_details" ] host_entity.AzureDetails["SubscriptionDetails"] = azure_api[ "sub_details" @@ -254,39 +254,49 @@ def _azure_api_details(az_cli, host_record): ) # Get details of attached disks and network interfaces disks = [ - disk["name"] - for disk in resource_details["properties"]["storageProfile"]["dataDisks"] + disk.get("name") + for disk in resource_details.get("properties", {}) + .get("storageProfile", {}) + .get("dataDisks", {}) ] network_ints = [ - net["id"] - for net in resource_details["properties"]["networkProfile"][ - "networkInterfaces" - ] + net.get("id") + for net in resource_details.get("properties", {}) + .get("networkProfile", {}) + .get("networkInterfaces") ] image = ( str( - resource_details["properties"]["storageProfile"] + resource_details.get("properties", {}) + .get("storageProfile", {}) .get("imageReference", {}) .get("offer", {}) ) + " " + str( - resource_details["properties"]["storageProfile"] + resource_details.get("properties", {}) + .get("storageProfile", {}) .get("imageReference", {}) .get("sku", {}) ) ) # Extract key details and add host_entity resource_details = { - "Azure Location": resource_details["location"], - "VM Size": resource_details["properties"]["hardwareProfile"]["vmSize"], + "Azure Location": resource_details.get("location"), + "VM Size": ( + resource_details.get("properties", {}) + .get("hardwareProfile", {}) + .get("vmSize") + ), "Image": image, "Disks": disks, - "Admin User": resource_details["properties"]["osProfile"]["adminUsername"], + "Admin User": resource_details.get("properties", {}) + .get("osProfile", {}) + .get("adminUsername"), "Network Interfaces": network_ints, - "Tags": str(resource_details["tags"]), + "Tags": str(resource_details.get("tags")), } - return {"resoure_details": resource_details, "sub_details": sub_details} + return {"resource_details": resource_details, "sub_details": sub_details} except CloudError: return None diff --git a/msticnb/nb/azsent/host/host_summary.yaml b/msticnb/nb/azsent/host/host_summary.yaml index b3e9324..561c8fc 100644 --- a/msticnb/nb/azsent/host/host_summary.yaml +++ b/msticnb/nb/azsent/host/host_summary.yaml @@ -3,8 +3,9 @@ metadata: description: Host summary default_options: - heartbeat: Query Heartbeat table for host information. - - azure_net: 'Query AzureNetworkAnalytics table for host - network topology information.' + - azure_net: ' + Query AzureNetworkAnalytics table for host + network topology information.' - alerts: Query any alerts for the host. - bookmarks: Query any bookmarks for the host. - azure_api: Query Azure API for VM information. diff --git a/msticnb/nb/azsent/network/ip_summary.py b/msticnb/nb/azsent/network/ip_summary.py index f896317..fd61c17 100644 --- a/msticnb/nb/azsent/network/ip_summary.py +++ b/msticnb/nb/azsent/network/ip_summary.py @@ -181,7 +181,7 @@ class IpAddressSummary(Notebooklet): # pylint: disable=too-many-branches @set_text(docs=_CELL_DOCS, key="run") # noqa: MC0001 - def run( + def run( # noqa: MC0001 self, value: Any = None, data: Optional[pd.DataFrame] = None, diff --git a/msticnb/nb/azsent/network/network_flow_summary.py b/msticnb/nb/azsent/network/network_flow_summary.py index d486492..85a511d 100644 --- a/msticnb/nb/azsent/network/network_flow_summary.py +++ b/msticnb/nb/azsent/network/network_flow_summary.py @@ -156,7 +156,7 @@ def __init__(self, data_providers: Optional[DataProviders] = None, **kwargs): # pylint: disable=too-many-branches @set_text(docs=_CELL_DOCS, key="run") # noqa: MC0001 - def run( + def run( # noqa: MC0001 self, value: Any = None, data: Optional[pd.DataFrame] = None, diff --git a/msticnb/nb_metadata.py b/msticnb/nb_metadata.py index 12e77bb..eade501 100644 --- a/msticnb/nb_metadata.py +++ b/msticnb/nb_metadata.py @@ -165,7 +165,7 @@ def _read_metadata_file(mod_path): if not md_path.is_file(): md_path = Path(str(mod_path).replace(".py", ".yml")) if md_path.is_file(): - with open(md_path, "r") as _md_file: + with open(md_path, "r", encoding="utf-8") as _md_file: return yaml.safe_load(_md_file) return None diff --git a/msticnb/nblib/azsent/host.py b/msticnb/nblib/azsent/host.py index fc330ba..065aa9f 100644 --- a/msticnb/nblib/azsent/host.py +++ b/msticnb/nblib/azsent/host.py @@ -100,7 +100,7 @@ def get_aznet_topology( @lru_cache() # noqa:MC0001 -def verify_host_name( +def verify_host_name( # noqa: MC0001 qry_prov: QueryProvider, host_name: str, timespan: TimeSpan = None, **kwargs ) -> HostNameVerif: """ diff --git a/msticnb/notebooklet.py b/msticnb/notebooklet.py index e6c0ea2..0d87b7b 100644 --- a/msticnb/notebooklet.py +++ b/msticnb/notebooklet.py @@ -455,7 +455,7 @@ def _get_timespan(self, timespan=None, **kwargs): def import_cell(cls): """Import the text of this module into a new cell.""" if cls.module_path: - with open(cls.module_path, "r") as mod_file: + with open(cls.module_path, "r", encoding="utf-8") as mod_file: mod_text = mod_file.read() if mod_text: # replace relative references with absolute paths @@ -553,7 +553,9 @@ def check_table_exists(self, table: str) -> bool: def get_methods(self) -> Dict[str, Callable[[Any], Any]]: """Return methods available for this class.""" meths = inspect.getmembers(self, inspect.ismethod) - cls_selector = f"bound method {self.__class__.__name__.rsplit('.')[0]}" + cls_selector = ( + f"bound method {self.__class__.__name__.rsplit('.', maxsplit=1)[0]}" + ) return { meth[0]: meth[1] for meth in meths diff --git a/setup.py b/setup.py index 139da3d..a1e8db2 100644 --- a/setup.py +++ b/setup.py @@ -9,14 +9,14 @@ import setuptools -with open("requirements.txt", "r") as fh: +with open("requirements.txt", "r", encoding="utf-8") as fh: INSTALL_REQUIRES = fh.readlines() # pylint: disable=locally-disabled, invalid-name -with open("README.md", "r") as fh: +with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() -with open("msticnb/_version.py", "r") as fd: +with open("msticnb/_version.py", "r", encoding="utf-8") as fd: v_match = re.search(r'^VERSION\s*=\s*[\'"]([^\'"]*)[\'"]', fd.read(), re.MULTILINE) __version__ = v_match.group(1) if v_match else "no version" # pylint: enable=locally-disabled, invalid-name diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 342976d..f24a93a 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -5,7 +5,7 @@ # -------------------------------------------------------------------------- """NB metadata test class.""" import pytest_check as check -from msticnb import data_providers +from msticnb import data_providers, nblts from msticnb.nb_metadata import NBMetadata, read_mod_metadata from msticnb.nb.azsent.host import host_summary @@ -33,6 +33,8 @@ def test_read_metadata(): def test_class_metadata(monkeypatch): """Test class correctly loads yaml metadata.""" monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock) + if "azuredata" in nblts.azsent.host.HostSummary.metadata.req_providers: + nblts.azsent.host.HostSummary.metadata.req_providers.remove("azuredata") data_providers.init( query_provider="LocalData", providers=["tilookup", "geolitelookip"] ) diff --git a/tests/test_nb_pivot.py b/tests/test_nb_pivot.py index 0014a36..0a89329 100644 --- a/tests/test_nb_pivot.py +++ b/tests/test_nb_pivot.py @@ -12,7 +12,7 @@ from msticpy.datamodel import entities from msticpy.datamodel.pivot import Pivot -from msticnb import data_providers +from msticnb import data_providers, nblts from msticnb.nb_pivot import add_pivot_funcs from msticnb.notebooklet import NotebookletResult @@ -42,6 +42,8 @@ def _init_pivot(monkeypatch): test_data = str(Path(TEST_DATA_PATH).absolute()) monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock) + if "azuredata" in nblts.azsent.host.HostSummary.metadata.req_providers: + nblts.azsent.host.HostSummary.metadata.req_providers.remove("azuredata") data_providers.init( query_provider="LocalData", providers=["geolitelookup"],