text: # Copyright 2015 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.
  """This module contains common test utilities for Timesketch."""
  
  from __future__ import unicode_literals
  
  import codecs
  import json
  import six
  
  from flask_testing import TestCase
  
  from timesketch.app import create_app
  from timesketch.lib.definitions import HTTP_STATUS_CODE_REDIRECT
  from timesketch.models import init_db
  from timesketch.models import drop_all
  from timesketch.models import db_session
  from timesketch.models.user import Group
  from timesketch.models.user import User
  from timesketch.models.sketch import Sketch
  from timesketch.models.sketch import Timeline
  from timesketch.models.sketch import SearchIndex
  from timesketch.models.sketch import SearchTemplate
  from timesketch.models.sketch import View
  from timesketch.models.sketch import Event
  from timesketch.models.sketch import Story
  
  
  class TestConfig(object):
      """Config for the test environment."""
      DEBUG = True
      SECRET_KEY = 'testing'
      SQLALCHEMY_DATABASE_URI = 'sqlite://'
      WTF_CSRF_ENABLED = False
      ELASTIC_HOST = None
      ELASTIC_PORT = None
      LABELS_TO_PREVENT_DELETION = ['protected', 'magic']
      UPLOAD_ENABLED = False
      GRAPH_BACKEND_ENABLED = False
      AUTO_INDEX_ANALYZERS = []
      AUTO_SKETCH_ANALYZERS = []
      SIMILARITY_DATA_TYPES = []
  
  
  class MockElasticClient(object):
      """A mock implementation of a ElasticSearch client."""
  
      def __init__(self):
          """Initialize the client."""
          self.indices = MockElasticIndices()
  
      def search(self, index, body, size):  # pylint: disable=unused-argument
          """Mock a client search, used for aggregations."""
          meta = {
              'es_time': 23,
              'es_total_count': 5621,
              'timed_out': False,
              'max_score': 0.0
          }
  
          objects = [{
              'my_aggregation': {'buckets': [
                  {'foobar': 1, 'second': 'foobar'},
                  {'foobar': 4, 'second': 'more stuff'},
                  {'foobar': 532, 'second': 'hvernig hefurdu thad'}]},
              'my_second_aggregation': {'buckets': [
                  {'foobar': 54, 'second': 'faranlegt', 'third': 'other text'},
                  {'foobar': 42, 'second': 'asnalegt'}]}
  
          }]
          return {'meta': meta, 'objects': objects}
  
  
  class MockElasticIndices(object):
      # pylint: disable=unused-argument
      def get_mapping(self, *args, **kwargs):
          """Mock get mapping call."""
          return {}
  
      def stats(self, *args, **kwargs):
          return {'indices': {}}
  
  
  class MockDataStore(object):
      """A mock implementation of a Datastore."""
  
      event_dict = {
          '_index': [],
          '_id': 'adc123',
          '_type': 'plaso_event',
          '_source': {
              'es_index': '',
              'es_id': '',
              'label': '',
              'timestamp': 1410895419859714,
              'timestamp_desc': '',
              'datetime': '2014-09-16T19:23:40+00:00',
              'source_short': '',
              'source_long': '',
              'message': '',
          }
      }
      search_result_dict = {
          'hits': {
              'hits': [{
                  'sort': [1410593223000],
                  '_type': 'plaso_event',
                  '_source': {
                      'timestamp':
                      1410593222543942,
                      'message':
                      'Test event',
                      'timesketch_label': [
                          {
                              'user_id': 1,
                              'name': '__ts_star',
                              'sketch_id': 1
                          },
                          {
                              'user_id': 2,
                              'name': '__ts_star',
                              'sketch_id': 99
                          },
                      ],
                      'timestamp_desc':
                      'Content Modification Time',
                      'datetime':
                      '2014-09-13T07:27:03+00:00'
                  },
                  '_score': 'null',
                  '_index': 'test',
                  '_id': 'test'
              }],
              'total':
              1,
              'max_score':
              'null'
          },
          '_shards': {
              'successful': 10,
              'failed': 0,
              'total': 10
          },
          'took': 5,
          'timed_out': False
      }
  
      def __init__(self, host, port):
          """Initialize the datastore.
          Args:
              host: Hostname or IP address to the datastore
              port: The port used by the datastore
          """
          self.client = MockElasticClient()
          self.host = host
          self.port = port
          # Dictionary containing event dictionaries.
          self.event_store = {}
  
      # pylint: disable=arguments-differ,unused-argument
      def search(self, *args, **kwargs):
          """Mock a search query.
          Returns:
              A dictionary with search result or integer if count is requested.
          """
          if kwargs.get('count'):
              # 4711 is sometimes used instead of 17, on occasions when you want
              # to denote a slightly larger number. Probably comes from the name
              # of 'genuine' Eau-de-cologne, 'No. 4711 ".
              # Ref: https://hack.org/mc/writings/hackerswe/hackerswe.html
              return 4711
          return self.search_result_dict
  
      def get_event(self, searchindex_id, event_id):
          """Mock returning a single event from the datastore.
  
          Args:
              searchindex_id: String of ElasticSearch index id
              event_id: String of ElasticSearch event id
  
          Returns:
              A dictionary with event data.
          """
          return self.event_dict
  
      @staticmethod
      def get_filter_labels(sketch_id, indices):
          """Mock returning a single event from the datastore.
  
          Returns:
              A list with label.
          """
          return []
  
      def set_label(self,
                    searchindex_id,
                    event_id,
                    event_type,
                    sketch_id,
                    user_id,
                    label,
                    toggle=False,
                    single_update=True):
          """Mock adding a label to an event."""
          return
  
      # pylint: disable=unused-argument
      def create_index(self, *args, **kwargs):
          """Mock creating an index."""
          return
  
      def import_event(self, index_name, event_type, event=None,
                       event_id=None, flush_interval=None):
          """Mock adding the event to Elasticsearch, instead add the event
          to event_store.
          Args:
              flush_interval: Number of events to queue up before indexing. (This
              functionality is not supported.)
              index_name: Name of the index in Elasticsearch
              event_type: Type of event (e.g. plaso_event)
              event: Event dictionary
              event_id: Event Elasticsearch ID
          """
  
          if event_id in self.event_store:
              self.event_store[event_id]['_source'].update(event)
              return
  
          new_event = {
              '_index': index_name,
              '_id': event_id,
              '_type': event_type,
              '_source': event
          }
          self.event_store[event_id] = new_event
  
      @property
      def version(self):
          """Get Elasticsearch version.
  
          Returns:
            Version number as a string.
          """
          return '6.0'
  
      # pylint: disable=unused-argument
      def search_stream(self, query_string, query_filter, query_dsl,
                        indices, return_fields, enable_scroll=True):
          for i in range(len(self.event_store)):
              yield self.event_store[str(i)]
  
  
  class MockGraphDatabase(object):
      """A mock implementation of a Datastore."""
  
      def __init__(self, host, username, password):
          """Initialize the datastore.
          Args:
              host: Neo4j host
              username: Neo4j username
              password: Neo4j password
          """
          self.host = host
          self.username = username
          self.password = password
  
      class MockQuerySequence(object):
          """A mock implementation of a QuerySequence."""
          MOCK_GRAPH = [{
              'nodes': [{
                  'id': '1',
                  'labels': ['User'],
                  'properties': {
                      'username': 'test',
                      'uid': '123456'
                  }
              }, {
                  'id': '2',
                  'labels': ['Machine'],
                  'properties': {
                      'hostname': 'test'
                  }
              }],
              'relationships': [{
                  'endNode': '2',
                  'id': '3',
                  'startNode': '1',
                  'properties': {
                      'method': 'Network'
                  },
                  'type': 'ACCESS'
              }]
          }]
          MOCK_ROWS = {}
          MOCK_STATS = {}
  
          def __init__(self):
              self.graph = self.MOCK_GRAPH
              self.rows = self.MOCK_ROWS
              self.stats = self.MOCK_ROWS
  
      class MockEmptyQuerySequence(object):
          def __init__(self):
              self.graph = None
              self.rows = {}
              self.stats = {}
  
      # pylint: disable=unused-argument
      def query(self, *args, **kwargs):
          """Mock a search query.
          Returns:
              A MockQuerySequence instance.
          """
          query = args[0]
          if query == 'empty':
              return self.MockEmptyQuerySequence()
          return self.MockQuerySequence()
  
  
  class BaseTest(TestCase):
      """Base class for tests."""
  
      COLOR_WHITE = 'FFFFFF'
  
      def create_app(self):
          """Setup the Flask application.
          Returns:
              Flask application (instance of flask.app.Flask)
          """
          app = create_app(TestConfig)
          return app
  
      def _commit_to_database(self, model):
          """Add object to the database session and commit.
          Args:
              model: Instance of timesketch.models.[model] object
          """
          db_session.add(model)
          db_session.commit()
  
      def _create_user(self, username, set_password=False):
          """Create a user in the database.
          Args:
              username: Username (string)
              set_password: Boolean value to decide if a password should be set
          Returns:
              A user (instance of timesketch.models.user.User)
          """
          user = User.get_or_create(username=username)
          if set_password:
              user.set_password(plaintext='test', rounds=4)
          self._commit_to_database(user)
          return user
  
      def _create_group(self, name, user):
          """Create a user in the database.
          Args:
              name: Group name
              user: A user (instance of timesketch.models.user.User)
          Returns:
              A group (instance of timesketch.models.user.Group)
          """
          group = Group.get_or_create(name=name)
          user.groups.append(group)
          self._commit_to_database(group)
          return group
  
      def _create_sketch(self, name, user, acl=False):
          """Create a sketch in the database.
          Args:
              name: Name of the sketch (string)
              user: A user (instance of timesketch.models.user.User)
              acl: Boolean value to decide if ACL permissions should be set
          Returns:
              A sketch (instance of timesketch.models.sketch.Sketch)
          """
          sketch = Sketch.get_or_create(name=name, description=name, user=user)
          if acl:
              for permission in ['read', 'write', 'delete']:
                  sketch.grant_permission(permission=permission, user=user)
          label = sketch.Label(label='Test label', user=user)
          status = sketch.Status(status='Test status', user=user)
          sketch.labels.append(label)
          sketch.status.append(status)
          self._commit_to_database(sketch)
          return sketch
  
      def _create_searchindex(self, name, user, acl=False):
          """Create a searchindex in the database.
          Args:
              name: Name of the searchindex (string)
              user: A user (instance of timesketch.models.user.User)
              acl: Boolean value to decide if ACL permissions should be set
          Returns:
              A searchindex (instance of timesketch.models.sketch.SearchIndex)
          """
          searchindex = SearchIndex.get_or_create(
              name=name, description=name, index_name=name, user=user)
          if acl:
              for permission in ['read', 'write', 'delete']:
                  searchindex.grant_permission(permission=permission, user=user)
          searchindex.set_status(status='ready')
          self._commit_to_database(searchindex)
          return searchindex
  
      def _create_event(self, sketch, searchindex, user):
          """Create an event in the database.
          Args:
              sketch: A sketch (instance of timesketch.models.sketch.Sketch)
              searchindex:
                  A searchindex (instance of timesketch.models.sketch.SearchIndex)
              user: A user (instance of timesketch.models.user.User)
          Returns:
              An event (instance of timesketch.models.sketch.Event)
          """
          event = Event.get_or_create(
              sketch=sketch, searchindex=searchindex, document_id='test')
          comment = event.Comment(comment='test', user=user)
          event.comments.append(comment)
          self._commit_to_database(event)
          return event
  
      def _create_story(self, sketch, user):
          """Create a story in the database.
          Args:
              sketch: A sketch (instance of timesketch.models.sketch.Sketch)
              user: A user (instance of timesketch.models.user.User)
          Returns:
              A story (instance of timesketch.models.story.Story)
          """
          story = Story.get_or_create(
              title='Test', content='Test', sketch=sketch, user=user)
          self._commit_to_database(story)
          return story
  
      def _create_timeline(self, name, sketch, searchindex, user):
          """Create a timeline in the database.
          Args:
              name: Name of the timeline (string)
              sketch: A sketch (instance of timesketch.models.sketch.Sketch)
              searchindex:
                  A searchindex (instance of timesketch.models.sketch.SearchIndex)
              user: A user (instance of timesketch.models.user.User)
          Returns:
              A timeline (instance of timesketch.models.sketch.Timeline)
          """
          timeline = Timeline(
              name=name,
              description=name,
              user=user,
              sketch=sketch,
              searchindex=searchindex,
              color=self.COLOR_WHITE)
          timeline.set_status(status='ready')
          self._commit_to_database(timeline)
          return timeline
  
      def _create_view(self, name, sketch, user):
          """Create a view in the database.
          Args:
              name: Name of the view (string)
              sketch: A sketch (instance of timesketch.models.sketch.Sketch)
              user: A user (instance of timesketch.models.user.User)
          Returns:
              A view (instance of timesketch.models.sketch.View)
          """
          view = View(
              name=name,
              query_string=name,
              query_filter=json.dumps(dict()),
              user=user,
              sketch=sketch)
          self._commit_to_database(view)
          return view
  
      def _create_searchtemplate(self, name, user):
          """Create a search template in the database.
          Args:
              name: Name of the view (string)
              user: A user (instance of timesketch.models.user.User)
          Returns:
              A search template (timesketch.models.sketch.SearchTemplate)
          """
          searchtemplate = SearchTemplate(
              name=name,
              query_string=name,
              query_filter=json.dumps(dict()),
              user=user)
          self._commit_to_database(searchtemplate)
          return searchtemplate
  
      def setUp(self):
          """Setup the test database."""
          init_db()
  
          self.user1 = self._create_user(username='test1', set_password=True)
          self.user2 = self._create_user(username='test2', set_password=False)
  
          self.group1 = self._create_group(name='test_group1', user=self.user1)
          self.group2 = self._create_group(name='test_group2', user=self.user1)
  
          self.sketch1 = self._create_sketch(
              name='Test 1', user=self.user1, acl=True)
          self.sketch2 = self._create_sketch(
              name='Test 2', user=self.user1, acl=False)
          self.sketch3 = self._create_sketch(
              name='Test 3', user=self.user1, acl=True)
  
          self.searchindex = self._create_searchindex(
              name='test', user=self.user1, acl=True)
          self.searchindex2 = self._create_searchindex(
              name='test2', user=self.user1, acl=True)
  
          self.timeline = self._create_timeline(
              name='Timeline 1',
              sketch=self.sketch1,
              searchindex=self.searchindex,
              user=self.user1)
  
          self.view1 = self._create_view(
              name='View 1', sketch=self.sketch1, user=self.user1)
          self.view2 = self._create_view(
              name='View 2', sketch=self.sketch2, user=self.user1)
          self.view3 = self._create_view(
              name='', sketch=self.sketch1, user=self.user2)
  
          self.searchtemplate = self._create_searchtemplate(
              name='template', user=self.user1)
  
          self.event = self._create_event(
              sketch=self.sketch1, searchindex=self.searchindex, user=self.user1)
  
          self.story = self._create_story(sketch=self.sketch1, user=self.user1)
  
      def tearDown(self):
          """Tear down the test database."""
          db_session.remove()
          drop_all()
  
      def login(self):
          """Authenticate the test user."""
          self.client.post(
              '/login/',
              data=dict(username='test1', password='test'),
              follow_redirects=True)
  
      def test_unauthenticated(self):
          """
          Generic test for all resources. It tests that no
          unauthenticated request are accepted.
          """
          if not getattr(self, 'resource_url', False):
              self.skipTest(self)
  
          response = self.client.get(self.resource_url)
          if response.status_code == 405:
              response = self.client.post(self.resource_url)
          if isinstance(response.data, six.binary_type):
              response_data = codecs.decode(response.data, 'utf-8')
          else:
              response_data = response.data
          self.assertIn('/login/', response_data)
          self.assertEqual(response.status_code, HTTP_STATUS_CODE_REDIRECT)
  
  
  class ModelBaseTest(BaseTest):
      """Base class for database model tests."""
  
      def _test_db_object(self, expected_result=None, model_cls=None):
          """Generic test that checks if the stored data is correct."""
          db_obj = model_cls.query.get(1)
          for x in expected_result:
              k, v = x[0], x[1]
              self.assertEqual(db_obj.__getattribute__(k), v)
  