diff --git a/msticpy/__init__.py b/msticpy/__init__.py index fa549965a..b0c151e05 100644 --- a/msticpy/__init__.py +++ b/msticpy/__init__.py @@ -133,6 +133,7 @@ ) from .common.utility import search_name as search from .init.logging import set_logging_level, setup_logging +from .lazy_importer import lazy_import __version__ = VERSION __author__ = "Ian Hellen, Pete Bryan, Ashwin Patil" @@ -144,71 +145,25 @@ if not os.environ.get("KQLMAGIC_EXTRAS_REQUIRES"): os.environ["KQLMAGIC_EXTRAS_REQUIRES"] = "jupyter-basic" -_STATIC_ATTRIBS = list(locals().keys()) - -_DEFAULT_IMPORTS = { - "az_connect": "msticpy.auth.azure_auth", - "current_providers": "msticpy.init.nbinit", - "ContextLookup": "msticpy.context.contextlookup", - "GeoLiteLookup": "msticpy.context.geoip", - "init_notebook": "msticpy.init.nbinit", - "reset_ipython_exception_handler": "msticpy.init.nbinit", - "MicrosoftSentinel": "msticpy.context.azure", - "MpConfigEdit": "msticpy.config.mp_config_edit", - "MpConfigFile": "msticpy.config.mp_config_file", - "QueryProvider": "msticpy.data", - "TILookup": "msticpy.context.tilookup", - "TimeSpan": "msticpy.common.timespan", - "WorkspaceConfig": "msticpy.common.wsconfig", - "entities": "msticpy.datamodel", - "Pivot": "msticpy.init.pivot", +_LAZY_IMPORTS = { + "msticpy.auth.azure_auth.az_connect", + "msticpy.common.timespan.TimeSpan", + "msticpy.common.wsconfig.WorkspaceConfig", + "msticpy.config.mp_config_edit.MpConfigEdit", + "msticpy.config.mp_config_file.MpConfigFile", + "msticpy.context.azure.MicrosoftSentinel", + "msticpy.context.contextlookup.ContextLookup", + "msticpy.context.geoip.GeoLiteLookup", + "msticpy.context.tilookup.TILookup", + "msticpy.data.QueryProvider", + "msticpy.datamodel.entities", + "msticpy.init.nbinit.current_providers", + "msticpy.init.nbinit.init_notebook", + "msticpy.init.nbinit.reset_ipython_exception_handler", + "msticpy.init.pivot.Pivot", } - -def __getattr__(attrib: str) -> Any: - """ - Import and return an attribute of MSTICPy. - - Parameters - ---------- - attrib : str - The attribute name - - Returns - ------- - Any - The attribute value. - - Raises - ------ - AttributeError - No attribute found. - - """ - if attrib in _DEFAULT_IMPORTS: - try: - return getattr(importlib.import_module(_DEFAULT_IMPORTS[attrib]), attrib) - except (MsticpyImportExtraError, MsticpyImportExtraError): - raise - except (ImportError, MsticpyException) as err: - warnings.warn(f"Unable to import module for 'msticpy.{attrib}'") - print( - f"WARNING. The msticpy attribute '{attrib}' is not loadable.", - "You may need to install one or more additional dependencies.\n", - "Please check the exception details below for more information.", - "\n".join( - traceback.format_exception( - type(err), value=err, tb=err.__traceback__ - ) - ), - ) - raise AttributeError(f"msticpy failed to load '{attrib}'") from err - raise AttributeError(f"msticpy has no attribute '{attrib}'") - - -def __dir__(): - """Return attribute list.""" - return sorted(set(_STATIC_ATTRIBS + list(_DEFAULT_IMPORTS))) +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) def load_plugins(plugin_paths: Union[str, Iterable[str]]): diff --git a/msticpy/common/utility/package.py b/msticpy/common/utility/package.py index d317bb18c..b5b34773c 100644 --- a/msticpy/common/utility/package.py +++ b/msticpy/common/utility/package.py @@ -323,7 +323,7 @@ def init_dir(static_attribs: List[str], dynamic_imports: Dict[str, str]): return sorted(set(static_attribs + list(dynamic_imports))) -def lazy_import(module: str, attrib: str, call: bool = False): +def delayed_import(module: str, attrib: str, call: bool = False): """Import attribute from module on demand.""" attribute = None diff --git a/msticpy/config/__init__.py b/msticpy/config/__init__.py index 3b4344d3a..3390bbea6 100644 --- a/msticpy/config/__init__.py +++ b/msticpy/config/__init__.py @@ -12,21 +12,12 @@ It use the ipywidgets package. """ -from ..common.utility.package import init_dir, init_getattr +from ..lazy_importer import lazy_import -_STATIC_ATTRIBS = list(locals().keys()) -_DEFAULT_IMPORTS = { - "MpConfigControls": "msticpy.config.mp_config_control", - "MpConfigEdit": "msticpy.config.mp_config_edit", - "MpConfigFile": "msticpy.config.mp_config_file", +_LAZY_IMPORTS = { + "msticpy.config.mp_config_control.MpConfigControls", + "msticpy.config.mp_config_edit.MpConfigEdit", + "msticpy.config.mp_config_file.MpConfigFile", } - -def __getattr__(attrib: str): - """Import and a dynamic attribute of module.""" - return init_getattr(__name__, _DEFAULT_IMPORTS, attrib) - - -def __dir__(): - """Return attribute list.""" - return init_dir(_STATIC_ATTRIBS, _DEFAULT_IMPORTS) +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/config/ce_azure_sentinel.py b/msticpy/config/ce_azure_sentinel.py index 3d574c4de..313444778 100644 --- a/msticpy/config/ce_azure_sentinel.py +++ b/msticpy/config/ce_azure_sentinel.py @@ -9,7 +9,7 @@ import ipywidgets as widgets from .._version import VERSION -from ..common.utility.package import lazy_import +from ..common.utility.package import delayed_import # from ..context.azure.sentinel_core import MicrosoftSentinel from .ce_common import ( @@ -24,7 +24,7 @@ __version__ = VERSION __author__ = "Ian Hellen" -ms_sentinel = lazy_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel") +ms_sentinel = delayed_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel") # pylint: disable=too-many-ancestors diff --git a/msticpy/config/ce_user_defaults.py b/msticpy/config/ce_user_defaults.py index e1e0c7d65..dcfd457fd 100644 --- a/msticpy/config/ce_user_defaults.py +++ b/msticpy/config/ce_user_defaults.py @@ -67,7 +67,8 @@ def __init__(self, mp_controls: MpConfigControls): """ super().__init__(mp_controls) - from ..data import DataEnvironment # pylint: disable=import-outside-toplevel + # pylint: disable=import-outside-toplevel + from ..data.core.query_defns import DataEnvironment self._data_env_enum = DataEnvironment diff --git a/msticpy/config/mp_config_file.py b/msticpy/config/mp_config_file.py index 82ef91488..fdcac8c0b 100644 --- a/msticpy/config/mp_config_file.py +++ b/msticpy/config/mp_config_file.py @@ -25,16 +25,13 @@ except ImportError: _KEYVAULT = False -try: - from ..context.azure.sentinel_core import MicrosoftSentinel - - _SENTINEL = True -except ImportError: - _SENTINEL = False from ..common.pkg_config import current_config_path, refresh_config, validate_config +from ..common.utility.package import delayed_import from .comp_edit import CompEditDisplayMixin, CompEditStatusMixin from .file_browser import FileBrowser +ms_sentinel = delayed_import("msticpy.context.azure.sentinel_core", "MicrosoftSentinel") + __version__ = VERSION __author__ = "Ian Hellen" @@ -306,7 +303,7 @@ def get_workspace_from_url(url: str) -> Dict[str, Dict[str, str]]: workspace. """ - return MicrosoftSentinel.get_workspace_details_from_url(url) + return ms_sentinel().get_workspace_details_from_url(url) def _show_sentinel_workspace(self, show: bool = True): """Fetch settings from Sentinel Portal URL.""" diff --git a/msticpy/context/__init__.py b/msticpy/context/__init__.py index ad8992169..d3dbaf2a7 100644 --- a/msticpy/context/__init__.py +++ b/msticpy/context/__init__.py @@ -6,10 +6,8 @@ """Context Providers Subpackage.""" from typing import Any -# flake8: noqa: F403 from ..common.utility import ImportPlaceholder -from .geoip import GeoLiteLookup, IPStackLookup -from .tilookup import TILookup +from ..lazy_importer import lazy_import from .vtlookupv3 import VT3_AVAILABLE vtlookupv3: Any @@ -20,3 +18,12 @@ vtlookupv3 = ImportPlaceholder( # type: ignore "vtlookupv3", ["vt-py", "vt-graph-api", "nest_asyncio"] ) + + +_LAZY_IMPORTS = { + "msticpy.context.geoip.GeoLiteLookup", + "msticpy.context.geoip.IPStackLookup", + "msticpy.context.tilookup.TILookup", +} + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/context/azure/__init__.py b/msticpy/context/azure/__init__.py index a2168e161..741bf3b1d 100644 --- a/msticpy/context/azure/__init__.py +++ b/msticpy/context/azure/__init__.py @@ -5,6 +5,11 @@ # -------------------------------------------------------------------------- """Data provider sub-package.""" -# flake8: noqa: F401 -from .azure_data import AzureData -from .sentinel_core import MicrosoftSentinel +from ...lazy_importer import lazy_import + +_LAZY_IMPORTS = { + "msticpy.context.azure.azure_data.AzureData", + "msticpy.context.azure.sentinel_core.MicrosoftSentinel", +} + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/context/azure/sentinel_dynamic_summary.py b/msticpy/context/azure/sentinel_dynamic_summary.py index 0644789ca..2109fd00e 100644 --- a/msticpy/context/azure/sentinel_dynamic_summary.py +++ b/msticpy/context/azure/sentinel_dynamic_summary.py @@ -15,7 +15,7 @@ from ..._version import VERSION from ...common.exceptions import MsticpyAzureConnectionError, MsticpyParameterError from ...common.pkg_config import get_config, get_http_timeout -from ...data import QueryProvider +from ...data.core.data_providers import QueryProvider from .azure_data import get_api_headers # pylint: disable=unused-import diff --git a/msticpy/context/azure/sentinel_workspaces.py b/msticpy/context/azure/sentinel_workspaces.py index 27bf4560a..9dec77fd1 100644 --- a/msticpy/context/azure/sentinel_workspaces.py +++ b/msticpy/context/azure/sentinel_workspaces.py @@ -18,7 +18,7 @@ from ...common.data_utils import df_has_data from ...common.pkg_config import get_http_timeout from ...common.utility import mp_ua_header -from ...data import QueryProvider +from ...data.core.data_providers import QueryProvider __version__ = VERSION __author__ = "Ian Hellen" diff --git a/msticpy/context/tiproviders/kql_base.py b/msticpy/context/tiproviders/kql_base.py index 047803b1c..ba52302b3 100644 --- a/msticpy/context/tiproviders/kql_base.py +++ b/msticpy/context/tiproviders/kql_base.py @@ -26,7 +26,7 @@ from ...common.exceptions import MsticpyConfigError from ...common.utility import export from ...common.wsconfig import WorkspaceConfig -from ...data import QueryProvider +from ...data.core.data_providers import QueryProvider from ..lookup_result import LookupStatus from ..provider_base import generate_items from .ti_provider_base import ResultSeverity, TIProvider diff --git a/msticpy/context/tiproviders/open_page_rank.py b/msticpy/context/tiproviders/open_page_rank.py index d265ce665..930a201b5 100644 --- a/msticpy/context/tiproviders/open_page_rank.py +++ b/msticpy/context/tiproviders/open_page_rank.py @@ -31,7 +31,11 @@ @export class OPR(HttpTIProvider): - """Open PageRank Lookup.""" + """ + Open PageRank Lookup. + + See https://www.domcop.com/openpagerank/what-is-openpagerank + """ _BASE_URL = "https://openpagerank.com" @@ -50,10 +54,6 @@ def __init__(self, **kwargs): super().__init__(**kwargs) self._provider_name = self.__class__.__name__ - print( - "Using Open PageRank.", - "See https://www.domcop.com/openpagerank/what-is-openpagerank", - ) async def lookup_iocs_async( self, diff --git a/msticpy/data/__init__.py b/msticpy/data/__init__.py index ae18c06b1..44ef6502e 100644 --- a/msticpy/data/__init__.py +++ b/msticpy/data/__init__.py @@ -21,10 +21,16 @@ """ from .._version import VERSION -from ..common.exceptions import MsticpyImportExtraError -# flake8: noqa: F403 -from .core.data_providers import QueryProvider -from .core.query_defns import DataEnvironment, DataFamily +# from ..common.exceptions import MsticpyImportExtraError +from ..lazy_importer import lazy_import __version__ = VERSION + +_LAZY_IMPORTS = { + "msticpy.data.core.data_providers.QueryProvider", + "msticpy.data.core.query_defns.DataEnvironment", + "msticpy.data.core.query_defns.DataFamily", +} + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/data/azure_data.py b/msticpy/data/azure_data.py index b0886d5be..f84381ece 100644 --- a/msticpy/data/azure_data.py +++ b/msticpy/data/azure_data.py @@ -12,7 +12,7 @@ # flake8: noqa: F403, F401 # pylint: disable=unused-import -from ..context.azure import AzureData +from ..context.azure.azure_data import AzureData WARN_MSSG = ( "This module has moved to msticpy.context.azure.azure_data\n" diff --git a/msticpy/data/core/data_providers.py b/msticpy/data/core/data_providers.py index 7984eb71d..f9ec3455e 100644 --- a/msticpy/data/core/data_providers.py +++ b/msticpy/data/core/data_providers.py @@ -14,7 +14,7 @@ from ..._version import VERSION from ...common.pkg_config import get_config from ...common.utility import export, valid_pyname -from ...nbwidgets import QueryTime +from ...nbwidgets.query_time import QueryTime from .. import drivers from ..drivers.driver_base import DriverBase, DriverProps from .param_extractor import extract_query_params diff --git a/msticpy/data/core/query_provider_utils_mixin.py b/msticpy/data/core/query_provider_utils_mixin.py index 18b202cd4..ea1d2b4b2 100644 --- a/msticpy/data/core/query_provider_utils_mixin.py +++ b/msticpy/data/core/query_provider_utils_mixin.py @@ -9,7 +9,7 @@ from typing import Dict, Iterable, List, NamedTuple, Optional, Pattern, Protocol, Union from ..._version import VERSION -from ...common.utility.package import lazy_import +from ...common.utility.package import delayed_import from ..drivers.driver_base import DriverBase from .query_defns import DataEnvironment from .query_source import QuerySource @@ -18,7 +18,7 @@ __version__ = VERSION __author__ = "Ian Hellen" -query_browser = lazy_import("msticpy.vis.query_browser", "browse_queries") +query_browser = delayed_import("msticpy.vis.query_browser", "browse_queries") # pylint: disable=too-few-public-methods diff --git a/msticpy/init/nbinit.py b/msticpy/init/nbinit.py index 170d606ed..2710a3b66 100644 --- a/msticpy/init/nbinit.py +++ b/msticpy/init/nbinit.py @@ -183,7 +183,7 @@ def _verbose(verbosity: Optional[int] = None) -> int: dict(pkg="IPython.display", tgt="display"), dict(pkg="IPython.display", tgt="HTML"), dict(pkg="IPython.display", tgt="Markdown"), - dict(pkg="ipywidgets", alias="widgets"), + # dict(pkg="ipywidgets", alias="widgets"), dict(pkg="pathlib", tgt="Path"), dict(pkg="numpy", alias="np"), ] @@ -193,22 +193,22 @@ def _verbose(verbosity: Optional[int] = None) -> int: _MP_IMPORTS = [ dict(pkg="msticpy"), dict(pkg="msticpy.data", tgt="QueryProvider"), - dict(pkg="msticpy.vis.foliummap", tgt="FoliumMap"), - dict(pkg="msticpy.context", tgt="TILookup"), - dict(pkg="msticpy.context", tgt="GeoLiteLookup"), - dict(pkg="msticpy.context", tgt="IPStackLookup"), - dict(pkg="msticpy.transform", tgt="IoCExtract"), + # dict(pkg="msticpy.vis.foliummap", tgt="FoliumMap"), + # dict(pkg="msticpy.context", tgt="TILookup"), + # dict(pkg="msticpy.context", tgt="GeoLiteLookup"), + # dict(pkg="msticpy.context", tgt="IPStackLookup"), + # dict(pkg="msticpy.transform", tgt="IoCExtract"), dict(pkg="msticpy.common.utility", tgt="md"), dict(pkg="msticpy.common.utility", tgt="md_warn"), dict(pkg="msticpy.common.wsconfig", tgt="WorkspaceConfig"), dict(pkg="msticpy.init.pivot", tgt="Pivot"), dict(pkg="msticpy.datamodel", tgt="entities"), dict(pkg="msticpy.init", tgt="nbmagics"), - dict(pkg="msticpy.nbtools", tgt="SecurityAlert"), + # dict(pkg="msticpy.nbtools", tgt="SecurityAlert"), dict(pkg="msticpy.vis", tgt="mp_pandas_plot"), - dict(pkg="msticpy.vis", tgt="nbdisplay"), + # dict(pkg="msticpy.vis", tgt="nbdisplay"), dict(pkg="msticpy.init", tgt="mp_pandas_accessors"), - dict(pkg="msticpy", tgt="nbwidgets"), + # dict(pkg="msticpy", tgt="nbwidgets"), ] _MP_IMPORT_ALL: List[Dict[str, str]] = [ @@ -230,36 +230,6 @@ def _verbose(verbosity: Optional[int] = None) -> int: _SYNAPSE_KWARGS = ["identity_type", "storage_svc_name", "tenant_id", "cloud"] -def _pr_output(*args): - """Output to IPython display or print.""" - if not _VERBOSITY(): - return - if is_ipython(): - display(HTML(" ".join([*args, "
"]).replace("\n", "
"))) - else: - print(*args) - - -def _err_output(*args): - """Output to IPython display or print - always output regardless of verbosity.""" - if is_ipython(): - display(HTML(" ".join([*args, "
"]).replace("\n", "
"))) - display( - HTML( - "For more info and options run:" - "
import msticpy as mp\nhelp(mp.nbinit)
" - ) - ) - else: - print(*args) - print( - "\nFor more info and options run:", - "\n import msticpy as mp", - "\n help(mp.nbinit)", - ) - - -# pylint: disable=too-many-statements def init_notebook( namespace: Optional[Dict[str, Any]] = None, def_imports: str = "all", @@ -383,8 +353,6 @@ def init_notebook( https://github.com/Azure/Azure-Sentinel-Notebooks/blob/master/ConfiguringNotebookEnvironment.ipynb """ - global current_providers # pylint: disable=global-statement, invalid-name - if namespace is None and get_ipython(): namespace = get_ipython().user_global_ns else: @@ -416,12 +384,7 @@ def init_notebook( check_aml_settings(*_get_aml_globals(namespace)) else: # If not in AML check and print version status - stdout_cap = io.StringIO() - with redirect_stdout(stdout_cap): - check_version() - output = stdout_cap.getvalue() - _pr_output(output) - logger.info("Check version failures: %s", output) + _check_msticpy_version() if _detect_env("synapse", **kwargs) and is_in_synapse(): synapse_params = { @@ -431,14 +394,9 @@ def init_notebook( # Handle required packages and imports _pr_output("Processing imports....") - stdout_cap = io.StringIO() - with redirect_stdout(stdout_cap): - imp_ok = _global_imports( - namespace, additional_packages, user_install, extra_imports, def_imports - ) - output = stdout_cap.getvalue() - _pr_output(output) - logger.info("Import failures: %s", output) + imp_ok = _import_packages( + namespace, def_imports, additional_packages, extra_imports, user_install + ) # Configuration check if no_config_check: @@ -462,36 +420,106 @@ def init_notebook( ) # load pivots - stdout_cap = io.StringIO() - with redirect_stdout(stdout_cap): - _pr_output("Loading pivots.") - _load_pivots(namespace=namespace) - output = stdout_cap.getvalue() - _pr_output(output) - logger.info("Pivot load failures: %s", output) + _load_pivot_functions(namespace) # User defaults + _load_user_defaults(namespace) + + # show any warnings + _show_init_warnings(imp_ok, conf_ok) + _pr_output("

Notebook initialization complete

") + logger.info("Notebook initialization complete") + + +def _pr_output(*args): + """Output to IPython display or print.""" + if not _VERBOSITY(): + return + if is_ipython(): + display(HTML(" ".join([*args, "
"]).replace("\n", "
"))) + else: + print(*args) + + +def _err_output(*args): + """Output to IPython display or print - always output regardless of verbosity.""" + if is_ipython(): + display(HTML(" ".join([*args, "
"]).replace("\n", "
"))) + display( + HTML( + "For more info and options run:" + "
import msticpy as mp\nhelp(mp.nbinit)
" + ) + ) + else: + print(*args) + print( + "\nFor more info and options run:", + "\n import msticpy as mp", + "\n help(mp.nbinit)", + ) + + +def _load_user_defaults(namespace): + """Load user defaults, if defined.""" + global current_providers # pylint: disable=global-statement, invalid-name stdout_cap = io.StringIO() with redirect_stdout(stdout_cap): _pr_output("Loading user defaults.") prov_dict = load_user_defaults() output = stdout_cap.getvalue() - _pr_output(output) - logger.info(output) - logger.info("User default load failures: %s", output) + if output.strip(): + _pr_output(output) + logger.info(output) + logger.info("User default load failures: %s", output) if prov_dict: namespace.update(prov_dict) current_providers = prov_dict _pr_output("Auto-loaded components:", ", ".join(prov_dict.keys())) - # show any warnings - _show_init_warnings(imp_ok, conf_ok) - _pr_output("

Notebook initialization complete

") - logger.info("Notebook initialization complete") + +def _load_pivot_functions(namespace): + """Load pivot functions.""" + stdout_cap = io.StringIO() + with redirect_stdout(stdout_cap): + _pr_output("Loading pivots.") + _load_pivots(namespace=namespace) + output = stdout_cap.getvalue() + if output.strip(): + _pr_output(output) + logger.info("Pivot load failures: %s", output) + + +def _import_packages( + namespace, def_imports, additional_packages, extra_imports, user_install +): + """Import packages from default set or supplied as parameters.""" + stdout_cap = io.StringIO() + with redirect_stdout(stdout_cap): + imp_ok = _global_imports( + namespace, additional_packages, user_install, extra_imports, def_imports + ) + output = stdout_cap.getvalue() + if output.strip(): + _pr_output(output) + logger.info("Import failures: %s", output) + return imp_ok + + +def _check_msticpy_version(): + """Check msticpy version.""" + stdout_cap = io.StringIO() + with redirect_stdout(stdout_cap): + check_version() + output = stdout_cap.getvalue() + if output.strip(): + _pr_output(output) + logger.info("Check version failures: %s", output) def _show_init_warnings(imp_ok, conf_ok): + """Show any warnings from init_notebook.""" if imp_ok and conf_ok: return True md("

Notebook setup completed with some warnings.

") @@ -580,6 +608,7 @@ def _global_imports( extra_imports: List[str] = None, def_imports: str = "all", ): + """Import packages from default set (defined statically).""" import_list = [] imports, imports_all = _build_import_list(def_imports) diff --git a/msticpy/init/pivot.py b/msticpy/init/pivot.py index ac067a115..c5e537744 100644 --- a/msticpy/init/pivot.py +++ b/msticpy/init/pivot.py @@ -15,12 +15,10 @@ import pkg_resources -# pylint: disable=unused-import -from .. import datamodel # noqa:F401 from .._version import VERSION from ..common.timespan import TimeSpan from ..context.tilookup import TILookup -from ..data import QueryProvider +from ..data.core.data_providers import QueryProvider from ..datamodel import entities with warnings.catch_warnings(): @@ -28,7 +26,7 @@ from ..datamodel import pivot as legacy_pivot from ..common.utility.types import SingletonClass -from ..nbwidgets import QueryTime +from ..nbwidgets.query_time import QueryTime from . import pivot_init # pylint: disable=unused-import, no-name-in-module diff --git a/msticpy/init/user_config.py b/msticpy/init/user_config.py index 1fc6fad0f..aef49a996 100644 --- a/msticpy/init/user_config.py +++ b/msticpy/init/user_config.py @@ -238,7 +238,7 @@ def _load_azure_data(comp_settings=None, **kwargs): def _load_azsent_api(comp_settings=None, **kwargs): del kwargs - from ..context.azure import MicrosoftSentinel + from ..context.azure.sentinel_core import MicrosoftSentinel res_id = comp_settings.pop("res_id", None) if res_id: diff --git a/msticpy/lazy_importer.py b/msticpy/lazy_importer.py new file mode 100644 index 000000000..1c0452f03 --- /dev/null +++ b/msticpy/lazy_importer.py @@ -0,0 +1,87 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +"""Lazy importer for msticpy sub-packages.""" +import importlib +from types import ModuleType +from typing import Callable, Iterable, Tuple + +from ._version import VERSION + +__version__ = VERSION +__author__ = "Ian Hellen" + + +def lazy_import( + importer_name: str, import_list: Iterable[str] +) -> Tuple[ModuleType, Callable, Callable]: + """ + Return the importing module and a callable for lazy importing. + + Parameters + ---------- + importer_name: str + The module performing the import to help facilitate resolving + relative imports. + import_list: Iterable[str] + Iterable of the modules to be potentially imported (absolute + or relative). + The basic form is "[mod_path.]import_item". This will import "import_item" + from the module specified by "mod_path". + The `as` form of importing is also supported, + e.g. "pkg.mod.func as spam_func". + + Returns + ------- + Tuple[module, Callable] + This function returns a tuple of two items. The first is the importer + module for easy reference within itself. The second item is a callable to be + set to `__getattr__` of the calling module. + + Notes + ----- + Code modified (slightly) from example by Brett Cannon. + https://snarky.ca/lazy-importing-in-python-3-7/ + + """ + module = importlib.import_module(importer_name) + static_attribs = set(dir(module)) + import_mapping = {} + for name in import_list: + importing, _, binding = name.partition(" as ") + if not binding: + _, _, binding = importing.rpartition(".") + import_mapping[binding] = importing + + def __getattr__(name: str): + """Return the imported module or module member.""" + if name not in import_mapping: + message = f"module {importer_name!r} has no attribute {name!r}" + raise AttributeError(message) + importing = import_mapping[name] + mod_name, _, attrib_name = importing.rpartition(".") + if mod_name == importer_name: + # avoid infinite recursion + raise AttributeError( + f"Recursive import of name '[{mod_name}].{name}' from '{importer_name}'." + ) + # importlib.import_module() implicitly sets submodules on this module as + # appropriate for direct imports. + try: + imported = importlib.import_module( + mod_name, module.__spec__.parent # type: ignore + ) + except ImportError as imp_err: + message = f"cannot import name '{mod_name}' from '{importer_name}'" + raise ImportError(message) from imp_err + mod_attrib = getattr(imported, attrib_name, None) + setattr(module, name, mod_attrib) + return mod_attrib + + def __dir__(): + """Return module attribute list combining static and dynamic attribs.""" + return sorted(set(import_mapping).union(static_attribs)) + + return module, __getattr__, __dir__ diff --git a/msticpy/nbtools/__init__.py b/msticpy/nbtools/__init__.py index 77a692a26..327f240a9 100644 --- a/msticpy/nbtools/__init__.py +++ b/msticpy/nbtools/__init__.py @@ -18,53 +18,74 @@ """ # flake8: noqa: F403 # pylint: disable=W0401 -import importlib -from typing import Any +# import importlib +# from typing import Any -from .. import nbwidgets +# from .. import nbwidgets from .._version import VERSION -from ..common import utility as utils -from ..common.wsconfig import WorkspaceConfig -from ..vis import nbdisplay -from .security_alert import SecurityAlert +from ..lazy_importer import lazy_import -try: - from IPython import get_ipython +# from ..common import utility as utils +# from ..common.wsconfig import WorkspaceConfig +# from ..vis import nbdisplay +# from .security_alert import SecurityAlert - from ..init import nbmagics -except ImportError as err: - pass +# try: +# from IPython import get_ipython -# pylint: enable=W0401 +# from ..init import nbmagics +# except ImportError as err: +# pass +# pylint: enable=W0401 __version__ = VERSION - -_DEFAULT_IMPORTS = {"nbinit": "msticpy.init.nbinit"} - - -def __getattr__(attrib: str) -> Any: - """ - Import and return an attribute of nbtools. - - Parameters - ---------- - attrib : str - The attribute name - - Returns - ------- - Any - The attribute value. - - Raises - ------ - AttributeError - No attribute found. - - """ - if attrib in _DEFAULT_IMPORTS: - module = importlib.import_module(_DEFAULT_IMPORTS[attrib]) - return module - raise AttributeError(f"msticpy has no attribute {attrib}") +# _DEFAULT_IMPORTS = {"nbinit": "msticpy.init.nbinit"} + +_LAZY_IMPORTS = { + "msticpy.init.nbinit", + "msticpy.common.utility as utils", + "msticpy.common.wsconfig.WorkspaceConfig", + "msticpy.nbtools.security_alert.SecurityAlert", + "msticpy.nbwidgets", + "msticpy.vis.nbdisplay", +} + +# def __getattr__(attrib: str) -> Any: +# """ +# Import and return an attribute of nbtools. + +# Parameters +# ---------- +# attrib : str +# The attribute name + +# Returns +# ------- +# Any +# The attribute value. + +# Raises +# ------ +# AttributeError +# No attribute found. + +# """ +# if attrib in _DEFAULT_IMPORTS: +# module = importlib.import_module(_DEFAULT_IMPORTS[attrib]) +# return module +# raise AttributeError(f"msticpy has no attribute {attrib}") + +# from .vtlookupv3 import VT3_AVAILABLE + +# vtlookupv3: Any +# if VT3_AVAILABLE: +# from .vtlookupv3 import vtlookupv3 +# else: +# # vtlookup3 will not load if vt package not installed +# vtlookupv3 = ImportPlaceholder( # type: ignore +# "vtlookupv3", ["vt-py", "vt-graph-api", "nest_asyncio"] +# ) + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/nbwidgets/__init__.py b/msticpy/nbwidgets/__init__.py index 3e2e0e889..b074eb871 100644 --- a/msticpy/nbwidgets/__init__.py +++ b/msticpy/nbwidgets/__init__.py @@ -6,18 +6,23 @@ """Widgets sub-package.""" from .._version import VERSION - -# pylint: disable=unused-import -from .core import IPyDisplayMixin, RegisteredWidget # noqa: F401 -from .get_environment_key import GetEnvironmentKey # noqa: F401 -from .get_text import GetText # noqa: F401 -from .lookback import Lookback # noqa: F401 -from .option_buttons import OptionButtons # noqa: F401 -from .progress import Progress # noqa: F401 -from .query_time import QueryTime # noqa: F401 -from .select_alert import SelectAlert # noqa: F401 -from .select_item import SelectItem # noqa: F401 -from .select_subset import SelectSubset # noqa: F401 +from ..lazy_importer import lazy_import __version__ = VERSION __author__ = "Ian Hellen" + +_LAZY_IMPORTS = { + "msticpy.nbwidgets.core.IPyDisplayMixin", + "msticpy.nbwidgets.core.RegisteredWidget", + "msticpy.nbwidgets.get_environment_key.GetEnvironmentKey", + "msticpy.nbwidgets.get_text.GetText", + "msticpy.nbwidgets.lookback.Lookback", + "msticpy.nbwidgets.option_buttons.OptionButtons", + "msticpy.nbwidgets.progress.Progress", + "msticpy.nbwidgets.query_time.QueryTime", + "msticpy.nbwidgets.select_alert.SelectAlert", + "msticpy.nbwidgets.select_item.SelectItem", + "msticpy.nbwidgets.select_subset.SelectSubset", +} + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/sectools/__init__.py b/msticpy/sectools/__init__.py index c948c4f5f..1e80bfc9f 100644 --- a/msticpy/sectools/__init__.py +++ b/msticpy/sectools/__init__.py @@ -31,20 +31,18 @@ The sectools sub-package will be removed in version 2.0.0 """ -import contextlib - -# from . import process_tree_utils as ptree from .._version import VERSION -from ..context.geoip import GeoLiteLookup, IPStackLookup, geo_distance -from ..context.tilookup import TILookup -from ..transform import base64unpack as base64 +from ..lazy_importer import lazy_import -# flake8: noqa: F403 -# pylint: disable=W0401 -from ..transform.iocextract import IoCExtract +__version__ = VERSION -with contextlib.suppress(ImportError): - from IPython import get_ipython +_LAZY_IMPORTS = { + "msticpy.context.geoip.GeoLiteLookup", + "msticpy.context.geoip.IPStackLookup", + "msticpy.context.geoip.geo_distance", + "msticpy.context.tilookup.TILookup", + "msticpy.transform.base64unpack as base64", + "msticpy.transform.iocextract.IoCExtract", +} - from ..init import nbmagics as sectool_magics -__version__ = VERSION +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/sectools/vtlookupv3/__init__.py b/msticpy/sectools/vtlookupv3/__init__.py index 68492c297..52ecdcdce 100644 --- a/msticpy/sectools/vtlookupv3/__init__.py +++ b/msticpy/sectools/vtlookupv3/__init__.py @@ -4,20 +4,3 @@ # license information. # -------------------------------------------------------------------------- """VirusTotal V3 Subpackage.""" - -from ..._version import VERSION - -# pylint: disable=unused-import -# flake8: noqa: F401 -from .vtfile_behavior import VTFileBehavior -from .vtlookupv3 import ( - VT_API_NOT_FOUND, - MsticpyVTNoDataError, - VTEntityType, - VTLookupV3, - VTObjectProperties, -) -from .vtobject_browser import VTObjectBrowser - -__version__ = VERSION -__author__ = "Ian Hellen" diff --git a/msticpy/transform/__init__.py b/msticpy/transform/__init__.py index 07ad4b064..ba4661e19 100644 --- a/msticpy/transform/__init__.py +++ b/msticpy/transform/__init__.py @@ -6,9 +6,16 @@ """MSTICPy Data Processing Tools.""" from .._version import VERSION - -# flake8: noqa: F401 -from . import base64unpack, process_tree_utils -from .iocextract import IoCExtract +from ..lazy_importer import lazy_import __version__ = VERSION + +_LAZY_IMPORTS = { + "msticpy.context.geoip.GeoLiteLookup", + "msticpy.context.geoip.IPStackLookup", + "msticpy.context.geoip.geo_distance", + "msticpy.context.tilookup.TILookup", + "msticpy.transform.iocextract.IoCExtract", +} + +module, __getattr__, __dir__ = lazy_import(__name__, _LAZY_IMPORTS) diff --git a/msticpy/vis/query_browser.py b/msticpy/vis/query_browser.py index 598456fd0..9e0cccbac 100644 --- a/msticpy/vis/query_browser.py +++ b/msticpy/vis/query_browser.py @@ -10,7 +10,7 @@ from IPython.display import HTML from .._version import VERSION -from ..nbwidgets import SelectItem +from ..nbwidgets.select_item import SelectItem __version__ = VERSION __author__ = "Ian Hellen" diff --git a/msticpy/vis/ti_browser.py b/msticpy/vis/ti_browser.py index 680ce1021..a5d4806e8 100644 --- a/msticpy/vis/ti_browser.py +++ b/msticpy/vis/ti_browser.py @@ -11,7 +11,7 @@ from IPython.display import HTML from .._version import VERSION -from ..nbwidgets import SelectItem +from ..nbwidgets.select_item import SelectItem __version__ = VERSION __author__ = "Ian Hellen" diff --git a/msticpy/vis/vtobject_browser.py b/msticpy/vis/vtobject_browser.py index 9c891d5f5..6e1ff60b9 100644 --- a/msticpy/vis/vtobject_browser.py +++ b/msticpy/vis/vtobject_browser.py @@ -12,7 +12,7 @@ from .._version import VERSION from ..context.vtlookupv3.vtlookupv3 import VTLookupV3, timestamps_to_utcdate -from ..nbwidgets import IPyDisplayMixin +from ..nbwidgets.core import IPyDisplayMixin __version__ = VERSION __author__ = "Ian Hellen" @@ -30,6 +30,7 @@ ) +# pylint: disable=too-few-public-methods class VTObjectBrowser(IPyDisplayMixin): """VirusTotal object attributes browser.""" diff --git a/tests/init/test_nbinit.py b/tests/init/test_nbinit.py index 594d4a90c..76a3d0655 100644 --- a/tests/init/test_nbinit.py +++ b/tests/init/test_nbinit.py @@ -31,10 +31,8 @@ def test_nbinit_no_params(): verbose=True, ) - check.is_in("pd", ns_dict) check.is_in("get_ipython", ns_dict) check.is_in("Path", ns_dict) - check.is_in("np", ns_dict) print(ns_dict.keys()) # Note - msticpy imports throw when exec'd from unit test @@ -42,11 +40,6 @@ def test_nbinit_no_params(): check.is_in("WIDGET_DEFAULTS", ns_dict) - check.equal(ns_dict["pd"].__name__, "pandas") - check.equal(ns_dict["np"].__name__, "numpy") - - check.equal(pd.get_option("display.max_columns"), 50) - def test_nbinit_imports(): """Test custom imports.""" @@ -62,7 +55,6 @@ def test_nbinit_imports(): check.is_in("pathlib", ns_dict) check.is_in("time", ns_dict) check.is_in("tdelta", ns_dict) - check.is_in("np", ns_dict) check.equal(timedelta, ns_dict["tdelta"]) check.equal(datetime.time, ns_dict["time"]) diff --git a/tests/test_msticpy.py b/tests/test_msticpy.py index 3a6e89240..9e49da894 100644 --- a/tests/test_msticpy.py +++ b/tests/test_msticpy.py @@ -8,7 +8,7 @@ import pytest_check as check import msticpy -from msticpy import _DEFAULT_IMPORTS +from msticpy import _LAZY_IMPORTS __author__ = "Ian Hellen" @@ -16,9 +16,10 @@ def test_getattr(): """Test fetching and importing dynamic attributes.""" - for attrib in _DEFAULT_IMPORTS: - check.is_in(attrib, dir(msticpy)) + for lazy_import in _LAZY_IMPORTS: + _, _, lazy_attrib = lazy_import.rpartition(".") + check.is_in(lazy_attrib, dir(msticpy)) - obj = getattr(msticpy, attrib) + obj = getattr(msticpy, lazy_attrib) if isinstance(obj, type) or callable(obj): - check.equal(obj.__name__, attrib) + check.equal(obj.__name__, lazy_attrib)