tsctl

#!/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 csv
import json
import sys
from uuid import uuid4

from flask import current_app
from flask_migrate import MigrateCommand
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 Group
from timesketch.models.user import User
from timesketch.models.sketch import SearchIndex

def get_human_size(size):
    for x in ['bytes', 'KB', 'MB', 'GB']:
        if size < 1024.0:
            return "%3.1f%s" % (size, x)
        size /= 1024.0
    return "%3.1f%s" % (size, 'TB')

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.get_or_create(username=username)
        user.set_password(plaintext=password)
        db_session.add(user)
        db_session.commit()
        sys.stdout.write(u'User {0:s} created/updated\n'.format(username))


class AddGroup(Command):
    """Create a new Timesketch group."""
    option_list = (
        Option(u'--name', u'-n', dest=u'name', required=True),
    )

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

    # pylint: disable=arguments-differ, method-hidden
    def run(self, name):
        """Creates the group."""
        name = unicode(name.decode(encoding=u'utf-8'))
        group = Group.get_or_create(name=name)
        db_session.add(group)
        db_session.commit()
        sys.stdout.write(u'Group {0:s} created\n'.format(name))


class GroupManager(Command):
    """Manage group memberships."""
    option_list = (
        Option(
            u'--add', u'-a', dest=u'add', action=u'store_true', required=False,
            default=False),
        Option(
            u'--remove', u'-r', dest=u'remove', action=u'store_true',
            required=False, default=False),
        Option(u'--group', u'-g', dest=u'group_name', required=True),
        Option(u'--user', u'-u', dest=u'user_name', required=True),
    )

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

    # pylint: disable=arguments-differ, method-hidden
    def run(self, add, remove, group_name, user_name):
        """Add the user to the group."""
        group_name = unicode(group_name.decode(encoding=u'utf-8'))
        user_name = unicode(user_name.decode(encoding=u'utf-8'))
        group = Group.query.filter_by(name=group_name).first()
        user = User.query.filter_by(username=user_name).first()

        # Add or remove user from group
        if remove:
            try:
                user.groups.remove(group)
                sys.stdout.write(
                    u'{0:s} removed from group {1:s}\n'.format(
                        user_name, group_name))
                db_session.commit()
            except ValueError:
                sys.stdout.write(
                    u'{0:s} is not a member of group {1:s}\n'.format(
                        user_name, group_name))
        else:
            user.groups.append(group)
            try:
                db_session.commit()
                sys.stdout.write(
                    u'{0:s} added to group {1:s}\n'.format(
                        user_name, group_name))
            except IntegrityError:
                sys.stdout.write(
                    u'{0:s} is already a member of group {1:s}\n'.format(
                        user_name, group_name))


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(u'read')
        db_session.add(searchindex)
        db_session.commit()
        sys.stdout.write(u'Search index {0:s} created\n'.format(name))


class CreateTimelineBase(Command):
    """Base class for file based ingestion of events."""
    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'file_path', 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(CreateTimelineBase, self).__init__()
        self.counter = Counter()
        self.events = []

    @staticmethod
    def create_search_index(index_name, timeline_name):
            """Create searchindex in the Timesketch database.

            Args:
                index_name: Then name of the index in Elasticsearch
                timeline_name: The name of the timeline in Timesketch
            """
            searchindex = SearchIndex.get_or_create(
                name=timeline_name, description=timeline_name,
                user=None, index_name=index_name)
            searchindex.grant_permission(u'read')
            db_session.add(searchindex)
            db_session.commit()

    def add_event(self, es, flush_interval, index_name, event_type, event):
        """Add event to Elasticsearch.

        Args:
            es = Elasticsearch client (instance of
            timesketch.lib.datastores.elastic.ElasticSearchDatastore)
            flush_interval: Number of events to queue up before indexing
            index_name: Name of the index in Elasticsearch
            event_type: Type of event (e.g. plaso_event)
            event: Event dictionary
        """
        # Header needed by Elasticsearch when bulk inserting.
        self.events.append({
            u'index': {
                u'_index': index_name, u'_type': event_type
            }
        })
        self.events.append(event)
        self.counter[u'events'] += 1
        if self.counter[u'events'] % int(flush_interval) == 0:
            sys.stdout.write(
                u'Events inserted: {0:d}\n'.format(self.counter[u'events']))
            es.client.bulk(
                index=index_name, doc_type=event_type,
                body=self.events, request_timeout=30)
            self.events = []

    def finish(self, es, timeline_name, index_name, event_type):
        """Send any last events for indexing and create the timeline.

        Args:
            es = Elasticsearch client (instance of
            timesketch.lib.datastores.elastic.ElasticSearchDatastore)
            timeline_name: The name of the timeline in Timesketch
            index_name: Name of the index in Elasticsearch
            event_type: Type of event (e.g. plaso_event)
        """
        if self.events:
            es.client.bulk(
                index=index_name, doc_type=event_type, body=self.events)

        self.create_search_index(index_name, timeline_name)
        sys.stdout.write(
            u'Timeline name: {0:s}\nElasticsearch index: {1:s}\n'
            u'Events inserted: {2:d}\n'.format(
                timeline_name, index_name, self.counter[u'events']))

    def run(self, timeline_name, index_name, file_path, event_type,
            flush_interval):
        """Flask-script entrypoint for running the command.

        Args:
            timeline_name: The name of the timeline in Timesketch
            index_name: Name of the index in Elasticsearch
            file_path: Path to the file to process
            event_type: Type of event (e.g. plaso_event)
            flush_interval: Number of events to queue up before bulk insert
        """
        return NotImplementedError


class CreateTimelineFromJson(CreateTimelineBase):
    """Create a new Timesketch timeline from a JSON file."""

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

    def run(self, timeline_name, index_name, file_path, event_type,
            flush_interval):
        """Create timeline from a JSON file.

        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: The name of the timeline in Timesketch
            index_name: Name of the index in Elasticsearch
            file_path: Path to the file to process
            event_type: Type of event (e.g. plaso_event)
            flush_interval: Number of events to queue up before bulk insert
        """
        from os.path import getsize
        print("[+] JSON file to ingest is: '{0}'".format(get_human_size(getsize(file_path))))
        
        timeline_name = unicode(timeline_name.decode(encoding=u'utf-8'))
        index_name = unicode(index_name.decode(encoding=u'utf-8'))
        es = ElasticSearchDataStore(
            host=current_app.config[u'ELASTIC_HOST'],
            port=current_app.config[u'ELASTIC_PORT'])

        with open(file_path, u'rb') as fh:
            # This is expensive, i.e. whole file is read into memory.
            # TODO: Check file size and bail out if big.
            es.create_index(index_name=index_name, doc_type=event_type)
            for event in json.load(fh):
                self.add_event(
                    es, flush_interval, index_name, event_type, event)
        # Insert any remaining events.
        self.finish(es, timeline_name, index_name, event_type)


class CreateTimelineFromCsv(CreateTimelineBase):
    """Create a new Timesketch timeline from a CSV file."""

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

    def run(self, timeline_name, index_name, file_path, event_type,
            flush_interval):
        """Create the timeline from a CSV file.

        Args:
            timeline_name: The name of the timeline in Timesketch
            index_name: Name of the index in Elasticsearch
            file_path: Path to the file to process
            event_type: Type of event (e.g. plaso_event)
            flush_interval: Number of events to queue up before bulk insert

        """
        timeline_name = unicode(timeline_name.decode(encoding=u'utf-8'))
        index_name = unicode(index_name.decode(encoding=u'utf-8'))
        es = ElasticSearchDataStore(
            host=current_app.config[u'ELASTIC_HOST'],
            port=current_app.config[u'ELASTIC_PORT'])

        def read_csv(path):
            """Generator for reading a CSV file.

            Args:
                path: Path to the CSV file
            """
            with open(path, u'rb') as fh:
                reader = csv.DictReader(fh)
                for row in reader:
                    yield row

        es.create_index(index_name=index_name, doc_type=event_type)
        for event in read_csv(file_path):
            self.add_event(es, flush_interval, index_name, event_type, event)
        self.finish(es, timeline_name, index_name, event_type)


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_group', AddGroup())
    shell_manager.add_command(u'manage_group', GroupManager())
    shell_manager.add_command(u'add_index', AddSearchIndex())
    shell_manager.add_command(u'csv2ts', CreateTimelineFromCsv())
    shell_manager.add_command(u'db', MigrateCommand)
    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()