Merge pull request #16 from microsoft/ianhelle/pivot-related-fixes-20…

…21-06-09

Fix to account_summary to handle NAs in values

azure-pipelines.yml

@@ -41,6 +41,8 @@ stages:
         echo Using $MSTICPYCONFIG or %MSTICPYCONFIG%
         pip install --upgrade pytest pytest-azurepipelines
         pip install --upgrade pytest-cov pytest-check
+        # need these two packages for tests
+        pip install matplotlib scikit-learn
         pytest tests --junitxml=junit/test-results.xml --cov=msticnb --cov-report=xml
       continueOnError: true
       condition: succeededOrFailed()
@@ -52,7 +54,7 @@ stages:
         MSTICPY_TEST_NOSKIP: 1
     - script: |
         pip install --upgrade black
-        black -t py36 --check --exclude venv . 
+        black -t py36 --check --exclude venv .
       displayName: Black
       continueOnError: true
       condition: succeededOrFailed()
@@ -69,7 +71,7 @@ stages:
       continueOnError: true
       condition: succeededOrFailed()
     - script: |
-        pip install --upgrade flake8 
+        pip install --upgrade flake8
         flake8 --max-line-length=90 --exclude=tests* . --ignore=E501,W503
       displayName: flake8
       continueOnError: true

msticnb/nb/azsent/account/account_summary.py

@@ -1172,6 +1172,7 @@ def _create_ip_summary(data, ip_col, geoip):
     group_cols = group_cols1 + group_cols2
     all_data = (
         data[[ip_col]]  # the property and the column we want
+        .dropna()
         .drop_duplicates()  # drop duplicates
         .pipe(
             (get_geoip_whois, "data"), geo_lookup=geoip, ip_col=ip_col

msticnb/nb/azsent/host/logon_session_rarity.yaml

@@ -1,6 +1,7 @@
 metadata:
   name: LogonSessionRarity
   description: Calculates sessions with most unusual process activity.
+  inputs: [data]
   default_options:
   other_options:
   keywords:

msticnb/nb/azsent/network/ip_summary.py

@@ -374,7 +374,7 @@ def _get_azure_netflow(self, src_ip, result, timespan):
         """Retrieve Azure netflow and activity events."""
         if self.check_table_exists("AzureNetworkAnalytics_CL"):
             _get_az_netflows(self.query_provider, src_ip, result, timespan)
-            _display_df_summary(result.az_network_flows, "Azure network flows")
+            _display_df_summary(result.az_network_flows, "Azure NSG network flows")
 
     @set_text(docs=_CELL_DOCS, key="get_az_activity")
     def _get_azure_activity(self, src_ip, result, timespan):
@@ -458,7 +458,7 @@ def _determine_ip_origin(result):
 # %%
 # Get Azure network flows
 def _get_az_netflows(qry_prov, src_ip, result, timespan):
-    nb_data_wait("AzureNetworkAnalytics flows")
+    nb_data_wait("Azure NSG flows")
     result.az_network_flows = qry_prov.Network.list_azure_network_flows_by_ip(
         timespan, ip_address_list=src_ip
     )
@@ -590,13 +590,13 @@ def _summarize_azure_activity(result):
 @set_text(docs=_CELL_DOCS, key="get_az_net_if")
 def _get_az_net_if(qry_prov, src_ip, result):
     """Get the AzureNetwork topology record for `src_ip`."""
-    nb_data_wait("AzureNetworkAnalytics topology")
+    nb_data_wait("Azure NSG topology")
     # Try to find the interface topology log entry
     result.az_network_if = qry_prov.Network.get_host_for_ip(  # type:ignore
         ip_address=src_ip
     )
     if not df_has_data(result.az_network_if):
-        nb_markdown("Could not get Azure network interface record")
+        nb_markdown("Could not get Azure NSG network interface record")
 
 
 @set_text(docs=_CELL_DOCS, key="get_heartbeat")

msticnb/nb_metadata.py

@@ -27,6 +27,7 @@ class NBMetadata:
     description: str = ""
     default_options: List[Union[str, Dict]] = Factory(list)
     other_options: List[Union[str, Dict]] = Factory(list)
+    inputs: List[str] = ["value"]
     entity_types: List[str] = Factory(list)
     keywords: List[str] = Factory(list)
     req_providers: List[str] = Factory(list)

msticnb/nb_pivot.py

@@ -51,6 +51,9 @@ def add_pivot_funcs(pivot: Pivot = None, **kwargs):
     for nb_name, nb_class in nblts.iter_classes():
         if not issubclass(nb_class, Notebooklet) or nb_name == "TemplateNB":
             continue
+        if "value" not in nb_class.metadata.inputs:
+            # This doesn't take a "value" input so can't use as a pivot
+            continue
         nb_obj = nb_class()
         run_func = getattr(nb_obj, "run")
         wrp_func = _wrap_run_func(run_func, pivot.get_timespan)

mypy.ini

@@ -81,3 +81,9 @@ ignore_missing_imports = True
 
 [mypy-pytest.*]
 ignore_missing_imports = True
+
+[mypy-yaml.*]
+ignore_missing_imports = True
+
+[mypy-requests.*]
+ignore_missing_imports = True

tests/nb/azsent/host/test_hostlogonsummary.py

@@ -4,6 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 """Test case for hostslogonsummary nblet."""
+import sys
 from datetime import datetime
 from pathlib import Path
 
@@ -20,6 +21,11 @@
 
 # nosec
 # pylint: disable=no-member
+if not sys.platform.startswith("win"):
+    pytest.skip(
+        "skipping Linux and Mac for these tests since Matplotlib fails with no gui",
+        allow_module_level=True,
+    )
 
 
 @pytest.fixture

tests/nb/azsent/host/test_logon_session_rarity.py

@@ -33,7 +33,7 @@ def test_logon_session_rarity_notebooklet(monkeypatch):
 
     check.is_true(hasattr(nblts.azsent.host, "LogonSessionsRarity"))
     if not hasattr(nblts.azsent.host, "LogonSessionsRarity"):
-        print(nblts.azsent.host())
+        print(nblts.azsent.host)
     test_nb = nblts.azsent.host.LogonSessionsRarity()
 
     result = test_nb.run(data=data)