Added details to the current_user call in the API client and improvin…

…g e2e tests (#1541)

* Adding slightly more details to the username in the API client.

* adding one more check for status in e2e tests, to make less flaky

* Adding to e2e tests, running analyzers and a more specific search

* Adding a small API client test

* Moving to a more flexible user object.

api_client/python/timesketch_api_client/client.py

@@ -35,6 +35,7 @@
 from . import error
 from . import index
 from . import sketch
+from . import user
 from . import version
 from . import sigma
 
@@ -96,7 +97,6 @@ def __init__(self,
         self.api_root = '{0:s}/api/v1'.format(host_uri)
         self.credentials = None
         self._flow = None
-        self._username = username
 
         if not create_session:
             self.session = None
@@ -114,8 +114,8 @@ def __init__(self,
 
     @property
     def current_user(self):
-        """Property that returns the username that is logged in."""
-        return self._username
+        """Property that returns the user object of the logged in user."""
+        return user.User(self)
 
     @property
     def version(self):
@@ -507,7 +507,11 @@ def list_searchindices(self):
         """
         indices = []
         response = self.fetch_resource_data('searchindices/')
-        for index_dict in response['objects'][0]:
+        response_objects = response.get('objects')
+        if not response_objects:
+            return indices
+
+        for index_dict in response_objects[0]:
             index_id = index_dict['id']
             index_name = index_dict['name']
             index_obj = index.SearchIndex(

api_client/python/timesketch_api_client/search.py

@@ -606,7 +606,7 @@ def from_manual(  # pylint: disable=arguments-differ
         if not (query_string or query_filter or query_dsl):
             raise RuntimeError('You need to supply a query')
 
-        self._username = self.api.current_user
+        self._username = self.api.current_user.username
         self._name = 'From Explore'
         self._description = 'From Explore'
 
@@ -658,12 +658,19 @@ def from_saved(self, search_id):  # pylint: disable=arguments-differ
         self._query_dsl = data.get('query_dsl', '')
         query_filter = data.get('query_filter', '')
         if query_filter:
-            self.query_filter = json.loads(query_filter)
+            filter_dict = json.loads(query_filter)
+            if 'fields' in filter_dict:
+                fields = filter_dict.pop('fields')
+                return_fields = [x.get('field') for x in fields]
+                self.return_fields = ','.join(return_fields)
+
+            self.query_filter = filter_dict
         self._query_string = data.get('query_string', '')
         self._resource_id = search_id
         self._searchtemplate = data.get('searchtemplate', 0)
         self._updated_at = data.get('updated_at', '')
         self._username = data.get('user', {}).get('username', 'System')
+
         self.resource_data = data
 
     @property
@@ -822,11 +829,25 @@ def save(self):
             resource_url = (
                 f'{self.api.api_root}/sketches/{self._sketch.id}/views/')
 
+        query_filter = self.query_filter
+        if self.return_fields:
+            sketch_data = self._sketch.data
+            sketch_meta = sketch_data.get('meta', {})
+            mappings = sketch_meta.get('mappings', [])
+
+            use_mappings = []
+            for field in self.return_fields.split(','):
+                field = field.strip().lower()
+                for map_entry in mappings:
+                    if map_entry.get('field', '').lower() == field:
+                        use_mappings.append(map_entry)
+            query_filter['fields'] = use_mappings
+
         data = {
             'name': self.name,
             'description': self.description,
             'query': self.query_string,
-            'filter': self.query_filter,
+            'filter': query_filter,
             'dsl': self.query_dsl,
             'labels': json.dumps(self.labels),
         }

api_client/python/timesketch_api_client/user.py

@@ -0,0 +1,83 @@
+# Copyright 2020 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Timesketch API client library."""
+import logging
+
+from . import resource
+
+
+logger = logging.getLogger('timesketch_api.user')
+
+
+class User(resource.BaseResource):
+    """User object."""
+
+    def __init__(self, api):
+        """Initializes the user object."""
+        self._object_data = None
+        resource_uri = 'users/me/'
+        super().__init__(api, resource_uri)
+
+    def _get_data(self):
+        """Returns dict from the first object of the resource data."""
+        if self._object_data:
+            return self._object_data
+
+        data = self.data
+        objects = data.get('objects')
+        if objects:
+            self._object_data = objects[0]
+        else:
+            self._object_data = {}
+
+        return self._object_data
+
+    @property
+    def groups(self):
+        """Property that returns the groups the user belongs to."""
+        data = self._get_data()
+        groups = data.get('groups', [])
+        return [x.get('name', '') for x in groups]
+
+    @property
+    def is_active(self):
+        """Property that returns bool indicating whether the user is active."""
+        data = self._get_data()
+        return data.get('active', True)
+
+    @property
+    def is_admin(self):
+        """Property that returns bool indicating whether the user is admin."""
+        data = self._get_data()
+        return data.get('admin', False)
+
+    @property
+    def username(self):
+        """Property that returns back the username of the current user."""
+        data = self._get_data()
+        return data.get('username', 'Unknown')
+
+    def __str__(self):
+        """Returns a string representation of the username."""
+        user_strings = [self.username]
+
+        if self.is_active:
+            user_strings.append('[active]')
+        else:
+            user_strings.append('[inactive]')
+
+        if self.is_admin:
+            user_strings.append('<is admin>')
+
+        return ' '.join(user_strings)

api_client/python/timesketch_api_client/version.py

@@ -14,7 +14,7 @@
 """Version information for Timesketch API Client."""
 
 
-__version__ = '20201217'
+__version__ = '20201219'
 
 
 def get_version():

end_to_end_tests/__init__.py

@@ -14,5 +14,6 @@
 """End to end test module."""
 
 # Register all tests by importing them.
+from . import client_test
 from . import graph_test
 from . import query_test

end_to_end_tests/client_test.py

@@ -0,0 +1,53 @@
+# Copyright 2020 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""End to end tests of Timesketch client functionality."""
+
+from . import interface
+from . import manager
+
+
+class ClientTest(interface.BaseEndToEndTest):
+    """End to end tests for client functionality."""
+
+    NAME = 'client_test'
+
+    def test_client(self):
+        """Client tests."""
+        expected_user = 'test'
+        user = self.api.current_user
+        self.assertions.assertEqual(user.username, expected_user)
+        self.assertions.assertEqual(user.is_admin, False)
+        self.assertions.assertEqual(user.is_active, True)
+
+        sketch_name = 'Testing'
+        sketch_description = 'This is truly a foobar'
+        new_sketch = self.api.create_sketch(
+            name=sketch_name, description=sketch_description)
+
+        self.assertions.assertEqual(new_sketch.name, sketch_name)
+        self.assertions.assertEqual(
+            new_sketch.description, sketch_description)
+
+        first_sketch = self.api.get_sketch(1)
+        self.assertions.assertEqual(
+            self.sketch.name, first_sketch.name)
+
+        sketches = list(self.api.list_sketches())
+        self.assertions.assertEqual(len(sketches), 2)
+
+        for index in self.api.list_searchindices():
+            self.assertions.assertTrue(bool(index.name))
+
+
+manager.EndToEndTestManager.register_test(ClientTest)

end_to_end_tests/graph_test.py

@@ -52,5 +52,4 @@ def test_graph(self):
         self.assertions.assertEqual(graph_saved.description, 'this is it')
 
 
-
 manager.EndToEndTestManager.register_test(GraphTest)

end_to_end_tests/interface.py

@@ -40,6 +40,7 @@ class BaseEndToEndTest(object):
     """
 
     NAME = 'name'
+    _ANALYZERS_COMPLETE_SET = frozenset(['ERROR', 'DONE'])
 
     def __init__(self):
         """Initialize the end-to-end test object."""
@@ -80,8 +81,9 @@ def import_timeline(self, filename):
                 raise TimeoutError
             _ = timeline.lazyload_data(refresh_cache=True)
             status = timeline.status
+
             # TODO: Do something with other statuses? (e.g. failed)
-            if status == 'ready':
+            if status == 'ready' and timeline.index.status == 'ready':
                 break
             retry_count += 1
             time.sleep(sleep_time_seconds)
@@ -97,6 +99,46 @@ def _get_test_methods(self):
             if name.startswith('test_'):
                 yield name, func
 
+    def run_analyzer(self, timeline_name, analyzer_name):
+        """Run an analyzer on an imported timeline.
+
+        Args:
+            timeline_name (str): the name of the imported timeline.
+            analyzer_name (str): the name of the analyzer to run.
+        """
+        timeline = None
+        for time_obj in self.sketch.list_timelines():
+            if time_obj.name == timeline_name:
+                timeline = time_obj
+                break
+        if not timeline:
+            print(
+                f'Unable to run analyzer: {analyzer_name} on {timeline_name}, '
+                'didn\'t find the timeline, timeline name correct?')
+            return
+
+        results = timeline.run_analyzer(analyzer_name)
+
+        # Poll the analyzer status to see when analyzer completes it's run.
+        max_time_seconds = 600  # Timeout after 10 min
+        sleep_time_seconds = 5  # Sleep between API calls
+        max_retries = max_time_seconds / sleep_time_seconds
+        retry_count = 0
+
+        while True:
+            if retry_count >= max_retries:
+                raise TimeoutError('Unable to wait for analyzer run to end.')
+
+            status_set = set()
+            for line in results.status.split('\n'):
+                status_set.add(line.split()[-1])
+
+            if status_set.issubset(self._ANALYZERS_COMPLETE_SET):
+                break
+
+            retry_count += 1
+            time.sleep(sleep_time_seconds)
+
     def setup(self):
         """Setup function that is run before any tests.
 

end_to_end_tests/query_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """End to end tests of Timesketch query functionality."""
 
+import pandas as pd
 from timesketch_api_client import search
 
 from . import interface
@@ -36,5 +37,57 @@ def test_wildcard_query(self):
         count = len(data_frame)
         self.assertions.assertEqual(count, 3205)
 
+    def test_specific_queries(self):
+        """Test few specific queries."""
+        search_obj = search.Search(self.sketch)
+        search_obj.query_string = 'message_identifier: "1073748864"'
+        search_obj.return_fields = 'computer_name,data_type,strings,user_sid'
+
+        data_frame = search_obj.table
+        self.assertions.assertEqual(len(data_frame), 204)
+
+        computers = list(data_frame.computer_name.unique())
+        self.assertions.assertEqual(len(computers), 1)
+        self.assertions.assertEqual(
+            computers[0], 'WKS-WIN764BITB.shieldbase.local')
+
+        def extract_strings(row):
+            strings = row.strings
+            return pd.Series({
+                'service': strings[0],
+                'state_from': strings[1],
+                'state_to': strings[2],
+                'by': strings[3]
+            })
+
+        strings_frame = data_frame.apply(extract_strings, axis=1)
+        services = set(strings_frame.service.unique())
+        expected_set = set([
+            'Background Intelligent Transfer Service',
+            'Windows Modules Installer'])
+        self.assertions.assertSetEqual(services, expected_set)
+
+        search_name = 'My First Search'
+        search_obj.name = search_name
+        search_obj.description = 'Can it be, is it really?'
+
+        search_obj.save()
+
+        _ = self.sketch.lazyload_data(refresh_cache=True)
+        saved_search = None
+        for search_obj in self.sketch.list_saved_searches():
+            if search_obj.name == search_name:
+                saved_search = search_obj
+                break
+
+        if search_obj is None:
+            raise RuntimeError('Unable to find the saved search.')
+        self.assertions.assertEqual(
+            saved_search.return_fields,
+            'computer_name,data_type,strings,user_sid')
+
+        self.assertions.assertEqual(
+            saved_search.query_string, 'message_identifier: "1073748864"')
+
 
 manager.EndToEndTestManager.register_test(QueryTest)