Added ability to modify aggregations and aggregation groups (#1172)

* Adding the ability to save aggregations and groups in the API client.

* Adding setters into the agg group.

* Mistake

* Minor changes

* Filtering out aggregations that are part of a group in list_aggregations

* Fixing list_aggregations for the case where there are no groups

api_client/python/setup.py

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

api_client/python/timesketch_api_client/aggregation.py

@@ -184,6 +184,7 @@ def from_aggregator_run(
         """
         self.type = 'aggregator_run'
         self._parameters = aggregator_parameters
+
         self.resource_data = self._run_aggregator(
             aggregator_name, aggregator_parameters, view_id, chart_type)
 
@@ -212,14 +213,19 @@ def description(self):
         """Property that returns the description string."""
         return self._aggregator_data.get('description', '')
 
+    @description.setter
+    def description(self, description):
+        """Set the description of an aggregation."""
+        self._aggregator_data['description'] = description
+
     @property
     def id(self):
         """Property that returns the ID of the aggregator, if possible."""
         agg_id = self._aggregator_data.get('id')
         if agg_id:
             return agg_id
 
-        return -1
+        return 0
 
     @property
     def name(self):
@@ -229,6 +235,11 @@ def name(self):
             return name
         return self.aggregator_name
 
+    @name.setter
+    def name(self, name):
+        """Set the name of the aggregation."""
+        self._aggregator_data['name'] = name
+
     @property
     def dict(self):
         """Property that returns back a Dict with the results."""
@@ -277,6 +288,27 @@ def generate_chart(self):
         vega_spec_string = json.dumps(vega_spec)
         return altair.Chart.from_json(vega_spec_string)
 
+    def save(self):
+        """Save the aggregation in the database."""
+        data = {
+            'name': self.name,
+            'description': self.description,
+            'agg_type': self.aggregator_name,
+            'parameters': self._parameters,
+            'chart_type': self.chart_type,
+            'view_id': self.view,
+        }
+
+        if self.id:
+            resource_url = '{0:s}/sketches/{1:d}/aggregation/{2:d}/'.format(
+                self.api.api_root, self._sketch.id, self.id)
+        else:
+            resource_url = '{0:s}/sketches/{1:d}/aggregation/'.format(
+                self.api.api_root, self._sketch.id)
+
+        response = self.api.session.post(resource_url, json=data)
+        return response.status_code in definitions.HTTP_STATUS_CODE_20X
+
 
 class AggregationGroup(resource.BaseResource):
     """Aggregation Group object.
@@ -316,21 +348,45 @@ def description(self):
         """Returns the description of the aggregation group."""
         return self._description
 
+    @description.setter
+    def description(self, description):
+        """Sets the description of the aggregation group."""
+        self._description = description
+        self.save()
+
     @property
     def name(self):
         """Returns the name of the aggregation group."""
         return self._name
 
+    @name.setter
+    def name(self, name):
+        """Sets the name of the aggregation group."""
+        self._name = name
+        self.save()
+
     @property
     def orientation(self):
         """Returns the chart orientation."""
         return self._orientation
 
+    @orientation.setter
+    def orientation(self, orientation):
+        """Sets the chart orientation."""
+        self._orientation = orientation
+        self.save()
+
     @property
     def parameters(self):
         """Returns a dict with the group parameters."""
         return self._parameters
 
+    @parameters.setter
+    def parameters(self, parameters):
+        """Sets the group parameters."""
+        self._parameters = parameters
+        self.save()
+
     @property
     def table(self):
         """Property that returns a pandas DataFrame."""
@@ -435,6 +491,30 @@ def get_tables(self):
         """Returns a list of pandas DataFrame from each aggregation."""
         return [x.table for x in self._aggregations]
 
+    def save(self):
+        """Save the aggregation group in the database."""
+        if not self._aggregations:
+            return False
+
+        data = {
+            'name': self._name,
+            'description': self._description,
+            'parameters': json.dumps(self._parameters),
+            'aggregations': json.dumps([x.id for x in self._aggregations]),
+            'orientation': self._orientation,
+        }
+
+        if self.id:
+            res_url = '{0:s}/sketches/{1:d}/aggregation/group/{2:d}/'.format(
+                self.api.api_root, self._sketch.id, self.id)
+        else:
+            res_url = '{0:s}/sketches/{1:d}/aggregation/group/'.format(
+                self.api.api_root, self._sketch.id)
+
+        response = self.api.session.post(res_url, json=data)
+        _ = self.lazyload_data(refresh_cache=True)
+        return response.status_code in definitions.HTTP_STATUS_CODE_20X
+
     def to_pandas(self):
         """Returns a pandas DataFrame.
 

api_client/python/timesketch_api_client/sketch.py

@@ -248,10 +248,22 @@ def list_aggregations(self):
         if not isinstance(first_object, dict):
             return aggregations
 
+        aggregation_groups = first_object.get('aggregationgroups')
+        if aggregation_groups:
+            aggregation_groups = aggregation_groups[0]
+            groups = [
+                x.get('id', 0) for x in aggregation_groups.get(
+                    'aggregations', [])]
+        else:
+            groups = tuple()
+
         for aggregation_dict in first_object.get('aggregations', []):
+            agg_id = aggregation_dict.get('id')
+            if agg_id in groups:
+                continue
             aggregation_obj = aggregation.Aggregation(
                 sketch=self, api=self.api)
-            aggregation_obj.from_store(aggregation_id=aggregation_dict['id'])
+            aggregation_obj.from_store(aggregation_id=agg_id)
             aggregations.append(aggregation_obj)
         return aggregations
 

timesketch/api/v1/resources.py

@@ -1168,7 +1168,7 @@ def get(self, sketch_id, group_id):
 
         parameters = {}
         if group.parameters:
-            parameters = json.loads(parameters)
+            parameters = json.loads(group.parameters)
 
         result_chart.title = parameters.get('chart_title', group.name)
         time_after = time.time()
@@ -1185,6 +1185,71 @@ def get(self, sketch_id, group_id):
         schema = {'meta': meta, 'objects': objects}
         return jsonify(schema)
 
+    @login_required
+    def post(self, sketch_id, group_id):
+        """Handles POST request to the resource.
+
+        Args:
+            sketch_id: Integer primary key for a sketch database model.
+            group_id: Integer primary key for an aggregation group database
+                model.
+        """
+        sketch = Sketch.query.get_with_acl(sketch_id)
+        group = AggregationGroup.query.get(group_id)
+        if not group:
+            abort(
+                HTTP_STATUS_CODE_NOT_FOUND, 'No Group found with this ID.')
+
+        if not sketch:
+            abort(
+                HTTP_STATUS_CODE_NOT_FOUND, 'No sketch found with this ID.')
+
+        # Check that this group belongs to the sketch
+        if group.sketch_id != sketch.id:
+            msg = (
+                'The sketch ID ({0:d}) does not match with the aggregation '
+                'group sketch ID ({1:d})'.format(sketch.id, group.sketch_id))
+            abort(HTTP_STATUS_CODE_FORBIDDEN, msg)
+
+        if not sketch.has_permission(user=current_user, permission='write'):
+            abort(
+                HTTP_STATUS_CODE_FORBIDDEN,
+                'The user does not have write permission on the sketch.')
+
+
+        form = request.json
+        if not form:
+            abort(
+                HTTP_STATUS_CODE_BAD_REQUEST,
+                'No JSON data, unable to process request to create '
+                'a new aggregation group.')
+
+        group.name = form.get('name', group.name)
+        group.description = form.get('description', group.description)
+        group.parameters = form.get('parameters', group.parameters)
+        group.orientation = form.get('orientation', group.orientation)
+        group.user = current_user
+        group.sketch = sketch
+
+        agg_ids = json.loads(form.get('aggregations', group.aggregations))
+        aggregations = []
+
+        for agg_id in agg_ids:
+            aggregation = Aggregation.query.get(agg_id)
+            if not aggregation:
+                abort(
+                    HTTP_STATUS_CODE_BAD_REQUEST,
+                    'No aggregation found for ID: {0:d}'.format(agg_id))
+            aggregations.append(aggregation)
+
+        group.aggregations = aggregations
+
+        db_session.add(group)
+        db_session.commit()
+
+        return self.to_json(group, status_code=HTTP_STATUS_CODE_CREATED)
+
+
     @login_required
     def delete(self, sketch_id, group_id):
         """Handles DELETE request to the resource.