Moving VT3 dependencies to extras. (#111)

* Moving VT3 dependencies to extras. (leaving Splunk for now since it's more tightly bound) * Fixing the text and URL in the exception
microsoft · Oct 15, 2020 · 1c8ecdd · 1c8ecdd
1 parent 524aaed
commit 1c8ecdd
Show file tree

Hide file tree

Showing 6 changed files with 83 additions and 65 deletions.
diff --git a/conda/conda-reqs.txt b/conda/conda-reqs.txt
@@ -15,7 +15,6 @@ keyring>=13.2.1
 matplotlib>=3.0.0
 msrest>=0.6.0
 msrestazure>=0.6.0
-nest_asyncio>=1.4.0
 networkx>=2.2
 numpy>=1.15.4
 pandas>=0.25.0
@@ -30,5 +29,3 @@ setuptools>=40.6.3
 statsmodels
 tqdm>=4.36.1
 urllib3>=1.23
-vt-py>=0.5.4
-vt-graph-api>=1.0.1
diff --git a/msticpy/common/exceptions.py b/msticpy/common/exceptions.py
@@ -143,7 +143,7 @@ def _repr_html_(self):
                 l_content, l_type = line
                 if l_type == "title":
                     content += f"<h3><p class='title'>{l_content}</p></h3>"
-                if l_type == "uri":
+                elif l_type == "uri":
                     if isinstance(l_content, tuple):
                         name, uri = l_content
                     else:
@@ -155,8 +155,7 @@ def _repr_html_(self):
                 text_line = line.replace("\n", "<br>")
                 content += f"{text_line}<br>"
 
-        ex_html = "".join((ex_style, div_tmplt.format(content=content)))
-        return ex_html
+        return "".join((ex_style, div_tmplt.format(content=content)))
 
     def _display_txt_exception(self):
         """Display text-only version of the exception text."""
@@ -167,7 +166,7 @@ def _display_txt_exception(self):
                     print("-" * len(l_content))
                     print(l_content)
                     print("-" * len(l_content))
-                if l_type == "uri":
+                elif l_type == "uri":
                     if isinstance(l_content, tuple):
                         print(f" - {': '.join(l_content)}")
                     else:
@@ -343,6 +342,44 @@ class MsticpyKqlConnectionError(MsticpyUserError):
     )
 
 
+class MsticpyImportExtraError(MsticpyUserError):
+    """Exception class for Imports that need an extra."""
+
+    DEF_HELP_URI = (
+        "Installing msticpy",
+        "https://msticpy.readthedocs.io/en/latest/getting_started/Installing.html",
+    )
+
+    def __init__(
+        self, *args, help_uri: Union[Tuple[str, str], str, None] = None, **kwargs
+    ):
+        """
+        Create import missing extra exception.
+
+        Parameters
+        ----------
+        help_uri : Union[Tuple[str, str], str, None], optional
+            Override the default help URI.
+        extra : str
+            The name of the setup extra that needs to be installed.
+
+        """
+        extra = kwargs.pop("extra", None)
+        if not extra:
+            raise AttributeError("Keyword argument 'extra' must be supplied")
+        mssg = "".join(
+            [
+                "This feature requires one or more additional packages",
+                "to be installed.\n",
+                "To do this run the command:\n",
+                f"pip install msticpy[{extra}]",
+            ]
+        )
+        add_args = [*args, mssg]
+        uri = help_uri or self.DEF_HELP_URI
+        super().__init__(*add_args, help_uri=uri, **kwargs)
+
+
 class MsticpyAzureConnectionError(MsticpyUserError):
     """Exception class for KqlConnection errors."""
 

diff --git a/msticpy/data/drivers/splunk_driver.py b/msticpy/data/drivers/splunk_driver.py
@@ -8,9 +8,6 @@
 from typing import Any, Tuple, Union, Dict, Iterable, Optional
 
 import pandas as pd
-import splunklib.client as sp_client
-import splunklib.results as sp_results
-from splunklib.client import AuthenticationError, HTTPError
 
 from .driver_base import DriverBase, QuerySource
 from ..._version import VERSION
@@ -19,9 +16,21 @@
     MsticpyConnectionError,
     MsticpyNotConnectedError,
     MsticpyUserConfigError,
+    MsticpyImportExtraError,
 )
 from ...common.provider_settings import get_provider_settings, ProviderSettings
 
+try:
+    import splunklib.client as sp_client
+    import splunklib.results as sp_results
+    from splunklib.client import AuthenticationError, HTTPError
+except ImportError as imp_err:
+    raise MsticpyImportExtraError(
+        "Cannot use this feature without splunk-sdk installed",
+        title="Error importing splunk-sdk",
+        extra="splunk",
+    ) from imp_err
+
 __version__ = VERSION
 __author__ = "Ashwin Patil"
 

diff --git a/msticpy/sectools/vtlookupv3.py b/msticpy/sectools/vtlookupv3.py
@@ -5,9 +5,18 @@
 import pandas as pd
 from IPython.display import HTML, display
 
-import vt
-from vt_graph_api import VTGraph
-from vt_graph_api import errors as vt_graph_errs
+from ..common.exceptions import MsticpyImportExtraError
+
+try:
+    import vt
+    from vt_graph_api import VTGraph
+    from vt_graph_api import errors as vt_graph_errs
+except ImportError as imp_err:
+    raise MsticpyImportExtraError(
+        "Cannot use this feature without vt-py and vt-graph-api packages installed.",
+        title="Error importing VirusTotal modules.",
+        extra="vt3",
+    ) from imp_err
 
 
 class MsticpyVTNoDataError(Exception):

diff --git a/requirements.txt b/requirements.txt
@@ -27,7 +27,6 @@ matplotlib>=3.0.0
 msal~=1.0.0
 msrest>=0.6.0
 msrestazure>=0.6.0
-nest_asyncio>=1.4.0
 networkx>=2.2
 numpy>=1.15.4
 pandas>=0.25.0
@@ -44,5 +43,3 @@ statsmodels>=0.11.1
 tldextract>=2.2.2
 tqdm>=4.36.1
 urllib3>=1.23
-vt-graph-api>=1.0.1
-vt-py>=0.5.4
diff --git a/setup.py b/setup.py
@@ -9,53 +9,6 @@
 import setuptools
 
 
-INSTALL_REQUIRES = [
-    "adal>=1.2.2",
-    "attrs>=18.2.0",
-    "azure-common>=1.1.18",
-    "azure-cli-core==2.5.0",
-    "azure-core>=1.2.2",
-    "azure-identity>=1.4.0",
-    "azure-keyvault-secrets>=4.0.0",
-    "azure-mgmt-compute>=4.6.2",
-    "azure-mgmt-keyvault>=2.0.0",
-    "azure-mgmt-monitor>=1.0.1",
-    "azure-mgmt-network>=2.7.0",
-    "azure-mgmt-resource>=2.2.0,<=10.1.0",
-    "azure-mgmt-subscription>=0.2.0",
-    "beautifulsoup4>=4.6.3",
-    "bokeh>=1.4.0",
-    "cryptography>=2.8",
-    "deprecated>=1.2.4",
-    "dnspython>=1.16.0",
-    "folium>=0.9.0",
-    "geoip2>=2.9.0",
-    "ipwhois>=1.1.0",
-    "ipython>=7.1.1",
-    "ipywidgets>=7.4.2",
-    "keyring>=13.2.1",
-    "Kqlmagic>=0.1.106",
-    "matplotlib>=3.0.0",
-    "msal~=1.0.0",
-    "msrest>=0.6.0",
-    "msrestazure>=0.6.0",
-    "networkx>=2.2",
-    "numpy>=1.15.4",
-    "pandas>=0.25.0",
-    "pytz>=2019.2",
-    "pyyaml>=3.13",
-    "requests>=2.21.1",
-    "scikit-learn>=0.20.2",
-    "scipy>=1.1.0",
-    "seaborn>=0.9.0",
-    "setuptools>=40.6.3",
-    "statsmodels>=0.11.1",
-    "tldextract>=2.2.2",
-    "tqdm>=4.36.1",
-    "urllib3>=1.23",
-]
-
-
 with open("README.md", "r") as fh:
     LONG_DESC = fh.read()
 
@@ -71,6 +24,13 @@
 with open("requirements-dev.txt", "r") as fh:
     INSTALL_DEV_REQUIRES = fh.readlines()
 
+# Extras definitions
+EXTRAS = {
+    "dev": INSTALL_DEV_REQUIRES,
+    "vt3": ["vt-py>=0.5.4", "vt-graph-api>=1.0.1", "nest_asyncio>=1.4.0"],
+    "splunk": ["splunk-sdk>=1.6.0"],
+}
+
 setuptools.setup(
     name="msticpy",
     version=__version__,
@@ -94,8 +54,17 @@
         "Development Status :: 4 - Beta",
     ],
     install_requires=INSTALL_REQUIRES,
-    extras_require={"dev": INSTALL_DEV_REQUIRES},
-    keywords=["security", "azure", "sentinel"],
+    extras_require=EXTRAS,
+    keywords=[
+        "security",
+        "azure",
+        "sentinel",
+        "mstic",
+        "cybersec",
+        "infosec",
+        "cyber",
+        "cybersecurity",
+    ],
     zip_safe=False,
     include_package_data=True,
 )