Some test and linter fixes

.pre-commit-config.yaml

@@ -1,13 +1,31 @@
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
+  - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v2.3.0
     hooks:
     -   id: check-yaml
     -   id: check-json
     -   id: trailing-whitespace
         args: [--markdown-linebreak-ext=md]
--   repo: https://github.com/ambv/black
-    rev: stable
+  - repo: https://github.com/ambv/black
+    rev: 20.8b1
     hooks:
-    - id: black
-      language: python
+      - id: black
+        language: python
+        args:
+          - -t
+          - py36
+  - repo: https://github.com/pre-commit/mirrors-pylint
+    rev: v2.6.0
+    hooks:
+      - id: pylint
+        args:
+          - --disable=E0401
+          - --ignore-patterns=test_
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.8.4
+    hooks:
+      - id: flake8
+        args:
+          - --extend-ignore=E0401,E501
+          - --max-line-length=90
+          - --exclude=tests,test*.py

msticnb/nb/azsent/account/account_summary.py

@@ -22,7 +22,6 @@
     nb_display,
     nb_markdown,
     set_text,
-    show_bokeh,
     df_has_data,
 )
 from ....nblib.azsent.alert import browse_alerts

msticnb/nb/azsent/network/ip_summary.py

@@ -41,7 +41,7 @@
 _CLS_METADATA, _CELL_DOCS = nb_metadata.read_mod_metadata(__file__, __name__)
 
 
-# pylint: disable=too-few-public-methods
+# pylint: disable=too-few-public-methods, too-many-instance-attributes
 # Rename this class
 class IpSummaryResult(NotebookletResult):
     """
@@ -154,7 +154,7 @@ def __init__(
         self.passive_dns: pd.DataFrame = None
 
 
-# pylint: enable=too-few-public-methods
+# pylint: enable=too-few-public-methods, too-many-instance-attributes
 
 
 # Rename this class
@@ -178,10 +178,7 @@ class IpAddressSummary(Notebooklet):
     __doc__ = nb_metadata.update_class_doc(__doc__, metadata)
     _cell_docs = _CELL_DOCS
 
-    # @set_text decorator will display the title and text every time
-    # this method is run.
-    # The key value refers to an entry in the `output` section of
-    # the notebooklet yaml file.
+    # pylint: disable=too-many-branches
     @set_text(docs=_CELL_DOCS, key="run")
     def run(
         self,
@@ -311,19 +308,25 @@ def browse_ti_results(self):
             return ti_lookup.browse_results(self._last_result.ti_results)
         return None
 
-    def netflow_by_protocol(self,) -> Figure:
+    def netflow_by_protocol(
+        self,
+    ) -> Figure:
         """Display netflows grouped by protocol."""
         if not self.check_valid_result_data("az_network_flows"):
             return None
         return _plot_netflow_by_protocol(self._last_result)
 
-    def netflow_total_by_protocol(self,) -> Figure:
+    def netflow_total_by_protocol(
+        self,
+    ) -> Figure:
         """Display netflows grouped by protocol."""
         if not self.check_valid_result_data("az_network_flows"):
             return None
         return _plot_netflow_values_by_protocol(self._last_result)
 
-    def netflow_by_direction(self,) -> Figure:
+    def netflow_by_direction(
+        self,
+    ) -> Figure:
         """Display netflows grouped by direction."""
         if not self.check_valid_result_data("az_network_flows"):
             return None
@@ -381,8 +384,10 @@ def _get_azure_activity(self, src_ip, result, timespan):
 
         if self.check_table_exists("AzureActivity"):
             nb_data_wait("AzureActivity")
-            result.azure_activity = self.query_provider.Azure.list_azure_activity_for_ip(
-                timespan, ip_address_list=src_ip
+            result.azure_activity = (
+                self.query_provider.Azure.list_azure_activity_for_ip(
+                    timespan, ip_address_list=src_ip
+                )
             )
         _display_df_summary(result.azure_activity, "Azure Activity")
 
@@ -417,8 +422,10 @@ def _get_related_bookmarks(
         self, src_ip, result, timespan: TimeSpan
     ) -> pd.DataFrame:
         nb_data_wait("Bookmarks")
-        result.related_bookmarks = self.query_provider.AzureSentinel.list_bookmarks_for_entity(
-            timespan, entity_id=src_ip
+        result.related_bookmarks = (
+            self.query_provider.AzureSentinel.list_bookmarks_for_entity(  # type: ignore
+                timespan, entity_id=src_ip
+            )
         )
         _display_df_summary(result.related_bookmarks, "related bookmarks")
 

msticnb/nb/azsent/network/network_flow_summary.py

@@ -26,7 +26,6 @@
     nb_markdown,
     nb_warn,
     set_text,
-    df_has_data,
 )
 from ....data_providers import DataProviders
 from ....nblib.azsent.host import get_aznet_topology, get_heartbeat
@@ -697,7 +696,7 @@ def _display_geo_map_all(flow_index, ip_locator, host_entity):
     return folium_map
 
 
-# pylint: disable=too-many-branches
+# pylint: disable=too-many-branches, too-many-locals
 @set_text(docs=_CELL_DOCS, key="display_geo_map")
 def _display_geo_map(flow_index, ip_locator, host_entity, ti_results, select_asn):
     folium_map = foliummap.FoliumMap(zoom_start=4)

msticnb/nblib/azsent/host.py

@@ -83,7 +83,7 @@ def get_aznet_topology(
 
     """
     if "AzureNetworkAnalytics_CL" not in qry_prov.schema_tables:
-        return
+        return None
     nb_data_wait("AzureNetworkAnalytics")
     az_net_df = None
     if host_name:

msticnb/nblib/iptools.py

@@ -21,7 +21,7 @@
 
 
 def get_ip_ti(
-    ti_lookup: "TILookup",  # type: ignore
+    ti_lookup,
     data: pd.DataFrame,
     ip_col: str,
 ) -> pd.DataFrame:
@@ -62,9 +62,7 @@ def _normalize_ip4(data, ip_col):
     )
 
 
-def get_geoip_whois(
-    geo_lookup: "GeoIpLookup", data: pd.DataFrame, ip_col: str
-):  # type: ignore
+def get_geoip_whois(geo_lookup, data: pd.DataFrame, ip_col: str):
     """
     Get GeoIP and WhoIs data for IPs.
 
@@ -94,7 +92,10 @@ def get_geoip_whois(
     return get_whois_df(geo_df, "IpAddress", whois_col="Whois_data")
 
 
-_VPS_URL = "https://raw.githubusercontent.com/Azure/Azure-Sentinel/master/Sample%20Data/Feeds/VPS_Networks.csv"
+_VPS_URL = (
+    "https://raw.githubusercontent.com/Azure/Azure-Sentinel/"
+    + "master/Sample%20Data/Feeds/VPS_Networks.csv"
+)
 _NET_DICT = defaultdict(list)
 
 

msticnb/notebooklet.py

@@ -28,6 +28,7 @@
 __author__ = "Ian Hellen"
 
 
+# pylint: disable=too-many-public-methods
 class Notebooklet(ABC):
     """Base class for Notebooklets."""
 

msticnb/notebooklet_result.py

@@ -5,7 +5,7 @@
 # --------------------------------------------------------------------------
 """Notebooklet Result base classes."""
 import inspect
-from typing import Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import pandas as pd
 from bokeh.models import LayoutDOM
@@ -29,7 +29,7 @@ def __init__(
         self,
         description: Optional[str] = None,
         timespan: Optional[TimeSpan] = None,
-        notebooklet: Optional["Notebooklet"] = None,  # type: ignore
+        notebooklet: Optional[Any] = None,  # type: ignore
     ):
         """
         Create new Notebooklet result instance.
@@ -40,7 +40,7 @@ def __init__(
             Result description, by default None
         timespan : Optional[TimeSpan], optional
             TimeSpan for the results, by default None
-        notebooklet : Optional[, optional
+        notebooklet : Optional[Notebooklet], optional
             Originating notebooklet, by default None
         """
         self.description = description or self.__class__.__qualname__
@@ -109,6 +109,7 @@ def __getattr__(self, name):
         """Proxy attributes of the notebooklet member."""
         if self.notebooklet:
             return getattr(self.notebooklet, name)
+        raise AttributeError(f"{self.__class__} has no attribute '{name}'")
 
     def _populate_attr_desc(self):
         indent = " " * 4

tests/test_dataprovider.py

@@ -11,40 +11,44 @@
 from msticpy.sectools.geoip import GeoLiteLookup
 from msticpy.sectools import TILookup
 
-from msticnb.data_providers import DataProviders, init
+from msticnb import data_providers
 
+from .unit_test_lib import GeoIPLiteMock
 
 # pylint: disable=no-member
 
 
-def test_init_data_providers():
+def test_init_data_providers(monkeypatch):
     """Test creating DataProviders instance."""
-    dprov = DataProviders(query_provider="LocalData")
+    monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock)
+    dprov = data_providers.DataProviders(query_provider="LocalData")
 
     check.is_not_none(dprov)
-    check.equal(dprov, DataProviders.current())
+    check.equal(dprov, data_providers.DataProviders.current())
 
     check.is_in("LocalData", dprov.providers)
     check.is_in("geolitelookup", dprov.providers)
     check.is_in("tilookup", dprov.providers)
     check.is_instance(dprov.providers["LocalData"], QueryProvider)
-    check.is_instance(dprov.providers["geolitelookup"], GeoLiteLookup)
+    check.is_instance(dprov.providers["geolitelookup"], GeoIPLiteMock)
     check.is_instance(dprov.providers["tilookup"], TILookup)
 
 
-def test_new_init_data_providers():
+def test_new_init_data_providers(monkeypatch):
     """Test creating new provider with new provider list."""
-    init(query_provider="LocalData", providers=[])
-    dprov = DataProviders.current()
-    init(query_provider="LocalData", providers=[])
-    dprov2 = DataProviders.current()
+    monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock)
+
+    data_providers.init(query_provider="LocalData", providers=[])
+    dprov = data_providers.DataProviders.current()
+    data_providers.init(query_provider="LocalData", providers=[])
+    dprov2 = data_providers.DataProviders.current()
     check.equal(dprov2, dprov)
 
     # specify provider
-    dprov = DataProviders(query_provider="LocalData")
-    init(query_provider="LocalData", providers=["tilookup"])
+    dprov = data_providers.DataProviders(query_provider="LocalData")
+    data_providers.init(query_provider="LocalData", providers=["tilookup"])
     msticnb = sys.modules["msticnb"]
-    dprov2 = DataProviders.current()
+    dprov2 = data_providers.DataProviders.current()
     pkg_providers = getattr(msticnb, "data_providers")
     check.not_equal(dprov2, dprov)
     check.is_in("LocalData", dprov2.providers)
@@ -59,17 +63,21 @@ def test_new_init_data_providers():
     check.is_instance(dprov2.providers["tilookup"], TILookup)
 
 
-def test_add_sub_data_providers():
+def test_add_sub_data_providers(monkeypatch):
     """Test intializing adding and subtracting providers."""
-    dprov = DataProviders(query_provider="LocalData")
-    init(query_provider="LocalData", providers=["tilookup"])
+    monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock)
+
+    dprov = data_providers.DataProviders(query_provider="LocalData")
+    data_providers.init(query_provider="LocalData", providers=["tilookup"])
     msticnb = sys.modules["msticnb"]
-    dprov2 = DataProviders.current()
+    dprov2 = data_providers.DataProviders.current()
 
     # Add and remove a provider from defaults
-    init(query_provider="LocalData", providers=["+ipstacklookup", "-geolitelookup"])
+    data_providers.init(
+        query_provider="LocalData", providers=["+ipstacklookup", "-geolitelookup"]
+    )
 
-    dprov3 = DataProviders.current()
+    dprov3 = data_providers.DataProviders.current()
     pkg_providers = getattr(msticnb, "data_providers")
     check.not_equal(dprov3, dprov)
     check.not_equal(dprov3, dprov2)