Adding labels to aggregations, adding ability to filter based on labe…

…ls. (#1176)

* Adding labels to aggregations.

* Getting labels completed.

* Cleanup

* Fixes

* Fixes after form deletion and test tool fixes.

* Minor changes

* Making labels JSON instead of CSV

api_client/python/setup.py

@@ -24,7 +24,7 @@
 
 setup(
     name='timesketch-api-client',
-    version='20200420',
+    version='20200424',
     description='Timesketch API client',
     license='Apache License, Version 2.0',
     url='http://www.timesketch.org/',

api_client/python/timesketch_api_client/aggregation.py

@@ -40,6 +40,7 @@ class Aggregation(resource.BaseResource):
     def __init__(self, sketch, api):
         self._sketch = sketch
         self._aggregator_data = {}
+        self._labels = []
         self._parameters = {}
         self.aggregator_name = ''
         self.chart_color = ''
@@ -135,6 +136,12 @@ def from_store(self, aggregation_id):
         self.aggregator_name = data.get('agg_type')
         self.type = 'stored'
 
+        label_string = data.get('label_string', '')
+        if label_string:
+            self._labels = json.loads(label_string)
+        else:
+            self._labels = []
+
         chart_type = data.get('chart_type')
         param_string = data.get('parameters', '')
         if param_string:
@@ -227,6 +234,11 @@ def id(self):
 
         return 0
 
+    @property
+    def labels(self):
+        """Property that returns a list of the aggregation labels."""
+        return self._labels
+
     @property
     def name(self):
         """Property that returns the name of the aggregation."""
@@ -250,6 +262,17 @@ def table(self):
         """Property that returns a pandas DataFrame."""
         return self.to_pandas()
 
+    def add_label(self, label):
+        """Add a label to the aggregation.
+
+        Args:
+            label (str): string with the label information.
+        """
+        if label in self._labels:
+            return
+        self._labels.append(label)
+        self.save()
+
     def to_dict(self):
         """Returns a dict."""
         entries = {}
@@ -296,8 +319,11 @@ def save(self):
             'agg_type': self.aggregator_name,
             'parameters': self._parameters,
             'chart_type': self.chart_type,
-            'view_id': self.view,
         }
+        if self.view:
+            data['view_id'] = self.view
+        if self._labels:
+            data['labels'] = json.dumps(self._labels)
 
         if self.id:
             resource_url = '{0:s}/sketches/{1:d}/aggregation/{2:d}/'.format(

api_client/python/timesketch_api_client/sketch.py

@@ -228,9 +228,17 @@ def list_aggregation_groups(self):
             groups.append(group)
         return groups
 
-    def list_aggregations(self):
+    def list_aggregations(self, include_labels=None, exclude_labels=None):
         """List all saved aggregations for this sketch.
 
+        Args:
+            include_labels (list): list of strings with labels. If defined
+                then only return aggregations that have the label in the list.
+            exclude_labels (list): list of strings with labels. If defined
+                then only return aggregations that don't have a label in the
+                list. include_labels will be processed first in case both are
+                defined.
+
         Returns:
             List of aggregations (instances of Aggregation objects)
         """
@@ -261,6 +269,20 @@ def list_aggregations(self):
             agg_id = aggregation_dict.get('id')
             if agg_id in groups:
                 continue
+            label_string = aggregation_dict.get('label_string', '')
+            if label_string:
+                labels = json.loads(label_string)
+            else:
+                labels = []
+
+            if include_labels:
+                if not any(x in include_labels for x in labels):
+                    continue
+
+            if exclude_labels:
+                if any(x in exclude_labels for x in labels):
+                    continue
+
             aggregation_obj = aggregation.Aggregation(
                 sketch=self, api=self.api)
             aggregation_obj.from_store(aggregation_id=agg_id)
@@ -815,32 +837,22 @@ def store_aggregation(
         Returns:
           A stored aggregation object or None if not stored.
         """
-        resource_url = '{0:s}/sketches/{1:d}/aggregation/'.format(
-            self.api.api_root, self.id)
-
-        form_data = {
-            'name': name,
-            'description': description,
-            'agg_type': aggregator_name,
-            'chart_type': chart_type,
-            'sketch': self.id,
-            'parameters': aggregator_parameters
-        }
-
-        response = self.api.session.post(resource_url, json=form_data)
-        if response.status_code not in definitions.HTTP_STATUS_CODE_20X:
-            error.error_message(
-                response, message='Error storing the aggregation',
-                error=RuntimeError)
-
-        response_dict = response.json()
+        # TODO: Deprecate this function.
+        logger.warning(
+            'This function is about to be deprecated, please use the '
+            '`.save()` function of an aggregation object instead')
+
+        aggregator_obj = self.run_aggregator(
+            aggregator_name, aggregator_parameters)
+        aggregator_obj.name = name
+        aggregator_obj.description = description
+        if chart_type:
+            aggregator_obj.chart_type = chart_type
+        if aggregator_obj.save():
+            _ = self.lazyload_data(refresh_cache=True)
+            return aggregator_obj
 
-        objects = response_dict.get('objects', [])
-        if not objects:
-            return None
-
-        _ = self.lazyload_data(refresh_cache=True)
-        return self.get_aggregation(objects[0].get('id'))
+        return None
 
     def comment_event(self, event_id, index, comment_text):
         """

api_client/python/timesketch_api_client/story.py

@@ -685,7 +685,7 @@ def to_export_format(self, export_format):
         logger.error(
             'Error exporting story: [{0:d}] {1!s} {2!s}'.format(
                 response.status_code, response.reason, response.text))
-        return ''
+        return {}
 
     def to_string(self):
         """Returns a string with the content of all the story."""

notebooks/jupyter-timesketch-demo.ipynb

@@ -728,7 +728,7 @@
     "\n",
     "Start by checking out whether there are any stored aggregations that we can just take a look at.\n",
     "\n",
-    "You can also store your own aggregations using the `gd_sketch.store_aggregation` function. However we are not going to do that in this colab."
+    "You can also store your own aggregations using the `.save()` function on the aggregation object. However we are not going to do that in this colab."
    ]
   },
   {