Added a decorator to quickly register a function as a feature. (#45)

* Added a decorator to quickly register a function as a feature.

* Moved feature additions to the lw object to a separate function.

* added marker to the pytest ini file.

* Added a bit more resilience into the decode accessor.

* Fixed error message for decoding base64.

Author: Kristinn

jupyter/docker/docker_build/10-functions.py

@@ -130,7 +130,15 @@ class LaceworkHelper:
 
     def __init__(self):
         self.ctx = LaceworkContext()
+        self.refresh_features()
 
+    def refresh_features(self):
+        """
+        Refresh all the features of the helper.
+
+        IF a new feature is added this function can be called to refresh
+        the features that are registered to the helper object.
+        """
         for feature, feature_name in manager.LaceworkManager.get_features():
             feature_fn = decorators.feature_decorator(feature, self.ctx)
             setattr(self, feature_name, feature_fn)

jupyter/laceworkjupyter/__init__.py

@@ -6,12 +6,18 @@
 """
 
 import base64
+import binascii
+import logging
+
 import urllib.parse
 
 import pandas as pd
 
 
-@pd.api.extensions.register_series_accessor('decode')
+logger = logging.getLogger("lacework_sdk.jupyter.accessor")
+
+
+@pd.api.extensions.register_series_accessor("decode")
 class DecodeAccessor:
     """
     Accessor class for decoding data.
@@ -26,10 +32,18 @@ def _decode_string_base64(self, string_value):
         :param str string_value: The base64 decoded string.
         :return: A decoded string.
         """
-        decoded_string = base64.b64decode(string_value)
         try:
-            return decoded_string.decode('utf8')
+            decoded_string = base64.b64decode(string_value)
+        except binascii.Error as exc:
+            logger.error(
+                "Unable to decode string: '{0:s}', error: {1!s}".format(
+                    string_value, exc)),
+            return string_value
+
+        try:
+            return decoded_string.decode("utf8")
         except UnicodeDecodeError:
+            logger.warning("Unable to decode string using UTF-8")
             return decoded_string
 
     def base64(self):

jupyter/laceworkjupyter/accessors.py

@@ -4,6 +4,7 @@
 from laceworkjupyter import manager
 
 
+@manager.register_feature
 def get_client(
         api_key=None, api_secret=None, account=None,
         subaccount=None, instance=None, base_domain=None,
@@ -21,6 +22,3 @@ def get_client(
 
     ctx.set_client(client)
     return client
-
-
-manager.LaceworkManager.add_feature(get_client, 'get_client')

jupyter/laceworkjupyter/features/client.py

@@ -5,6 +5,7 @@
 from laceworkjupyter import utils
 
 
+@manager.register_feature
 def parse_date_offset(offset_string, ctx=None):
     """
     Parse date offset string and return a start and end time.
@@ -18,6 +19,3 @@ def parse_date_offset(offset_string, ctx=None):
     ctx.add("start_time", start_time)
     ctx.add("end_time", end_time)
     return start_time, end_time
-
-
-manager.LaceworkManager.add_feature(parse_date_offset, "parse_date_offset")

jupyter/laceworkjupyter/features/date.py

@@ -247,6 +247,7 @@ def _execute_query(button):
         f"'<b><i>lw.ctx.get(\"lql_results\")</i></b>'")
 
 
+@manager.register_feature
 def cloud_hunt(ctx=None):
     """
     Displays a UI to build LQL queries to do threat hunting.
@@ -342,6 +343,3 @@ def cloud_hunt(ctx=None):
     lw_ctx = ctx
 
     display(grid)  # noqa: F821
-
-
-manager.LaceworkManager.add_feature(cloud_hunt, "cloud_hunt")

jupyter/laceworkjupyter/features/hunt.py

@@ -12,6 +12,7 @@
 logger = logging.getLogger("lacework_sdk.jupyter.feature.policies")
 
 
+@manager.register_feature
 def list_available_queries(ctx=None):
     """
     Returns a DataFrame with the available LQL queries.
@@ -21,6 +22,7 @@ def list_available_queries(ctx=None):
     return ctx.client.queries.get()
 
 
+@manager.register_feature
 def query_stored_lql(query_id, start_time="", end_time="", ctx=None):
     """
     Returns the results from running a LQL query.
@@ -58,8 +60,3 @@ def query_stored_lql(query_id, start_time="", end_time="", ctx=None):
 
     return client.queries.execute_by_id(
         query_id=query_id, arguments=arguments)
-
-
-manager.LaceworkManager.add_feature(
-    list_available_queries, "list_available_queries")
-manager.LaceworkManager.add_feature(query_stored_lql, "query_stored_lql")

jupyter/laceworkjupyter/features/policies.py

@@ -39,3 +39,19 @@ def get_features(cls):
         """
         for feature_fn, feature_name in cls._features.items():
             yield (feature_name, feature_fn)
+
+
+def register_feature(fn, name=""):
+    """
+    Decorator that can be used to register a feature.
+
+    :param function fn: The function to register.
+    :param str name: Optional string with the name of the function
+        as it should be registered. If not provided the name of the
+        function is used.
+    """
+    if not name:
+        name = fn.__name__
+
+    LaceworkManager.add_feature(fn, name)
+    return fn

jupyter/laceworkjupyter/manager.py

@@ -11,8 +11,13 @@ def test_decode_accessor():
     Tests the decode accessor.
     """
     lines = [
-        {'value': 12, 'some_string': 'VGhpcyBpcyBhIHN0cmluZw==', 'uri': 'http://mbl.is/%3Fstuff=r+1%20af'},
-        {'value': 114, 'some_string': 'VGhpcyBpcyBhIGEgc2VjcmV0', 'uri': 'http://mbl.is/%3Fsfi=r+1%20af'},
+        {
+            'value': 12, 'some_string': 'VGhpcyBpcyBhIHN0cmluZw==',
+            'uri': 'http://mbl.is/%3Fstuff=r+1%20af'
+        }, {
+            'value': 114, 'some_string': 'VGhpcyBpcyBhIGEgc2VjcmV0',
+            'uri': 'http://mbl.is/%3Fsfi=r+1%20af'
+        },
     ]
     frame = pd.DataFrame(lines)
 

jupyter/tests/laceworkjupyter/test_accessors.py

@@ -1,6 +1,7 @@
 [pytest]
 markers =
     ci_exempt: mark a test as exempt from CI.
+    flaky_test: mark a test as potentially flaky.
 
 addopts =
     --ignore=tests/api/v2/test_policies.py