Added a helper to extract data from JSON structs inside data frames (#54)

Author: Kristinn

jupyter/docker/docker_build/00-import.py

@@ -1,4 +1,5 @@
 """This is an import file that runs on every startup of the Jupyter runtime."""
+# flake8: noqa
 
 import altair as alt
 import pandas as pd

jupyter/docker/docker_build/jupyter_notebook_config.py

@@ -1,4 +1,5 @@
 # Configuration file for jupyter-notebook.
+# flake8: noqa
 
 ## Use a regular expression for the Access-Control-Allow-Origin header
 #  

jupyter/laceworkjupyter/features/__init__.py

@@ -1,9 +1,11 @@
 """A simple loading of helper functions."""
+# flake8: noqa
 
-from . import cache  # noqa: F401
-from . import client  # noqa: F401
-from . import date  # noqa: F401
-from . import hunt  # noqa: F401
-from . import policies  # noqa: F401
-from . import query  # noqa: F401
-from . import query_builder  # noqa: F401
+from . import cache
+from . import client
+from . import date
+from . import helper
+from . import hunt
+from . import policies
+from . import query
+from . import query_builder

jupyter/laceworkjupyter/features/helper.py

@@ -0,0 +1,61 @@
+"""
+File that contains data frame helpers for the Lacework notebook environment.
+"""
+
+import json
+import logging
+
+import pandas as pd
+import numpy as np
+
+from laceworkjupyter import manager
+
+
+logger = logging.getLogger("lacework_sdk.jupyter.feature.helper")
+
+
+@manager.register_feature
+def deep_extract_field(data_frame, column, field_string, ctx=None):
+    """
+    Extract a field from a JSON struct inside a DataFrame.
+
+    Usage example:
+        df['hostname'] = lw.deep_extract_field(
+            df, 'properties', 'host.hostname')
+
+    :param DataFrame data_frame: The data frame to extract from.
+    :param str column: The name of the column that contains the JSON struct.
+    :param str field_string: String that contains the field to extract from,
+        this is a dot delimited string, eg: key.foo.bar, that will extract
+        a value from {'key': 'foo': {'bar': 'value'}}.
+    :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))

jupyter/laceworkjupyter/features/utils.py

@@ -141,7 +141,7 @@ def get_query_definition(table_name, filters):
     }
 
 
-def build_lql_query(query_name, query_dict, join_support=True):
+def build_lql_query(query_name, query_dict, join_support=True):  # noqa: C901
     """
     Build a LQL query and return evaluator ID and the query.