#!/usr/bin/env python
# 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 is for management of the timesketch application."""

from collections import Counter
import json
import sys
from uuid import uuid4

from flask import current_app
from flask_script import Command
from flask_script import Manager
from flask_script import Server
from flask_script import Option
from flask_script import prompt_bool
from flask_script import prompt_pass
from sqlalchemy.exc import IntegrityError

from timesketch import create_app
from timesketch.lib.datastores.elastic import ElasticSearchDataStore
from timesketch.models import db_session
from timesketch.models import drop_all
from timesketch.models.user import User
from timesketch.models.sketch import SearchIndex


class DropDataBaseTables(Command):
    """Drop all database tables."""
    def __init__(self):
        super(DropDataBaseTables, self).__init__()

    # pylint: disable=method-hidden
    def run(self):
        """Drop all tables after user ha verified."""
        verified = prompt_bool(
            u'Do you really want to drop all the database tables?')
        if verified:
            sys.stdout.write(u'All tables dropped. Database is now empty.\n')
            drop_all()


class AddUser(Command):
    """Create a new Timesketch user."""
    option_list = (
        Option(u'--username', u'-u', dest=u'username', required=True),
        Option(u'--password', u'-p', dest=u'password', required=False),
    )

    def __init__(self):
        super(AddUser, self).__init__()

    def get_password_from_prompt(self):
        """Get password from the command line prompt."""
        first_password = prompt_pass(u'Enter password')
        second_password = prompt_pass(u'Enter password again')
        if first_password != second_password:
            sys.stderr.write(u'Passwords don\'t match, try again.\n')
            self.get_password_from_prompt()
        return first_password

    # pylint: disable=arguments-differ, method-hidden
    def run(self, username, password):
        """Creates the user."""
        if not password:
            password = self.get_password_from_prompt()
        password = unicode(password.decode(encoding=u'utf-8'))
        username = unicode(username.decode(encoding=u'utf-8'))
        user = User(username=username, name=username)
        user.set_password(plaintext=password)
        try:
            db_session.add(user)
            db_session.commit()
            sys.stdout.write(u'User {0:s} created\n'.format(username))
        except IntegrityError:
            sys.stderr.write(
                u'The username ({0:s}) is already taken, '
                u'try another one.\n'.format(username))


class AddSearchIndex(Command):
    """Create a new Timesketch searchindex."""
    option_list = (
        Option(u'--name', u'-n', dest=u'name', required=True),
        Option(u'--index', u'-i', dest=u'index', required=True),
        Option(u'--user', u'-u', dest=u'username', required=True),
    )

    def __init__(self):
        super(AddSearchIndex, self).__init__()

    # pylint: disable=arguments-differ, method-hidden
    def run(self, name, index, username):
        """Create the SearchIndex."""
        es = ElasticSearchDataStore(
            host=current_app.config[u'ELASTIC_HOST'],
            port=current_app.config[u'ELASTIC_PORT'])
        user = User.query.filter_by(username=username).first()
        if not user:
            sys.stderr.write(u'User does not exist\n')
            sys.exit(1)
        if not es.client.indices.exists(index=index):
            sys.stderr.write(u'Index does not exist in the datastore\n')
            sys.exit(1)
        if SearchIndex.query.filter_by(name=name, index_name=index).first():
            sys.stderr.write(
                u'Index with this name already exist in Timesketch\n')
            sys.exit(1)
        searchindex = SearchIndex(
            name=name, description=name, user=user, index_name=index)
        searchindex.grant_permission(None, u'read')
        db_session.add(searchindex)
        db_session.commit()
        sys.stdout.write(u'Search index {0:s} created\n'.format(name))


class CreateTimelineFromJson(Command):
    """Create a new Timesketch timeline from a JSON file."""
    DEFAULT_FLUSH_INTERVAL = 1000
    DEFAULT_EVENT_TYPE = u'generic_event'
    DEFAULT_INDEX_NAME = uuid4().hex
    option_list = (
        Option(u'--name', u'-n', dest=u'timeline_name', required=True,
               help=u'Name of the timeline as it will appear in the '
                    u'Timesketch UI.'),
        Option(u'--file', u'-f', dest=u'filename', required=True,
               help=u'Path to the JSON file to process'),
        Option(
            u'--index_name', dest=u'index_name', required=False,
            default=DEFAULT_INDEX_NAME,
            help=u'OPTIONAL: Name of the Elasticsearch index. Specify an '
                 u'existing one to append to it. Default: unique UUID'),
        Option(
            u'--event_type', dest=u'event_type', required=False,
            default=DEFAULT_EVENT_TYPE,
            help=u'OPTIONAL: Type of event. This is what becomes the '
                 u'document type in Elasticsearch. '
                 u'Default: {0:s}.'.format(DEFAULT_EVENT_TYPE)),

        Option(
            u'--flush_interval', dest=u'flush_interval', required=False,
            default=DEFAULT_FLUSH_INTERVAL,
            help=u'OPTIONAL: How often to bulk insert events to Elasticsearch. '
                 u'Default: Every {0:d} event.'.format(DEFAULT_FLUSH_INTERVAL))
    )

    def __init__(self):
        super(CreateTimelineFromJson, self).__init__()

    def run(
            self, timeline_name, index_name, filename, event_type,
            flush_interval):
        """Create the timeline.

        NOTE: The way the built in json module works is by loading the whole
        document into memory. This is of course not optimal and expensive.
        Don't try to ingest too big of a document. We will implement a CSV
        importer in the future to address this.

        Elasticsearch is very forgiving about how your JSON is structured, but
        Timesketch has a minimum set of attributes that needs to be present to
        render correctly in the UI. These are:

        * message
        * timestamp
        * datetime
        * timestamp_desc

        You can of course have more attributes, and these will be indexed
        automatically and shown in the detailed view of the event.

        Example (minimal) JSON structure:
        [
            {
              "message": "foo",
              "timestamp": "1432293395",
              "datetime": "2015-05-22T13:16:35+00:00",
              "timestamp_desc": "Some timestamp",
            },
            {
              "message": "bar",
              "timestamp": "1432293443",
              "datetime": "2015-05-22T13:17:23+00:00",
              "timestamp_desc": "Some timestamp",
            }
        ]

        Args:
            timeline_name: Name for the new timeline.
            index_name: Name for the index in Elasticsearch.
            filename: Path to JSON file.
            event_type: Type of event, i.e. doc_type in Elasticsearch.
            flush_interval: Number of events to queue up before bulk insert.
        """
        counter = Counter()
        events = []
        es = ElasticSearchDataStore(
            host=current_app.config[u'ELASTIC_HOST'],
            port=current_app.config[u'ELASTIC_PORT'])
        timeline_name = unicode(timeline_name.decode(encoding=u'utf-8'))
        flush_interval = int(flush_interval)

        def _insert_events_to_elasticsearch(_events, _index, _event_type):
            """Bulk insert events into Elasticsearch.

            Args:
                _events: List of events to insert.
                _index: Elasticsearch index to add the events to.
                _event_type: Document type for the events.
            """
            es.client.bulk(
                index=_index, doc_type=_event_type, body=_events)

        try:
            with open(filename, u'rb') as fh:
                # This is expensive, i.e. whole file is read into memory.
                # TODO: Implement CSV importer that can be made more efficient.
                # TODO: Check file size and bail out if big.
                index_name, event_type = es.create_index(
                    index_name=index_name, doc_type=event_type)
                for event in json.load(fh):
                    # Header needed by Elasticsearch when bulk inserting.
                    events.append({
                        u'index': {
                            u'_index': index_name, u'_type': event_type
                        }
                    })
                    events.append(event)
                    counter[u'events'] += 1
                    if counter[u'events'] % flush_interval == 0:
                        _insert_events_to_elasticsearch(
                            events, index_name, event_type)
                        events = []
                # Insert any remaining events.
                if events:
                    _insert_events_to_elasticsearch(
                        events, index_name, event_type)

            # Create the searchindex in the Timesketch database.
            searchindex = SearchIndex.get_or_create(
                name=timeline_name, description=timeline_name,
                user=None, index_name=index_name)
            searchindex.grant_permission(None, u'read')
            db_session.add(searchindex)
            db_session.commit()
            sys.stdout.write(
                u'Timeline name: {0:s}\nElasticsearch index: {1:s}\n'
                u'Events inserted: {2:d}\n'.format(
                    timeline_name, index_name, counter[u'events']))
        except IOError as exception:
            sys.stderr.write(u'Error: {0:s}\n'.format(exception))


if __name__ == '__main__':
    # Setup Flask-script command manager and register commands.
    shell_manager = Manager(create_app)
    shell_manager.add_command(u'add_user', AddUser())
    shell_manager.add_command(u'add_index', AddSearchIndex())
    shell_manager.add_command(u'drop_db', DropDataBaseTables())
    shell_manager.add_command(u'json2ts', CreateTimelineFromJson())
    shell_manager.add_command(u'runserver', Server(
        host=u'127.0.0.1', port=5000))
    shell_manager.add_option(
        u'-c', u'--config', dest=u'config', default=u'/etc/timesketch.conf',
        required=False)
    shell_manager.run()
