Merge pull request #33 from lacework/datasources

Added the new DataSources API into the SDK and made changes to the Jupyter wrapper

jupyter/README.md

@@ -50,6 +50,26 @@ $ docker-compose pull
 $ docker-compose up -d
 ```
 
+The other option is to create a file called `docker-compose.yml` with the content of:
+
+```
+version: '3'
+services:
+  lacebook:
+    container_name: lacebook
+    image: docker.io/lacework/lacebook:latest
+    ports:
+      - 127.0.0.1:8899:8899
+    restart: on-failure
+    volumes:
+      - $HOME/.lacework.toml:/home/lacework/.lacework.toml
+      - /$HOME/data/:/usr/local/src/lacedata/
+```
+
+(*if you use this docker compose file you will either need to create the folder $HOME/data that is
+readable and writeable by a user with the UID/GID 1000:1000. or change the line to point to a folder
+with that permission*)
+
 This will start up a lacebook container which starts a Jupyter container listening on port 8899.
 To access the lacebook container visit http://localhost:8899. When prompted for a password
 use `lacework`.
@@ -75,3 +95,5 @@ and enter the backend URL: `http://localhost:8899/?token=lacework`
 ## How-To
 
 More to come here.
+
+One way to start exploring is to run this [Colab notebook](https://colab.research.google.com/github/lacework/python-sdk/blob/master/jupyter/notebooks/colab_sample.ipynb)

jupyter/laceworkjupyter/__init__.py

@@ -4,7 +4,8 @@
 from laceworksdk import LaceworkClient
 
 from . import config
-from . import utils
+from . import decorators
+from . import plugins
 
 
 logger = logging.getLogger('lacework_sdk.jupyter.client')
@@ -15,11 +16,22 @@ class APIWrapper:
     API Wrapper class that takes an API wrapper and decorates functions.
     """
 
-    def __init__(self, api_wrapper):
+    def __init__(self, api_wrapper, wrapper_name):
         self._api_wrapper = api_wrapper
+        self._api_name = wrapper_name
+
         for func_name in [f for f in dir(api_wrapper) if not f.startswith('_')]:
             func = getattr(api_wrapper, func_name)
-            setattr(self, func_name, utils.dataframe_decorator(func))
+
+            decorator_plugin = plugins.PLUGINS.get(
+                f"{self._api_name}.{func_name}")
+            if decorator_plugin:
+                setattr(
+                    self,
+                    func_name,
+                    decorators.plugin_decorator(func, decorator_plugin))
+            else:
+                setattr(self, func_name, decorators.dataframe_decorator(func))
 
 
 class LaceworkJupyterHelper:
@@ -45,7 +57,7 @@ def __init__(
         wrappers = [w for w in dir(self.sdk) if not w.startswith('_')]
         for wrapper in wrappers:
             wrapper_object = getattr(self.sdk, wrapper)
-            api_wrapper = APIWrapper(wrapper_object)
+            api_wrapper = APIWrapper(wrapper_object, wrapper_name=wrapper)
             setattr(self, wrapper, api_wrapper)
 
     def __enter__(self):

jupyter/laceworkjupyter/decorators.py

@@ -0,0 +1,37 @@
+import functools
+
+import pandas as pd
+
+from . import config
+
+
+def dataframe_decorator(function):
+    """
+    A decorator used to convert Lacework JSON API output into a dataframe.
+    """
+    @functools.wraps(function)
+    def get_output(*args, **kwargs):
+        data = function(*args, **kwargs)
+
+        if isinstance(data, dict):
+            df = pd.DataFrame(data.get('data', []))
+            if 'SEVERITY' in df:
+                df['SEVERITY'] = df.SEVERITY.apply(
+                    lambda x: config.SEVERITY_DICT.get(x, x))
+            return df
+
+        return data
+
+    return get_output
+
+
+def plugin_decorator(function, output_plugin):
+    """
+    A decorator used to use a plugin to convert Lacework JSON API output.
+    """
+    @functools.wraps(function)
+    def get_output(*args, **kwargs):
+        data = function(*args, **kwargs)
+        return output_plugin(data)
+
+    return get_output

jupyter/laceworkjupyter/plugins/__init__.py

@@ -0,0 +1,11 @@
+"""A simple loading of plugins."""
+
+from . import alert_rules
+from . import datasource
+
+
+PLUGINS = {
+    'alert_rules.get': alert_rules.process_alert_rules,
+    'datasource.list_data_sources': datasource.process_list_data_sources,
+    'datasource.get_datasource_schema': datasource.process_datasource_schema,
+}

jupyter/laceworkjupyter/plugins/alert_rules.py

@@ -0,0 +1,21 @@
+"""An output plugin for the Lacework Alert Rules API."""
+
+
+import pandas as pd
+
+
+def process_alert_rules(data):
+    """
+    Returns a Pandas DataFrame from the API call.
+
+    :return: A pandas DataFrame.
+    """
+    data_dicts = data.get("data", [])
+    lines = []
+    for data_dict in data_dicts:
+        filter_dict = data_dict.get("filters", {})
+        filter_dict["mcGuid"] = data_dict.get("mcGuid")
+        filter_dict["intgGuidList"] = data_dict.get("intgGuidList")
+        filter_dict["type"] = data_dict.get("type")
+        lines.append(filter_dict)
+    return pd.DataFrame(lines)

jupyter/laceworkjupyter/plugins/datasource.py

@@ -0,0 +1,26 @@
+"""An output plugin for the Lacework DataSource API."""
+
+
+import pandas as pd
+
+
+def process_list_data_sources(data):
+    """
+    Returns a Pandas DataFrame from the API call.
+
+    :return: A pandas DataFrame.
+    """
+    lines = [{'name': x, 'description': y} for x, y in data]
+    return pd.DataFrame(lines)
+
+
+def process_datasource_schema(data):
+    """
+    Returns a Pandas DataFrame from the output of the API call.
+
+    :return: A pandas DataFrame.
+    """
+    data_dict = data.get('data', {})
+    schemas = data_dict.get('resultSchema', [])
+
+    return pd.DataFrame(schemas)

laceworksdk/api/__init__.py

@@ -20,6 +20,7 @@
 from .container_registries import ContainerRegistriesAPI
 from .contract_info import ContractInfoAPI
 from .custom_compliance_config import CustomComplianceConfigAPI
+from .datasource import DataSourceAPI
 from .download_file import DownloadFileAPI
 from .events import EventsAPI
 from .integrations import IntegrationsAPI
@@ -132,6 +133,7 @@ def __init__(self,
         self.compliance.config = CustomComplianceConfigAPI(self._session)
         self.container_registries = ContainerRegistriesAPI(self._session)
         self.contract_info = ContractInfoAPI(self._session)
+        self.datasource = DataSourceAPI(self._session)
         self.events = EventsAPI(self._session)
         self.files = DownloadFileAPI(self._session)
         self.integrations = IntegrationsAPI(self._session)

laceworksdk/api/datasource.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+"""
+Lacework DataSource API wrapper.
+"""
+
+import logging
+
+import bleach
+
+logger = logging.getLogger(__name__)
+
+
+class DataSourceAPI:
+    """
+    Lacework DataSource API.
+    """
+
+    _DEFAULT_DESCRIPTION = "No description available."
+
+    def __init__(self, session):
+        """
+        Initializes the DataSource object.
+
+        :param session: An instance of the HttpSession class
+
+        :return DataSourceAPI object.
+        """
+
+        super(DataSourceAPI, self).__init__()
+
+        self._session = session
+
+    def get_datasource_schema(
+            self, data_source):
+        """
+        A method to get the schema for a particular data source.
+
+        :param data_source: A string representing the data source to check for.
+
+        :return response json
+        """
+
+        logger.info(
+            "Getting the schema for a particular datasource from Lacework...")
+
+        data_clean = bleach.clean(data_source)
+        api_uri = f"/api/v2/Datasources/{data_clean}"
+
+        response = self._session.get(api_uri)
+        return response.json()
+
+
+    def list_data_sources(self):
+        """
+        A method to list the data sources that are available.
+
+        :return A list of tuples with two entries, source name and description.
+        """
+        logger.info("Getting list of data sources Lacework...")
+
+        api_uri = "/api/v2/Datasources"
+        response = self._session.get(api_uri)
+
+        response_json = response.json()
+
+        return_sources = []
+        data_sources = response_json.get("data", [])
+        for data_source in data_sources:
+            description = data_source.get(
+                "description", self._DEFAULT_DESCRIPTION)
+            if description == 'None':
+                description = self._DEFAULT_DESCRIPTION
+
+            return_sources.append(
+                (data_source.get("name", "No name"), description))
+
+        return return_sources

requirements.txt

@@ -1,3 +1,4 @@
 python-dotenv~=0.15
 requests~=2.25
 configparser~=5.0
+bleach~=4.0

setup.py

@@ -27,6 +27,7 @@
     keywords=["lacework", "api", "sdk", "python", "api"],
     install_requires=[
         "python-dotenv",
+        "bleach",
         "requests",
         "configparser",
     ],