Added a simple Mitre Att&ck class as well as a helper for Alerts API. (…

…#57)

jupyter/docker/docker_build/Dockerfile

@@ -32,7 +32,7 @@ COPY --chown=1000:1000 docker/docker_build/lacework /home/lacework/lacenv/share/
 RUN pip install --upgrade pip setuptools wheel && \
     pip install --upgrade ipywidgets jupyter_contrib_nbextensions jupyter_http_over_ws ipydatetime tabulate && \
     pip install --upgrade scikit-learn matplotlib python-evtx Evtx timesketch_import_client "snowflake-connector-python[secure-local-storage,pandas]" && \
-    pip install --upgrade ipyaggrid keras nbformat numpy pandas pyparsing qgrid ruamel.yaml sklearn && \
+    pip install --upgrade ipyaggrid keras nbformat numpy pandas pyparsing qgrid ruamel.yaml sklearn mitreattack-python && \
     pip install --upgrade tensorflow tqdm traitlets xmltodict ds4n6-lib picatrix timesketch_api_client openpyxl && \
     cd /home/lacework/code && pip install -e . && \
     cd /home/lacework/code/jupyter && pip install -e . && \

jupyter/laceworkjupyter/features/__init__.py

@@ -6,6 +6,7 @@
 from . import date
 from . import helper
 from . import hunt
+from . import mitre
 from . import policies
 from . import query
 from . import query_builder

jupyter/laceworkjupyter/features/helper.py

@@ -14,6 +14,36 @@
 logger = logging.getLogger("lacework_sdk.jupyter.feature.helper")
 
 
+def extract_json_field(json_obj, item):
+    """
+    Extract a field from a JSON struct.
+
+    :param json_obj: Either a JSON string or a dict.
+    :param str item: The item string, dot delimited.
+    :return: The extracted field.
+    """
+    if isinstance(json_obj, str):
+        try:
+            json_obj = json.loads(json_obj)
+        except json.JSONDecodeError:
+            logger.error("Unable to decode JSON string: %s", json_obj)
+            return np.nan
+
+    if not isinstance(json_obj, dict):
+        logger.error("Unable to extract, not a dict: %s", type(json_obj))
+        return np.nan
+
+    data = json_obj
+    for point in item.split("."):
+        if not isinstance(data, dict):
+            logger.error(
+                "Sub-item %s is not a dict (%s)", point, type(data))
+            return np.nan
+
+        data = data.get(point)
+    return data
+
+
 @manager.register_feature
 def deep_extract_field(data_frame, column, field_string, ctx=None):
     """
@@ -31,31 +61,9 @@ def deep_extract_field(data_frame, column, field_string, ctx=None):
     :param obj ctx: The context object.
     :return: A pandas Series with the extracted value.
     """
-    def _extract_function(json_obj, item):
-        if isinstance(json_obj, str):
-            try:
-                json_obj = json.loads(json_obj)
-            except json.JSONDecodeError:
-                logger.error("Unable to decode JSON string: %s", json_obj)
-                return np.nan
-
-        if not isinstance(json_obj, dict):
-            logger.error("Unable to extract, not a dict: %s", type(json_obj))
-            return np.nan
-
-        data = json_obj
-        for point in item.split("."):
-            if not isinstance(data, dict):
-                logger.error(
-                    "Sub-item %s is not a dict (%s)", point, type(data))
-                return np.nan
-
-            data = data.get(point)
-        return data
-
     if column not in data_frame:
         logger.error("Column does not exist in the dataframe.")
         return pd.Series()
 
     return data_frame[column].apply(
-        lambda x: _extract_function(x, field_string))
+        lambda x: extract_json_field(x, field_string))

jupyter/laceworkjupyter/features/mitre.py

@@ -0,0 +1,183 @@
+"""
+File that provides access to Mitre ATT&CK data.
+"""
+
+import pandas as pd
+
+from taxii2client.v20 import Server
+from stix2 import CompositeDataSource
+
+from laceworkjupyter import manager
+from laceworkjupyter import utils as main_utils
+from laceworkjupyter.features import client
+from laceworkjupyter.features import utils
+
+
+class MitreAttack:
+    """
+    Simple class that controls access to Mitre Att&ck datasets.
+    """
+
+    DEFAULT_SERVER = "https://cti-taxii.mitre.org/taxii/"
+
+    def __init__(self):
+        self._api_root = None
+        self._server = None
+        self._source = None
+        self._frames = {}
+
+    def _get_frame_from_collection(self, collection):
+        """
+        Returns a DataFrame with information from a Mitre Att&ck collection.
+        """
+        data = collection.get_objects()
+        lines = []
+        for item in data.get('objects', []):
+            line = item
+            external = item.get('external_references', [])
+            for source in external:
+                if not source.get("url", "").startswith("https://attack.mitre.org"):
+                    continue
+                line["mitre_source"] = source.get("source_name")
+                line["mitre_id"] = source.get("external_id")
+                url = source.get("url")
+                url_items = url.split("/")
+                line["mitre_collection"] = url_items[-2]
+            lines.append(line)
+        return pd.DataFrame(lines)
+
+    @property
+    def server(self):
+        """
+        Returns a copy of the Mitre server object.
+        """
+        if self._server:
+            return self._server
+
+        self._server = Server(self.DEFAULT_SERVER)
+        self._api_root = self._server.api_roots[0]
+        return self._server
+
+    @property
+    def collection_titles(self):
+        """
+        Returns a list of collection titles and ID that are available.
+        """
+        _ = self.server
+        return [(c.title, c.id) for c in self._api_root.collections]
+
+    @property
+    def source(self):
+        """
+        Returns the data source object for Att&ck data.
+        """
+        if self._source:
+            return self._source
+
+        self._source = CompositeDataSource()
+        self._source.add_data_sources(self._api_root.collections)
+
+        return self._source
+
+    def get_collection(self, collection_title):
+        """
+        Returns a DataFrame with the content of a particular collection.
+
+        :param str collection_title: The tile or a substring of a collection
+            entry in Mitre Att&ck framework. The first match of available
+            collections will be used.
+        :return: A pandas DataFrame with the content or an empty frame if the
+            collection was not found.
+        """
+        _ = self.server
+        for collection in self._api_root.collections:
+            if collection_title.lower() not in collection.title.lower():
+                continue
+
+            title = collection.title
+            if title not in self._frames:
+                self._frames[title] = self._get_frame_from_collection(
+                    collection)
+            return self._frames[title]
+
+        return pd.DataFrame()
+
+    def get_technique(self, mitre_id, collection_title='Enterprise ATT&CK'):
+        """
+        Returns a pandas Series with the content of a single technique.
+
+        :param str mitre_id: This is the corresponding Mitre Att&ck ID for the
+            technique to look up.
+        :param str collection_title: The title of the collection this
+            particular technique exists in. By default uses the Enterprise
+            collection.
+        :return: A pandas Series with the information, or an empty series
+            if not found.
+        """
+        titles = [x[0] for x in self.collection_titles]
+        use_title = ""
+        for title in titles:
+            if collection_title.lower() in title.lower():
+                use_title = title
+                break
+        if not use_title:
+            raise ValueError("Collection Title %s not found", collection_title)
+
+        if use_title not in self._frames:
+            _ = self.get_collection(use_title)
+        frame = self._frames.get(use_title)
+
+        frame_slice = frame[frame.mitre_id == mitre_id.upper()]
+        if not frame_slice.empty:
+            return frame_slice.iloc[0].dropna()
+
+        return pd.Series()
+
+
+@manager.register_feature
+def get_mitre_client(ctx=None):
+    """
+    Returns a Mitre class object.
+    """
+    if ctx:
+        mitre_client = ctx.get("mitre_client")
+        if mitre_client:
+            return mitre_client
+
+        mitre_client = MitreAttack()
+        ctx.add("mitre_client", mitre_client)
+        return mitre_client
+
+    return MitreAttack()
+
+
+@manager.register_feature
+def get_alerts_data_with_mitre(start_time="", end_time="", ctx=None):
+    """
+    Returns a DataFrame from the Alerts API call with Att&ck information.
+
+    :param str start_time: The start time, in ISO format.
+    :param str end_time: The end time, in ISO format.
+    :return: A pandas DataFrame with the content of the Alerts API call
+        merged with Att&CK information, if applicable.
+    """
+    if not ctx:
+        raise ValueError("The context is required for this operation.")
+
+    attack_mappings = pd.DataFrame(utils.load_yaml_file("mitre.yaml"))
+    mitre_client = get_mitre_client(ctx=ctx)
+    mitre_enterprise = mitre_client.get_collection("enterprise")
+    mitre_joined_df = attack_mappings.merge(mitre_enterprise, how="left")
+
+    lw_client = ctx.client
+    if not lw_client:
+        lw_client = client.get_client()
+
+    default_start, default_end = main_utils.parse_date_offset("LAST 2 DAYS")
+    if not start_time:
+        start_time = ctx.get("start_time", default_start)
+    if not end_time:
+        end_time = ctx.get("end_time", default_end)
+
+    alert_df = lw_client.alerts.get(start_time=start_time, end_time=end_time)
+    return alert_df.merge(mitre_joined_df, how="left")

jupyter/laceworkjupyter/features/mitre.yaml

@@ -0,0 +1,5 @@
+- mitre_id: T1204
+  alertType: SuspiciousApplicationLaunched
+
+- mitre_id: T1531
+  alertType: NetworkSecurityGroupCreatedOrUpdated

jupyter/laceworkjupyter/features/query_builder.py

@@ -171,7 +171,7 @@ def generate_filters(change):  # noqa: C901
     table_name = datasource['name']
     lw_ctx.add_state("query_builder", "query_custom_table_name", table_name)
 
-    table_schema = lw_ctx.client.datasources.get_datasource_schema(table_name)
+    table_schema = lw_ctx.client.datasources.get_datasource(table_name)
 
     checkboxes = []
     return_fields = []

jupyter/laceworkjupyter/plugins/__init__.py

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

jupyter/laceworkjupyter/plugins/alerts.py

@@ -0,0 +1,31 @@
+"""An output plugin for the Lacework Alerts API."""
+
+
+import pandas as pd
+
+from laceworkjupyter.features import helper
+
+
+def process_alerts(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:
+        data_dict["alertDescription"] = helper.extract_json_field(
+            data_dict.get("alertProps", {}), "description.descriptionId")
+
+        description_dict = helper.extract_json_field(
+            data_dict.get("alertProps", {}), "description.descriptionObj")
+        data_dict.update(description_dict)
+
+        alert_context = helper.extract_json_field(
+            data_dict.get("keys", {}), "src.keys.alert")
+        if alert_context:
+            data_dict.update(alert_context)
+
+        lines.append(data_dict)
+    return pd.DataFrame(lines)

jupyter/laceworkjupyter/version.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.2'
+__version__ = '0.1.5'
 
 
 def get_version():

jupyter/requirements.txt

@@ -2,3 +2,4 @@ laceworksdk>=0.9.23
 pandas>=1.0.1
 pyyaml>=5.4.1
 ipywidgets>=7.6.5
+mitreattack-python>=1.4.2

jupyter/setup.py

@@ -35,6 +35,8 @@
         "pyyaml",
         "laceworksdk",
         "ipywidgets",
+        "mitreattack-python",
+        "taxii2client",
         "pandas"
     ],
     classifiers=[