diff --git a/api_client/python/timesketch_api_client/client.py b/api_client/python/timesketch_api_client/client.py index 5dbeef0194..c01789758f 100644 --- a/api_client/python/timesketch_api_client/client.py +++ b/api_client/python/timesketch_api_client/client.py @@ -460,6 +460,58 @@ def explore(self, response = self.api.session.post(resource_url, json=form_data) return response.json() + def label_events(self, events, label_name): + """Labels one or more events with label_name. + + Args: + events: Array of JSON objects representing events. + label_name: String to label the event with. + + Returns: + Dictionary with query results. + """ + form_data = { + 'annotation': label_name, + 'annotation_type': 'label', + 'events': events + } + resource_url = u'{0:s}/sketches/{1:d}/event/annotate/'.format( + self.api.api_root, self.id) + response = self.api.session.post(resource_url, json=form_data) + return response.json() + + def search_by_label(self, label_name): + """Searches for all events containing a given label. + + Args: + label_name: A string representing the label to search for. + + Returns: + A dictionary with query results. + """ + query = { + "nested": { + "path": "timesketch_label", + "query": { + "bool": { + "must": [ + { + "term": { + "timesketch_label.name": label_name + } + }, + { + "term": { + "timesketch_label.sketch_id": self.id + } + } + ] + } + } + } + } + return self.explore(query_dsl=json.dumps({'query': query})) + class SearchIndex(BaseResource): """Timesketch searchindex object. diff --git a/api_client/python/timesketch_api_client/client_test.py b/api_client/python/timesketch_api_client/client_test.py index 96091bdc37..4b5753ed33 100644 --- a/api_client/python/timesketch_api_client/client_test.py +++ b/api_client/python/timesketch_api_client/client_test.py @@ -52,7 +52,7 @@ def get(*args, **kwargs): @staticmethod def post(*args, **kwargs): """Mock POST request handler.""" - return + return mock_response(*args, **kwargs) return MockSession() @@ -88,10 +88,8 @@ def json(self): u'objects': [{ u'id': 1, - u'name': - u'test', - u'description': - u'test', + u'name': u'test', + u'description': u'test', u'timelines': [{ u'id': 1, u'name': u'test', @@ -130,6 +128,8 @@ def json(self): MockResponse(json_data=sketch_data), u'http://127.0.0.1/api/v1/sketches/1/timelines/1': MockResponse(json_data=timeline_data), + u'http://127.0.0.1/api/v1/sketches/1/explore/': + MockResponse(json_data=timeline_data), } return url_router.get(args[0], MockResponse(None, 404)) @@ -177,7 +177,6 @@ def setUp(self): self.sketch = self.api_client.get_sketch(1) # TODO: Add test for upload() - # TODO: Add test for explore() def test_get_views(self): """Test to get a view.""" @@ -193,6 +192,12 @@ def test_get_timelines(self): self.assertEqual(len(timelines), 2) self.assertIsInstance(timelines[0], client.Timeline) + def test_explore(self): + """Tests to explore a timeline.""" + results = self.sketch.explore(query_string="description:test") + self.assertEqual(len(results['objects']), 1) + self.assertIsInstance(results['objects'], list) + class ViewTest(unittest.TestCase): """Test View object."""