# 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 holds version 1 of the Timesketch API.

The timesketch API is a RESTful API that exposes the following resources:

GET /sketches/
GET /sketches/:sketch_id/
GET /sketches/:sketch_id/explore/
GET /sketches/:sketch_id/event/
GET /sketches/:sketch_id/views/
GET /sketches/:sketch_id/views/:view_id/

POST /sketches/:sketch_id/event/
POST /sketches/:sketch_id/event/annotate/
POST /sketches/:sketch_id/views/
"""

import datetime
import json
import os
import uuid

from flask import abort
from flask import jsonify
from flask import request
from flask_login import current_user
from flask_login import current_app
from flask_login import login_required
from flask_restful import fields
from flask_restful import marshal
from flask_restful import reqparse
from flask_restful import Resource
from sqlalchemy import desc
from sqlalchemy import not_

from timesketch.lib.aggregators import heatmap
from timesketch.lib.definitions import HTTP_STATUS_CODE_OK
from timesketch.lib.definitions import HTTP_STATUS_CODE_CREATED
from timesketch.lib.definitions import HTTP_STATUS_CODE_BAD_REQUEST
from timesketch.lib.definitions import HTTP_STATUS_CODE_FORBIDDEN
from timesketch.lib.definitions import HTTP_STATUS_CODE_NOT_FOUND
from timesketch.lib.datastores.elastic import ElasticSearchDataStore
from timesketch.lib.errors import ApiHTTPError
from timesketch.lib.forms import AddTimelineForm
from timesketch.lib.forms import AggregationForm
from timesketch.lib.forms import SaveViewForm
from timesketch.lib.forms import NameDescriptionForm
from timesketch.lib.forms import EventAnnotationForm
from timesketch.lib.forms import ExploreForm
from timesketch.lib.forms import UploadFileForm
from timesketch.models import db_session
from timesketch.models.sketch import Event
from timesketch.models.sketch import SearchIndex
from timesketch.models.sketch import Sketch
from timesketch.models.sketch import Timeline
from timesketch.models.sketch import View


class ResourceMixin(object):
    """Mixin for API resources."""
    # Schemas for database model resources

    searchindex_status_fields = {
        u'id': fields.Integer,
        u'status': fields.String,
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    searchindex_fields = {
        u'id': fields.Integer,
        u'name': fields.String,
        u'index_name': fields.String,
        u'status': fields.Nested(searchindex_status_fields),
        u'deleted': fields.Boolean,
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    timeline_fields = {
        u'id': fields.Integer,
        u'name': fields.String,
        u'description': fields.String,
        u'color': fields.String,
        u'searchindex': fields.Nested(searchindex_fields),
        u'deleted': fields.Boolean,
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    view_fields = {
        u'id': fields.Integer,
        u'name': fields.String,
        u'query_string': fields.String,
        u'query_filter': fields.String,
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    user_fields = {
        u'username': fields.String
    }

    sketch_fields = {
        u'id': fields.Integer,
        u'name': fields.String,
        u'description': fields.String,
        u'user': fields.Nested(user_fields),
        u'timelines': fields.Nested(timeline_fields),
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    comment_fields = {
        u'comment': fields.String,
        u'user': fields.Nested(user_fields),
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    label_fields = {
        u'name': fields.String,
        u'user': fields.Nested(user_fields),
        u'created_at': fields.DateTime,
        u'updated_at': fields.DateTime
    }

    fields_registry = {
        u'searchindex': searchindex_fields,
        u'timeline': timeline_fields,
        u'view': view_fields,
        u'user': user_fields,
        u'sketch': sketch_fields,
        u'event_comment': comment_fields,
        u'event_label': label_fields
    }

    @property
    def datastore(self):
        """Property to get an instance of the datastore backend.

        Returns:
            Instance of timesketch.lib.datastores.elastic.ElasticSearchDatastore
        """
        return ElasticSearchDataStore(
            host=current_app.config[u'ELASTIC_HOST'],
            port=current_app.config[u'ELASTIC_PORT'])

    def to_json(
            self, model, model_fields=None, meta=None,
            status_code=HTTP_STATUS_CODE_OK):
        """Create json response from a database models.

        Args:
            model: Instance of a timesketch database model
            model_fields: Dictionary describing the resulting schema
            meta: Dictionary holding any metadata for the result
            status_code: Integer used as status_code in the response

        Returns:
            Response in json format (instance of flask.wrappers.Response)
        """
        if not meta:
            meta = dict()
        if not model_fields:
            try:
                model_fields = self.fields_registry[model.__tablename__]
            except AttributeError:
                model_fields = self.fields_registry[model[0].__tablename__]

        schema = {
            u'meta': meta,
            u'objects': [marshal(model, model_fields)]
        }
        response = jsonify(schema)
        response.status_code = status_code
        return response


class SketchListResource(ResourceMixin, Resource):
    """Resource for listing sketches."""
    def __init__(self):
        super(SketchListResource, self).__init__()
        self.parser = reqparse.RequestParser()
        self.parser.add_argument(u'name', type=unicode, required=True)
        self.parser.add_argument(u'description', type=unicode, required=False)

    @login_required
    def get(self):
        """Handles GET request to the resource.

        Returns:
            List of sketches (instance of flask.wrappers.Response)
        """
        # TODO: Handle offset parameter
        sketches = Sketch.all_with_acl()
        paginated_result = sketches.paginate(1, 10, False)
        meta = {
            u'next': paginated_result.next_num,
            u'previous': paginated_result.prev_num,
            u'offset': paginated_result.page,
            u'limit': paginated_result.per_page
        }
        if not paginated_result.has_prev:
            meta[u'previous'] = None
        if not paginated_result.has_next:
            meta[u'next'] = None
        result = self.to_json(paginated_result.items, meta=meta)
        return result

    @login_required
    def post(self):
        """Handles POST request to the resource.

        Returns:
            A sketch in JSON (instance of flask.wrappers.Response)
        """
        form = NameDescriptionForm.build(request)
        if form.validate_on_submit():
            sketch = Sketch(
                name=form.name.data, description=form.description.data,
                user=current_user)
            sketch.status.append(sketch.Status(user=None, status=u'new'))
            # Give the requesting user permissions on the new sketch.
            sketch.grant_permission(permission=u'read', user=current_user)
            sketch.grant_permission(permission=u'write', user=current_user)
            sketch.grant_permission(permission=u'delete', user=current_user)
            db_session.add(sketch)
            db_session.commit()
            return self.to_json(sketch, status_code=HTTP_STATUS_CODE_CREATED)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)


class SketchResource(ResourceMixin, Resource):
    """Resource to get a sketch."""
    @login_required
    def get(self, sketch_id):
        """Handles GET request to the resource.

        Returns:
            A sketch in JSON (instance of flask.wrappers.Response)
        """
        sketch = Sketch.query.get_with_acl(sketch_id)
        meta = dict(
            views=[
                {
                    u'name': view.name,
                    u'id': view.id
                } for view in sketch.get_named_views.all()
            ])
        return self.to_json(sketch, meta=meta)

    @login_required
    def post(self, sketch_id):
        """Handles POST request to the resource.

        Returns:
            A sketch in JSON (instance of flask.wrappers.Response)

        Raises:
            ApiHTTPError
        """
        sketch = Sketch.query.get_with_acl(sketch_id)
        searchindices_in_sketch = [t.searchindex.id for t in sketch.timelines]
        indices = SearchIndex.all_with_acl(
            current_user).order_by(
                desc(SearchIndex.created_at)).filter(
                    not_(SearchIndex.id.in_(searchindices_in_sketch)))

        add_timeline_form = AddTimelineForm.build(request)
        add_timeline_form.timelines.choices = set(
            (i.id, i.name) for i in indices.all())

        if add_timeline_form.validate_on_submit():
            if not sketch.has_permission(current_user, u'write'):
                abort(HTTP_STATUS_CODE_FORBIDDEN)
            for searchindex_id in add_timeline_form.timelines.data:
                searchindex = SearchIndex.query.get_with_acl(searchindex_id)
                if searchindex not in [t.searchindex for t in sketch.timelines]:
                    _timeline = Timeline(
                        name=searchindex.name,
                        description=searchindex.description,
                        sketch=sketch,
                        user=current_user,
                        searchindex=searchindex)
                    db_session.add(_timeline)
                    sketch.timelines.append(_timeline)
            db_session.commit()
            return self.to_json(sketch, status_code=HTTP_STATUS_CODE_CREATED)
        else:
            raise ApiHTTPError(
                message=add_timeline_form.errors,
                status_code=HTTP_STATUS_CODE_BAD_REQUEST)


class ViewListResource(ResourceMixin, Resource):
    """Resource to create a View."""
    @login_required
    def post(self, sketch_id):
        """Handles POST request to the resource.

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            A view in JSON (instance of flask.wrappers.Response)
        """
        form = SaveViewForm.build(request)
        if form.validate_on_submit():
            sketch = Sketch.query.get_with_acl(sketch_id)
            view = View(
                name=form.name.data, sketch=sketch, user=current_user,
                query_string=form.query.data,
                query_filter=json.dumps(form.filter.data, ensure_ascii=False))
            db_session.add(view)
            db_session.commit()
            return self.to_json(view, status_code=HTTP_STATUS_CODE_CREATED)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)


class ViewResource(ResourceMixin, Resource):
    """Resource to get a view."""
    @login_required
    def get(self, sketch_id, view_id):
        """Handles GET request to the resource.

        Args:
            sketch_id: Integer primary key for a sketch database model
            view_id: Integer primary key for a view database model

        Returns:
            A view in JSON (instance of flask.wrappers.Response)
        """
        sketch = Sketch.query.get_with_acl(sketch_id)
        view = View.query.get(view_id)

        # Check that this view belongs to the sketch
        if view.sketch_id != sketch.id:
            abort(HTTP_STATUS_CODE_NOT_FOUND)

        # If this is a user state view, check that it
        # belongs to the current_user
        if view.name == u'' and view.user != current_user:
            abort(HTTP_STATUS_CODE_FORBIDDEN)
        return self.to_json(view)


class ExploreResource(ResourceMixin, Resource):
    """Resource to search the datastore based on a query and a filter."""
    @login_required
    def post(self, sketch_id):
        """Handles POST request to the resource.
        Handler for /api/v1/sketches/:sketch_id/explore/

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            JSON with list of matched events
        """
        sketch = Sketch.query.get_with_acl(sketch_id)
        form = ExploreForm.build(request)

        if form.validate_on_submit():
            query_filter = form.filter.data
            sketch_indices = [
                t.searchindex.index_name for t in sketch.timelines]
            indices = query_filter.get(u'indices', sketch_indices)

            # Make sure that the indices in the filter are part of the sketch
            if set(indices) - set(sketch_indices):
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            # Make sure we have a query string or star filter
            if not form.query.data and not query_filter.get(u'star'):
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            result = self.datastore.search(
                sketch_id, form.query.data, query_filter, indices,
                aggregations=None, return_results=True)

            # Get labels for each event that matches the sketch.
            # Remove all other labels.
            for event in result[u'hits'][u'hits']:
                event[u'selected'] = False
                event[u'_source'][u'label'] = []
                try:
                    for label in event[u'_source'][u'timesketch_label']:
                        if sketch.id != label[u'sketch_id']:
                            continue
                        event[u'_source'][u'label'].append(label[u'name'])
                    del event[u'_source'][u'timesketch_label']
                except KeyError:
                    pass

            # Update or create user state view. This is used in the UI to let
            # the user get back to the last state in the explore view.
            view = View.get_or_create(
                user=current_user, sketch=sketch, name=u'')
            view.query_string = form.query.data
            view.query_filter = json.dumps(query_filter)
            db_session.add(view)
            db_session.commit()

            # Add metadata for the query result. This is used by the UI to
            # render the event correctly and to display timing and hit count
            # information.
            tl_colors = {}
            tl_names = {}
            for timeline in sketch.timelines:
                tl_colors[timeline.searchindex.index_name] = timeline.color
                tl_names[timeline.searchindex.index_name] = timeline.name

            try:
                buckets = result[u'aggregations'][u'data_type'][u'buckets']
            except KeyError:
                buckets = None

            es_total_count_unfiltered = 0
            if buckets:
                for bucket in buckets:
                    es_total_count_unfiltered += bucket[u'doc_count']

            meta = {
                u'es_time': result[u'took'],
                u'es_total_count': result[u'hits'][u'total'],
                u'es_total_count_unfiltered': es_total_count_unfiltered,
                u'timeline_colors': tl_colors,
                u'timeline_names': tl_names,
                u'histogram': buckets
            }
            schema = {
                u'meta': meta,
                u'objects': result[u'hits'][u'hits']
            }
            return jsonify(schema)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)


class AggregationResource(ResourceMixin, Resource):
    """Resource to query for aggregated results."""
    @login_required
    def post(self, sketch_id):
        """Handles POST request to the resource.
        Handler for /api/v1/sketches/:sketch_id/aggregation/

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            JSON with aggregation results
        """
        sketch = Sketch.query.get_with_acl(sketch_id)
        form = AggregationForm.build(request)

        if form.validate_on_submit():
            query_filter = form.filter.data
            sketch_indices = [
                t.searchindex.index_name for t in sketch.timelines]
            indices = query_filter.get(u'indices', sketch_indices)

            # Make sure that the indices in the filter are part of the sketch
            if set(indices) - set(sketch_indices):
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            # Make sure we have a query string or star filter
            if not form.query.data and not query_filter.get(u'star'):
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            result = []
            if form.aggtype.data == u'heatmap':
                result = heatmap(
                    es_client=self.datastore, sketch_id=sketch_id,
                    query=form.query.data, query_filter=query_filter,
                    indices=indices)
            else:
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            schema = {
                u'objects': result
            }
            return jsonify(schema)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)


class EventResource(ResourceMixin, Resource):
    """Resource to get a single event from the datastore.

    HTTP Args:
        searchindex_id: The datastore searchindex id as string
        event_id: The datastore event id as string
    """
    def __init__(self):
        super(EventResource, self).__init__()
        self.parser = reqparse.RequestParser()
        self.parser.add_argument(u'searchindex_id', type=unicode, required=True)
        self.parser.add_argument(u'event_id', type=unicode, required=True)

    @login_required
    def get(self, sketch_id):
        """Handles GET request to the resource.
        Handler for /api/v1/sketches/:sketch_id/event/

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            JSON of the datastore event
        """

        args = self.parser.parse_args()
        sketch = Sketch.query.get_with_acl(sketch_id)
        searchindex_id = args.get(u'searchindex_id')
        searchindex = SearchIndex.query.filter_by(
            index_name=searchindex_id).first()
        event_id = args.get(u'event_id')
        indices = [t.searchindex.index_name for t in sketch.timelines]

        # Check if the requested searchindex is part of the sketch
        if searchindex_id not in indices:
            abort(HTTP_STATUS_CODE_BAD_REQUEST)

        result = self.datastore.get_event(searchindex_id, event_id)

        event = Event.query.filter_by(
            sketch=sketch, searchindex=searchindex,
            document_id=event_id).first()

        # Comments for this event
        comments = []
        if event:
            for comment in event.comments:
                comment_dict = {
                    u'user': {
                        u'username': comment.user.username,
                    },
                    u'created_at': comment.created_at,
                    u'comment': comment.comment
                }
                comments.append(comment_dict)

        schema = {
            u'meta': {
                u'comments': comments
            },
            u'objects': result[u'_source']
        }
        return jsonify(schema)


class EventAnnotationResource(ResourceMixin, Resource):
    """Resource to create an annotation for an event."""
    @login_required
    def post(self, sketch_id):
        """Handles POST request to the resource.

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            An annotation in JSON (instance of flask.wrappers.Response)
        """
        form = EventAnnotationForm.build(request)
        if form.validate_on_submit():
            annotations = []
            sketch = Sketch.query.get_with_acl(sketch_id)
            indices = [t.searchindex.index_name for t in sketch.timelines]
            annotation_type = form.annotation_type.data
            events = form.events.raw_data

            for _event in events:
                searchindex_id = _event[u'_index']
                searchindex = SearchIndex.query.filter_by(
                    index_name=searchindex_id).first()
                event_id = _event[u'_id']
                event_type = _event[u'_type']

                if searchindex_id not in indices:
                    abort(HTTP_STATUS_CODE_BAD_REQUEST)

                # Get or create an event in the SQL database to have something
                # to attach the annotation to.
                event = Event.get_or_create(
                    sketch=sketch, searchindex=searchindex,
                    document_id=event_id)

                # Add the annotation to the event object.
                if u'comment' in annotation_type:
                    annotation = Event.Comment(
                        comment=form.annotation.data, user=current_user)
                    event.comments.append(annotation)
                    self.datastore.set_label(
                        searchindex_id, event_id, event_type, sketch.id,
                        current_user.id, u'__ts_comment', toggle=False)

                elif u'label' in annotation_type:
                    annotation = Event.Label.get_or_create(
                        label=form.annotation.data, user=current_user)
                    if annotation not in event.labels:
                        event.labels.append(annotation)
                    toggle = False
                    if u'__ts_star' in form.annotation.data:
                        toggle = True
                    self.datastore.set_label(
                        searchindex_id, event_id, event_type, sketch.id,
                        current_user.id, form.annotation.data, toggle=toggle)
                else:
                    abort(HTTP_STATUS_CODE_BAD_REQUEST)

                annotations.append(annotation)
                # Save the event to the database
                db_session.add(event)
                db_session.commit()
            return self.to_json(
                annotations, status_code=HTTP_STATUS_CODE_CREATED)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)


class UploadFileResource(ResourceMixin, Resource):
    """Resource that processes uploaded files."""
    @login_required
    def post(self):
        """Handles POST request to the resource.

        Returns:
            A view in JSON (instance of flask.wrappers.Response)

        Raises:
            ApiHTTPError
        """
        UPLOAD_ENABLED = current_app.config[u'UPLOAD_ENABLED']
        UPLOAD_FOLDER = current_app.config[u'UPLOAD_FOLDER']

        form = UploadFileForm()
        if form.validate_on_submit() and UPLOAD_ENABLED:
            from timesketch.lib.tasks import run_plaso
            file_storage = form.file.data
            timeline_name = form.name.data
            # We do not need a human readable filename or
            # datastore index name, so we use UUIDs here.
            filename = unicode(uuid.uuid4().hex)
            index_name = unicode(uuid.uuid4().hex)

            file_path = os.path.join(UPLOAD_FOLDER, filename)
            file_storage.save(file_path)

            search_index = SearchIndex.get_or_create(
                name=timeline_name, description=timeline_name, user=None,
                index_name=index_name)
            search_index.grant_permission(permission=u'read')
            search_index.set_status(u'processing')
            db_session.add(search_index)
            db_session.commit()

            run_plaso.apply_async(
                (file_path, timeline_name, index_name), task_id=index_name)

            return self.to_json(
                search_index, status_code=HTTP_STATUS_CODE_CREATED)
        else:
            raise ApiHTTPError(
                message=form.errors[u'file'][0],
                status_code=HTTP_STATUS_CODE_BAD_REQUEST)


class TaskResource(ResourceMixin, Resource):
    """Resource to get information on celery task."""
    def __init__(self):
        super(TaskResource, self).__init__()
        from timesketch import create_celery_app
        self.celery = create_celery_app()

    @login_required
    def get(self):
        """Handles GET request to the resource.

        Returns:
            A view in JSON (instance of flask.wrappers.Response)
        """
        TIMEOUT_THRESHOLD_SECONDS = current_app.config.get(
            u'CELERY_TASK_TIMEOUT', 7200)
        indices = SearchIndex.query.filter(SearchIndex.status.any(
            status=u'processing')).all()
        schema = {u'objects': [], u'meta': {}}
        for search_index in indices:
            # pylint: disable=too-many-function-args
            celery_task = self.celery.AsyncResult(search_index.index_name)
            task = dict(
                task_id=celery_task.task_id, state=celery_task.state,
                successful=celery_task.successful(), name=search_index.name,
                result=False)
            if celery_task.state == u'SUCCESS':
                task[u'result'] = celery_task.result
            elif celery_task.state == u'PENDING':
                time_pending = (
                    search_index.updated_at - datetime.datetime.now())
                if time_pending.seconds > TIMEOUT_THRESHOLD_SECONDS:
                    search_index.set_status(u'timeout')
            schema[u'objects'].append(task)
        return jsonify(schema)
