# 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.
"""Tests for v1 of the Timesketch API."""

import os
import shutil
import tempfile
import json
from unittest import mock

from timesketch.lib.definitions import HTTP_STATUS_CODE_BAD_REQUEST
from timesketch.lib.definitions import HTTP_STATUS_CODE_CREATED
from timesketch.lib.definitions import HTTP_STATUS_CODE_NOT_FOUND
from timesketch.lib.definitions import HTTP_STATUS_CODE_OK
from timesketch.lib.definitions import HTTP_STATUS_CODE_FORBIDDEN
from timesketch.lib.definitions import HTTP_STATUS_CODE_GATEWAY_TIMEOUT
from timesketch.lib.errors import DatastoreTimeoutError
from timesketch.lib.testlib import BaseTest
from timesketch.lib.testlib import MockDataStore
from timesketch.lib.dfiq import DFIQCatalog
from timesketch.api.v1.resources import scenarios
from timesketch.models.sketch import Scenario
from timesketch.models.sketch import InvestigativeQuestion
from timesketch.models.sketch import InvestigativeQuestionApproach
from timesketch.models.sketch import Facet
from timesketch.models.sketch import Timeline
from timesketch.models.sketch import SearchIndex
from timesketch.models.sketch import Sketch
from timesketch.models.user import User, Group
from timesketch.models import db_session
from timesketch.api.v1.resources import ResourceMixin
from timesketch.api.v1.resources import upload


class ResourceMixinTest(BaseTest):
    """Test ResourceMixin."""

    def test_to_json_empty_list(self):
        """Behavior of to_json when given an empty list."""
        response = ResourceMixin().to_json([])
        self.assertEqual(
            response.json,
            {
                "meta": {},
                "objects": [],
            },
        )


class InvalidResourceTest(BaseTest):
    """Test an Invalid Resource."""

    invalid_resource_url = "api/v1/invalidresource"

    def test_invalid_endpoint(self):
        """Authenticated request to get a non existent API endpoint"""
        self.login()
        response = self.client.get(self.invalid_resource_url)
        self.assert404(response)


class SketchListResourceTest(BaseTest):
    """Test SketchListResource."""

    resource_url = "/api/v1/sketches/"

    def test_sketch_list_resource(self):
        """Authenticated request to get list of sketches."""
        self.login()
        response = self.client.get(self.resource_url)
        self.assertEqual(len(response.json["objects"]), 2)
        result = sorted(i["name"] for i in response.json["objects"])
        self.assertEqual(result, ["Test 1", "Test 3"])
        self.assert200(response)

    def test_sketch_post_resource(self):
        """Authenticated request to create a sketch."""
        self.login()
        data = {"name": "test", "description": "test"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

    def test_sketch_list_all_scope(self):
        """Authenticated request to get a list of all sketches for a user."""
        # 1. Login as user1 and create a new sketch.
        self.login()
        sketch_name_by_user1 = "Sketch by User1 to be shared"
        data = {"name": sketch_name_by_user1, "description": "sharing is caring"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)
        new_sketch_id = response.json["objects"][0]["id"]

        # 2. Share it with user2.
        collaborator_url = f"/api/v1/sketches/{new_sketch_id}/collaborators/"
        collaborator_data = {"users": ["test2"]}
        response = self.client.post(
            collaborator_url,
            data=json.dumps(collaborator_data),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)

        # 3. Login as user2 and create a new sketch.
        # user2 also has access to sketch1 from user1 via setUp.
        self.login(username="test2", password="test")
        sketch_name_by_user2 = "My Sketch by User2"
        data = {"name": sketch_name_by_user2, "description": "My own sketch"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

        # 4. Get all sketches for user2.
        response = self.client.get(self.resource_url, query_string={"scope": "all"})
        self.assert200(response)

        # 5. Should be 2 sketches: one owned, one shared from user1 via the
        # collaborator API. The sketch shared via a View in setUp is not
        # included in this scope.
        self.assertEqual(len(response.json["objects"]), 2)
        names = {s["name"] for s in response.json["objects"]}
        expected_names = {sketch_name_by_user1, sketch_name_by_user2}
        self.assertSetEqual(names, expected_names)

    def test_sketch_list_no_duplicates(self):
        """Test that users in multiple shared groups don't see duplicate counts."""
        # 1. Use user2
        user = self.user2

        # 2. Create two groups and add the user to both
        group1 = Group(name="repro_group1", display_name="repro_group1")
        group2 = Group(name="repro_group2", display_name="repro_group2")
        db_session.add(group1)
        db_session.add(group2)
        db_session.commit()

        user.groups.append(group1)
        user.groups.append(group2)
        db_session.add(user)
        db_session.commit()

        # 3. Create a sketch owned by user1
        sketch = self._create_sketch(name="Repro Sketch", user=self.user1)

        # 4. Share the sketch with both groups
        sketch.grant_permission(permission="read", group=group1)
        sketch.grant_permission(permission="read", group=group2)
        db_session.add(sketch)
        db_session.commit()

        # 5. Login as test2
        self.login(username="test2", password="test")

        # 6. Call the /api/v1/sketches/ endpoint
        response = self.client.get(self.resource_url, query_string={"scope": "all"})

        # 7. Check total_items
        total_items = response.json["meta"]["total_items"]
        objects = response.json["objects"]
        objects_count = len(objects)

        self.assertEqual(objects_count, 1, "Expected exactly 1 sketch in objects list")
        self.assertEqual(total_items, 1, "Expected total_items to be exactly 1")

    def test_sketch_list_user_scope(self):
        """Authenticated request to get a list of the user's own sketches."""
        # user2 has access to sketch1 from user1 via setUp.
        # We want to ensure that scope=user only returns sketches owned by user2.
        self.login(username="test2", password="test")

        sketch_name = "My Sketch Owned By User2"
        data = {"name": sketch_name, "description": "This is a test sketch for user2"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

        # Get sketches with scope=user
        response = self.client.get(self.resource_url, query_string={"scope": "user"})
        self.assert200(response)

        # The list should only contain the sketch created by user2.
        # The shared sketch ("Test 1") should not be included.
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], sketch_name)

    def test_sketch_list_pagination(self):
        """Authenticated request to test pagination of sketches."""
        self.login()
        # User1 has two sketches with ACLs: Test 1, Test 3.
        response = self.client.get(self.resource_url, query_string={"per_page": 1})
        self.assert200(response)
        meta = response.json["meta"]
        objects = response.json["objects"]

        self.assertEqual(len(objects), 1)
        name1 = objects[0]["name"]
        self.assertIn(name1, ["Test 1", "Test 3"])
        self.assertEqual(meta["total_pages"], 2)
        self.assertTrue(meta["has_next"])
        self.assertFalse(meta["has_prev"])

        # Get page 2
        response = self.client.get(
            self.resource_url, query_string={"per_page": 1, "page": 2}
        )
        self.assert200(response)
        meta = response.json["meta"]
        objects = response.json["objects"]

        self.assertEqual(len(objects), 1)
        name2 = objects[0]["name"]
        self.assertIn(name2, ["Test 1", "Test 3"])
        self.assertNotEqual(name1, name2)
        self.assertFalse(meta["has_next"])
        self.assertTrue(meta["has_prev"])


class SketchResourceTest(BaseTest):
    """Test SketchResource."""

    resource_url = "/api/v1/sketches/1/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_sketch_resource(self):
        """Authenticated request to get a sketch."""
        self.login()
        response = self.client.get(self.resource_url)
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(len(response.json["objects"][0]["timelines"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], "Test 1")
        self.assertIsInstance(response.json["meta"]["emojis"], dict)
        self.assert200(response)

    def test_sketch_acl(self):
        """
        Authenticated request to get a sketch that the user do not have read
        permission on.
        """
        self.login()
        response = self.client.get("/api/v1/sketches/2/")
        self.assert403(response)

    def test_create_a_sketch(self):
        """Authenticated request to create a sketch."""
        self.login()
        data = {"name": "test_create_a_sketch", "description": "test_create_a_sketch"}
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(HTTP_STATUS_CODE_CREATED, response.status_code)

        # check the created sketch

        response = self.client.get("/api/v1/sketches/")
        self.assertEqual(len(response.json["objects"]), 3)
        self.assertIn(b"test_create_a_sketch", response.data)
        self.assert200(response)

    def test_append_label_to_sketch(self):
        """Authenticated request to append a label to a sketch."""
        self.login()

        data = {"labels": ["test_append_label_to_sketch"], "label_action": "add"}

        response = self.client.post(
            "/api/v1/sketches/3/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )

        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

        # check the result in content
        response = self.client.get("/api/v1/sketches/3/")
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], "Test 3")
        self.assertIn(
            "test_append_label_to_sketch", response.json["objects"][0]["label_string"]
        )
        self.assert200(response)

    def test_archive_sketch(self):
        """Authenticated request to archive a sketch."""
        self.login()

        # Create sketch to test with
        data = {"name": "test_archive_sketch", "description": "test_archive_sketch"}
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        created_id = response.json["objects"][0]["id"]

        self.assertEqual(HTTP_STATUS_CODE_CREATED, response.status_code)

        # Pull sketch
        response = self.client.get(f"/api/v1/sketches/{created_id}/")
        self.assertEqual(HTTP_STATUS_CODE_OK, response.status_code)
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], "test_archive_sketch")

        # Archive sketch
        resource_url = f"/api/v1/sketches/{created_id}/archive/"
        data = {"action": "archive"}
        response = self.client.post(
            resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assert200(response)

        # Pull the sketch again to get the status
        response = self.client.get(f"/api/v1/sketches/{created_id}/")
        self.assertEqual(
            response.json["objects"][0]["name"],
            "test_archive_sketch",
        )
        self.assert200(response)
        self.assertIn("archived", response.json["objects"][0]["status"][0]["status"])

    def test_unarchive_sketch_with_failed_timeline(self):
        """Tests that a sketch can be successfully unarchived even if it contains
        a timeline in a 'fail' state. This ensures that users can recover
        sketches with problematic timelines and fix them (e.g., by deleting
        the failed timeline) after unarchiving.
        """
        self.login()

        # Create sketch to test with
        data = {"name": "test_unarchive_fail", "description": "desc"}
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        created_id = response.json["objects"][0]["id"]

        sketch = Sketch.get_by_id(created_id)

        # Create a dummy search index
        searchindex = SearchIndex(
            name="failed_index",
            description="failed_index",
            user=sketch.user,
            index_name="failed_index",
        )
        searchindex.set_status("fail")
        db_session.add(searchindex)

        timeline = Timeline(
            name="failed_timeline",
            description="failed_timeline",
            sketch=sketch,
            user=sketch.user,
            searchindex=searchindex,
        )
        timeline.set_status("fail")
        db_session.add(timeline)
        db_session.commit()

        # Manually archive the sketch (bypass API check for failed timeline
        # during archive if any). Actually, archive API might block it too
        # if we don't fix archive API. But we are testing UNARCHIVE here.
        # So let's force the state.
        sketch.set_status("archived")
        timeline.set_status("fail")
        db_session.commit()

        # Try to unarchive
        resource_url = f"/api/v1/sketches/{created_id}/archive/"
        data = {"action": "unarchive"}
        response = self.client.post(
            resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )

        # This should now succeed with 200 OK instead of 500
        self.assert200(response)

        # Verify sketch is ready
        sketch = Sketch.get_by_id(created_id)
        self.assertEqual(sketch.get_status.status, "ready")

        # Verify failed timeline is still failed (or whatever the logic does)
        self.assertEqual(timeline.get_status.status, "fail")

    def test_unarchive_sketch_with_mixed_states(self):
        """Tests that a sketch can be successfully unarchived even if it contains
        both 'fail' and 'processing' timelines. This confirms that the
        unarchive operation is robust against various non-ready timeline
        states, allowing users to regain access to the sketch and address
        each problematic timeline individually.
        """
        self.login()

        # Create sketch
        data = {"name": "test_unarchive_mixed", "description": "desc"}
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        created_id = response.json["objects"][0]["id"]

        sketch = Sketch.get_by_id(created_id)

        # Create a failed timeline
        idx_fail = SearchIndex(name="idx_fail", user=sketch.user, index_name="idx_fail")
        idx_fail.set_status("fail")
        db_session.add(idx_fail)
        tl_fail = Timeline(
            name="tl_fail",
            sketch=sketch,
            user=sketch.user,
            searchindex=idx_fail,
        )
        tl_fail.set_status("fail")
        db_session.add(tl_fail)

        # Create a processing timeline
        idx_proc = SearchIndex(name="idx_proc", user=sketch.user, index_name="idx_proc")
        idx_proc.set_status("processing")
        db_session.add(idx_proc)
        tl_proc = Timeline(
            name="tl_proc",
            sketch=sketch,
            user=sketch.user,
            searchindex=idx_proc,
        )
        tl_proc.set_status("processing")
        db_session.add(tl_proc)

        # Set sketch to archived to simulate the stuck state
        sketch.set_status("archived")
        db_session.commit()

        # Try to unarchive
        resource_url = f"/api/v1/sketches/{created_id}/archive/"
        data = {"action": "unarchive"}
        response = self.client.post(
            resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )

        # Should succeed
        self.assert200(response)

        # Verify sketch is ready
        sketch = Sketch.get_by_id(created_id)
        self.assertEqual(sketch.get_status.status, "ready")

        # Verify timelines retained their stuck states (so user can see/fix
        # them)
        self.assertEqual(tl_fail.get_status.status, "fail")
        self.assertEqual(tl_proc.get_status.status, "processing")

    def test_sketch_delete_not_existant_sketch(self):
        """Authenticated request to delete a sketch that does not exist."""
        self.login()
        response = self.client.delete("/api/v1/sketches/99/")
        self.assert404(response)

    def test_sketch_delete_no_acl(self):
        """Authenticated request to delete a sketch that the User has no read
        permission on.
        """
        self.login()
        response = self.client.delete("/api/v1/sketches/2/")
        self.assert403(response)

    def test_attempt_to_delete_protected_sketch(self):
        """Authenticated request to delete a protected sketch."""
        self.login()
        data = {
            "name": "test_attempt_to_delete_protected_sketch",
            "description": "test_attempt_to_delete_protected_sketch",
        }
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(HTTP_STATUS_CODE_CREATED, response.status_code)
        data = {"labels": ["protected"], "label_action": "add"}
        response = self.client.post(
            "/api/v1/sketches/4/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )

        self.assertEqual(
            response.json["objects"][0]["name"],
            "test_attempt_to_delete_protected_sketch",
        )
        self.assertIn("protected", response.json["objects"][0]["label_string"])

        response = self.client.delete("/api/v1/sketches/4/")
        self.assert403(response)

    def test_attempt_to_delete_archived_sketch(self):
        """Test attempting to delete an archived sketch."""
        self.login()

        # Create sketch to test with
        data = {
            "name": "test_delete_archive_sketch",
            "description": "test_delete_archive_sketch",
        }
        response = self.client.post(
            "/api/v1/sketches/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        created_id = response.json["objects"][0]["id"]

        self.assertEqual(HTTP_STATUS_CODE_CREATED, response.status_code)
        response = self.client.get(f"/api/v1/sketches/{created_id}/")
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(
            response.json["objects"][0]["name"], "test_delete_archive_sketch"
        )
        self.assertEqual(HTTP_STATUS_CODE_OK, response.status_code)

        # Archive sketch
        resource_url = f"/api/v1/sketches/{created_id}/archive/"
        data = {"action": "archive"}
        response = self.client.post(
            resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assert200(response)

        # Pull the sketch again to get the status
        response = self.client.get(f"/api/v1/sketches/{created_id}/")
        self.assertEqual(
            response.json["objects"][0]["name"],
            "test_delete_archive_sketch",
        )
        self.assert200(response)
        status_list = response.json["objects"][0]["status"]
        self.assertIn("archived", status_list[0]["status"])

        # Soft delete should fail
        response = self.client.delete(f"/api/v1/sketches/{created_id}/")
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertEqual(
            "Unable to delete a sketch that is already archived.",
            response.json["message"],
        )

        # Force delete should also fail because the sketch is archived
        # We need to give the admin permission to delete the sketch first

        sketch = Sketch.get_by_id(created_id)
        user_admin = User.query.filter_by(username="testadmin").first()

        for permission in ["read", "write", "delete"]:
            sketch.grant_permission(permission=permission, user=user_admin)

        db_session.commit()

        # Check that the admin user now has delete permission.
        self.assertTrue(sketch.has_permission(user_admin, "delete"))

        self.login_admin()
        resource_url = f"/api/v1/sketches/{created_id}/?force=true"
        response = self.client.delete(resource_url)
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertEqual(
            "Unable to delete a sketch that is already archived.",
            response.json["message"],
        )

        # Unarchive the sketch (should succeed now even if we had failed
        # timelines, though this test sketch has no timelines)
        self.login()
        resource_url = f"/api/v1/sketches/{created_id}/archive/"
        data = {"action": "unarchive"}
        response = self.client.post(
            resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assert200(response)

        # Or force delete (admin)
        self.login_admin()
        resource_url = f"/api/v1/sketches/{created_id}/?force=true"
        response = self.client.delete(resource_url)
        self.assertEqual(HTTP_STATUS_CODE_OK, response.status_code)


class ViewListResourceTest(BaseTest):
    """Test ViewListResource."""

    resource_url = "/api/v1/sketches/1/views/"

    def test_post_view_list_resource(self):
        """Authenticated request to create a view."""
        self.login()
        data = {
            "name": "test",
            "new_searchtemplate": False,
            "query": "test",
            "filter": {},
            "dsl": {},
        }
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        data["from_searchtemplate_id"] = 1
        response_with_searchtemplate = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)
        self.assertEqual(
            response_with_searchtemplate.status_code, HTTP_STATUS_CODE_CREATED
        )


class ViewResourceTest(BaseTest):
    """Test ViewResource."""

    resource_url = "/api/v1/sketches/1/views/1/"

    def test_view_resource(self):
        """Authenticated request to get a view."""
        self.login()
        response = self.client.get(self.resource_url)
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], "View 1")
        self.assert200(response)

    def test_post_view_resource(self):
        """Authenticated request to update a view."""
        self.login()
        data = {"name": "test", "query": "test", "filter": "{}"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

    def test_invalid_user_in_view(self):
        """Authenticated request to get a view for another user."""
        self.login()
        response = self.client.get("/api/v1/sketches/1/views/3/")
        self.assert403(response)

    def test_invalid_view(self):
        """Authenticated request to get a view for non existing view."""
        self.login()
        response = self.client.get("/api/v1/sketches/1/views/2/")
        self.assert404(response)


class SearchTemplateResourceTest(BaseTest):
    """Test Search template resource."""

    resource_url = "/api/v1/searchtemplates/1/"

    def test_searchtemplate_resource(self):
        """Authenticated request to get a search template."""
        self.login()
        response = self.client.get(self.resource_url)
        self.assertEqual(len(response.json["objects"]), 1)
        self.assertEqual(response.json["objects"][0]["name"], "template")
        self.assert200(response)

    def test_invalid_searchtemplate(self):
        """Authenticated request to get a non existing search template."""
        self.login()
        response = self.client.get("/api/v1/searchtemplates/2/")
        self.assert404(response)


class ExploreResourceTest(BaseTest):
    """Test ExploreResource."""

    resource_url = "/api/v1/sketches/1/explore/"
    expected_response = {
        "meta": {
            "es_time": 5,
            "es_total_count": 1,
            "es_total_count_complete": 0,
            "timeline_colors": {"test": "FFFFFF"},
            "timeline_names": {"test": "Timeline 1"},
            "count_per_index": {},
            "count_per_timeline": {},
            "count_over_time": {"data": {}, "interval": ""},
            "scroll_id": "",
            "search_node": {
                "children": [],
                "description": None,
                "id": 1,
                "labels": [],
                "parent": None,
                "query_dsl": None,
                "query_filter": "{}",
                "query_result_count": 0,
                "query_string": "test",
                "scenario": None,
                "facet": None,
                "question": None,
            },
        },
        "objects": [
            {
                "sort": [1410593223000],
                "_type": "plaso_event",
                "_source": {
                    "timestamp": 1410593222543942,
                    "message": "Test event",
                    "label": ["__ts_star"],
                    "timestamp_desc": "Content Modification Time",
                    "datetime": "2014-09-13T07:27:03+00:00",
                    "__ts_timeline_id": 1,
                    "comment": ["test"],
                },
                "_score": "null",
                "selected": False,
                "_index": "test",
                "_id": "test",
            }
        ],
    }

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_search(self):
        """Authenticated request to query the datastore."""
        self.login()
        data = {"query": "test", "filter": {}}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        response_json = response.json
        # Remove flaky properties (dynamically generated)
        del response_json["meta"]["search_node"]["created_at"]
        del response_json["meta"]["search_node"]["query_time"]
        self.assertDictEqual(response_json, self.expected_response)
        self.assert200(response)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_search_timeout(self):
        """Test search timeout handling."""
        self.login()
        data = {"query": "test", "filter": {}}

        # Mock the search method to raise DatastoreTimeoutError
        with mock.patch.object(
            MockDataStore, "search", side_effect=DatastoreTimeoutError("Timeout")
        ):
            response = self.client.post(
                self.resource_url,
                data=json.dumps(data, ensure_ascii=False),
                content_type="application/json",
            )
            self.assertEqual(response.status_code, HTTP_STATUS_CODE_GATEWAY_TIMEOUT)
            self.assertIn("Timeout", response.json["message"])


class AggregationExploreResourceTest(BaseTest):
    """Test AggregationExploreResource."""

    resource_url = "/api/v1/sketches/1/aggregation/explore/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_heatmap_aggregation(self):
        """Authenticated request to get aggregation requests."""
        self.login()
        data = {"aggregation_dsl": "test"}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assert200(response)


class EventResourceTest(BaseTest):
    """Test EventResource."""

    resource_url = "/api/v1/sketches/1/event/"
    maxDiff = None
    expected_response = {
        "objects": {
            "timestamp_desc": "",
            "_id": "adc123",
            "_index": [],
            "timestamp": 1410895419859714,
            "label": "",
            "source_long": "",
            "source_short": "",
            "es_index": "",
            "es_id": "",
            "message": "",
            "datetime": "2014-09-16T19:23:40+00:00",
            "__ts_timeline_id": 1,
            "comment": ["test"],
        }
    }

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_get_event(self):
        """Authenticated request to get an event from the datastore."""
        self.login()
        response = self.client.get(
            self.resource_url + "?searchindex_id=test&event_id=test"
        )
        response_json = response.json
        del response_json["meta"]
        print(f"response_json: {response_json}")
        print(f"self.expected_response: {self.expected_response}")
        self.assertEqual(response.json, response.json | self.expected_response)
        self.assert200(response)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_invalid_index(self):
        """
        Authenticated request to get an event from the datastore, but in the
        wrong index.
        """
        self.login()
        response_400 = self.client.get(
            self.resource_url + "?searchindex_id=wrong_index&event_id=test"
        )
        self.assert400(response_400)


class EventAddAttributeResourceTest(BaseTest):
    """Test EventAddAttributeResource."""

    resource_url = "/api/v1/sketches/1/event/attributes/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_add_attributes(self):
        """Test add attributes with a well formed request."""
        self.login()

        attrs = [{"attr_name": "foo", "attr_value": "bar"}]
        events = {
            "events": [
                {
                    "_id": "1",
                    "_type": "_doc",
                    "_index": self.searchindex.index_name,
                    "attributes": attrs,
                },
                {
                    "_id": "2",
                    "_type": "_doc",
                    "_index": self.searchindex.index_name,
                    "attributes": attrs,
                },
            ]
        }

        expected_response = {
            "meta": {
                "attributes_added": 2,
                "chunks_per_index": {self.searchindex.index_name: 1},
                "error_count": 0,
                "last_10_errors": [],
                "events_modified": 2,
            },
            "objects": [],
        }

        response = self.client.post(self.resource_url, json=events)
        self.assertEqual(expected_response, response.json)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_incorrect_content_type(self):
        """Test that a content-type other than application/json is handled."""
        self.login()
        response = self.client.post(self.resource_url)
        self.assertEqual(400, response.status_code)
        self.assertIn(b"Request must be in JSON format.", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_missing_events(self):
        """Test that a request without events is handled."""
        self.login()

        response = self.client.post(self.resource_url, json={"events": []})
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Request must contain an events field.", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_events_type(self):
        """Test that that the wrong type for events is handled."""
        self.login()

        response = self.client.post(self.resource_url, json={"events": "a string"})
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Events field must be a list.", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_max_events(self):
        """Test that an event list larger than max events is handled."""
        self.login()

        response = self.client.post(self.resource_url, json={"events": ["a"] * 100001})
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Request exceeds maximum events", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_missing_id(self):
        """Test that an event without an _id field is handled."""
        self.login()

        response = self.client.post(self.resource_url, json={"events": [{}]})
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Event missing field _id.", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_attributes_type(self):
        """Test that event attributes of the wrong type is handled."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_index": "1",
                        "_type": "_doc",
                        "attributes": "a string",
                    }
                ]
            },
        )
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Attributes must be a list.", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_max_attributes(self):
        """Test that too many attributes is handled."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_index": "1",
                        "_type": "_doc",
                        "attributes": ["a"] * 11,
                    }
                ]
            },
        )
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Attributes for event exceeds maximum", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_events_json_attribute_fields(self):
        """Test that an attribute with a missing field is handled."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_index": "1",
                        "_type": "_doc",
                        "attributes": [{"attr_name": "abc"}],
                    }
                ]
            },
        )
        self.assertEqual(HTTP_STATUS_CODE_BAD_REQUEST, response.status_code)
        self.assertIn(b"Attribute missing field", response.data)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_chunk_calculation(self):
        """Tests that chunks are properly calculated."""
        self.login()

        attrs = [{"attr_name": "foo", "attr_value": "bar"}]

        # One chunk when event count is the same as chunk size.
        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": attrs,
                    }
                ]
                * 1000
            },
        )
        self.assertEqual(
            {self.searchindex.index_name: 1}, response.json["meta"]["chunks_per_index"]
        )

        # Two chunks when event count is chunk size + 1.
        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": attrs,
                    }
                ]
                * 1001
            },
        )
        self.assertEqual(
            {self.searchindex.index_name: 2}, response.json["meta"]["chunks_per_index"]
        )

        # Chunk per index with multiple indexes.
        self._create_timeline(
            name="Timeline 2",
            sketch=self.sketch1,
            searchindex=self.searchindex2,
            user=self.user1,
        )

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": attrs,
                    },
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex2.index_name,
                        "attributes": attrs,
                    },
                ]
            },
        )
        self.assertEqual(
            {self.searchindex.index_name: 1, self.searchindex2.index_name: 1},
            response.json["meta"]["chunks_per_index"],
        )

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_add_existing_attributes(self):
        """Tests existing attributes cannot be overridden."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": [{"attr_name": "exists", "attr_value": "yes"}],
                    }
                ]
            },
        )
        self.assertIn(
            "Attribute 'exists' already exists for event_id '1'.",
            response.json["meta"]["last_10_errors"],
        )

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_add_invalid_attributes_underscore(self):
        """Tests attributes beginning with an underscore cannot be added."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": [{"attr_name": "_invalid", "attr_value": "yes"}],
                    }
                ]
            },
        )
        self.assertIn(
            "Attribute '_invalid' for event_id '1' invalid, cannot start with '_'",
            response.json["meta"]["last_10_errors"],
        )

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_add_invalid_attributes_disallowed_name(self):
        """Tests attributes cannot be added with a disallowed name."""
        self.login()

        response = self.client.post(
            self.resource_url,
            json={
                "events": [
                    {
                        "_id": "1",
                        "_type": "_doc",
                        "_index": self.searchindex.index_name,
                        "attributes": [{"attr_name": "message", "attr_value": "yes"}],
                    }
                ]
            },
        )
        self.assertIn(
            "Cannot add 'message' for event_id '1', name not allowed.",
            response.json["meta"]["last_10_errors"],
        )

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_add_attributes_forbidden(self):
        """Test that adding attributes to an event in an index that does not
        belong to the sketch returns a 403 Forbidden error. This is a security
        check to prevent modifying events in unauthorized indices.
        """
        self.login()

        # searchindex2 is NOT associated with sketch 1 by default.
        events = {
            "events": [
                {
                    "_id": "1",
                    "_type": "_doc",
                    "_index": self.searchindex2.index_name,
                    "attributes": [{"attr_name": "foo", "attr_value": "bar"}],
                }
            ]
        }

        response = self.client.post(self.resource_url, json=events)
        self.assertEqual(HTTP_STATUS_CODE_FORBIDDEN, response.status_code)
        self.assertIn(b"does not belong to the sketch", response.data)


class EventTaggingResourceTest(BaseTest):
    """Test EventTaggingResource."""

    resource_url = "/api/v1/sketches/1/event/tagging/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_tagging(self):
        """Test that tagging an event with a valid request succeeds with an
        HTTP 200 OK response and the correct number of tags applied.
        """
        self.login()

        data = {
            "tag_string": '["foo", "bar"]',
            "events": [{"_id": "1", "_index": self.searchindex.index_name}],
        }
        response = self.client.post(self.resource_url, json=data)
        self.assertEqual(HTTP_STATUS_CODE_OK, response.status_code)
        self.assertEqual(response.json["meta"]["tags_applied"], 2)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_tagging_forbidden(self):
        """Test that attempting to tag an event in an index not associated with
        the current sketch returns a 403 Forbidden error.
        """
        self.login()

        data = {
            "tag_string": '["foo"]',
            "events": [{"_id": "1", "_index": self.searchindex2.index_name}],
        }
        response = self.client.post(self.resource_url, json=data)
        self.assertEqual(HTTP_STATUS_CODE_FORBIDDEN, response.status_code)
        self.assertIn(b"does not belong to the sketch", response.data)


class EventUnTagResourceTest(BaseTest):
    """Test EventUnTagResource."""

    resource_url = "/api/v1/sketches/1/event/untag/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_untagging(self):
        """Test untagging events."""
        self.login()

        data = {
            "tags_to_remove": ["foo"],
            "events": [{"_id": "1", "_index": self.searchindex.index_name}],
        }
        response = self.client.post(self.resource_url, json=data)
        self.assertEqual(HTTP_STATUS_CODE_OK, response.status_code)

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_untagging_forbidden(self):
        """Test that untagging an event in an index that does not belong to the
        sketch returns 403 Forbidden.
        """
        self.login()

        data = {
            "tags_to_remove": ["foo"],
            "events": [{"_id": "1", "_index": self.searchindex2.index_name}],
        }
        response = self.client.post(self.resource_url, json=data)
        self.assertEqual(HTTP_STATUS_CODE_FORBIDDEN, response.status_code)
        self.assertIn(b"does not belong to the sketch", response.data)


class EventAnnotationResourceTest(BaseTest):
    """Test EventAnnotationResource."""

    resource_url = "/api/v1/sketches/1/event/annotate/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_post_annotate_resource(self):
        """Authenticated request to create an annotation."""
        self.login()
        for annotation_type in ["comment", "label"]:
            event = {"_type": "test_event", "_index": "test", "_id": "test"}
            data = {
                "annotation": "test",
                "annotation_type": annotation_type,
                "events": [event],
            }
            response = self.client.post(
                self.resource_url,
                data=json.dumps(data),
                content_type="application/json",
            )
            self.assertIsInstance(response.json, dict)
            self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)

    def test_post_annotate_invalid_index_resource(self):
        """
        Authenticated request to create an annotation, but in the wrong index.
        """
        self.login()
        data = {
            "annotation": "test",
            "annotation_type": "comment",
            "event_id": "test",
            "searchindex_id": "invalid_searchindex",
        }
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)


class SearchIndexResourceTest(BaseTest):
    """Test SearchIndexResource."""

    resource_url = "/api/v1/searchindices/"

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_post_create_searchindex(self):
        """Authenticated request to create a searchindex."""
        self.login()
        data = {"searchindex_name": "test3", "es_index_name": "test3", "public": False}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertIsInstance(response.json, dict)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)


class TimelineListResourceTest(BaseTest):
    """Test TimelineList resource."""

    resource_url = "/api/v1/sketches/1/timelines/"

    def test_add_existing_timeline_resource(self):
        """Authenticated request to add a timeline to a sketch."""
        self.login()
        data = {"timeline": 1}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)

    def test_add_new_timeline_resource(self):
        """Authenticated request to add a timeline to a sketch."""
        self.login()
        data = {"timeline": 2}
        response = self.client.post(
            self.resource_url,
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)


class SigmaRuleResourceTest(BaseTest):
    """Test Sigma Rule resource."""

    expected_response = {
        "objects": {
            "description": "Detects suspicious installation of bbbbbb",
            "id": "5266a592-b793-11ea-b3de-bbbbbb",
            "level": "high",
            "logsource": {"product": "linux", "service": "shell"},
            "title": "Suspicious Installation of bbbbbb",
        }
    }

    @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
    def test_post_sigma_resource(self):
        """Authenticated request to POST an sigma rule."""
        MOCK_SIGMA_RULE = """
title: Suspicious Installation of bbbbbb
id: 5266a592-b793-11ea-b3de-bbbbbb
description: Detects suspicious installation of bbbbbb
references:
    - https://rmusser.net/docs/ATT&CK-Stuff/ATT&CK/Discovery.html
author: Alexander Jaeger
date: 2020/06/26
modified: 2022/06/12
logsource:
    product: linux
    service: shell
detection:
    keywords:
        # Generic suspicious commands
        - '*apt-get install bbbbbb*'
    condition: keywords
falsepositives:
    - Unknown
level: high
"""

        self.login()

        sigma = {
            "rule_uuid": "5266a592-b793-11ea-b3de-bbbbbb",
            "title": "Suspicious Installation of bbbbbb",
            "description": "Detects suspicious installation of bbbbbb",
            "rule_yaml": MOCK_SIGMA_RULE,
        }

        # Create a first rule
        response = self.client.post(
            "/api/v1/sigmarules/",
            data=json.dumps(sigma),
            content_type="application/json",
        )
        self.assertIn("bbbbbb", response.json["objects"][0]["rule_uuid"])
        self.assertIn(
            "bbbbbb",
            response.json["objects"][0]["description"],
        )
        self.assertIn(
            "shell",
            response.json["objects"][0]["rule_yaml"],
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)
        # Now GET the resources
        response = self.client.get("/api/v1/sigmarules/5266a592-b793-11ea-b3de-bbbbbb/")

        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        self.assertIn("bbbbbb", response.json["objects"][0]["rule_uuid"])
        self.assertIn(
            "bbbbbb",
            response.json["objects"][0]["description"],
        )
        self.assertIn(
            "shell",
            response.json["objects"][0]["rule_yaml"],
        )

    def test_get_sigma_rule(self):
        """Authenticated request to get an sigma rule."""
        self.login()
        response = self.client.get(
            "/api/v1/sigmarules/5266a592-b793-11ea-b3de-0242ac130004"
        )
        self.assertIsNotNone(response)

    def test_get_sigma_rule_that_does_not_exist(self):
        """Fetch a Sigma rule that does not exist."""
        self.login()
        response = self.client.get("/api/v1/sigmarules/foobar/")
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_NOT_FOUND)

    def test_put_sigma_rule(self):
        """Authenticated request to update sigma rule."""
        self.login()

        sigma = {
            "rule_uuid": "5266a592-b793-11ea-b3de-bbbbbb",
            "title": "Suspicious Installation of bbbbbb",
            "description": "Detects suspicious installation of bbbbbb",
            "rule_yaml": """
title: Suspicious Installation of bbbbbb
id: 5266a592-b793-11ea-b3de-bbbbbb
description: Detects suspicious installation of bbbbbb
references:
    - https://rmusser.net/docs/ATT&CK-Stuff/ATT&CK/Discovery.html
author: Alexander Jaeger
date: 2020/06/26
modified: 2022/06/12
logsource:
    product: linux
    service: shell
detection:
    keywords:
        # Generic suspicious commands
        - '*apt-get install bbbbbb*'
    condition: keywords
falsepositives:
    - Unknown
level: high
""",
        }

        # Create a first rule
        response = self.client.post(
            "/api/v1/sigmarules/",
            data=json.dumps(sigma),
            content_type="application/json",
        )
        self.assertIn("bbbbbb", response.json["objects"][0]["rule_uuid"])
        self.assertIn(
            "bbbbbb",
            response.json["objects"][0]["description"],
        )
        self.assertIn(
            "shell",
            response.json["objects"][0]["rule_yaml"],
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_CREATED)
        response = self.client.put(
            "/api/v1/sigmarules/5266a592-b793-11ea-b3de-bbbbbb/",
            data=json.dumps(
                {
                    "rule_uuid": "5266a592-b793-11ea-b3de-bbbbbb",
                    "title": "Suspicious Installation of cccccc",
                    "description": "Detects suspicious installation of cccccc",
                    "rule_yaml": """
title: Suspicious Installation of cccccc
id: 5266a592-b793-11ea-b3de-bbbbbb
description: Detects suspicious installation of cccccc
references:
    - https://rmusser.net/docs/ATT&CK-Stuff/ATT&CK/Discovery.html
author: Alexander Jaeger
date: 2020/06/26
modified: 2022/06/12
logsource:
    product: linux
    service: shell
detection:
    keywords:
        # Generic suspicious commands
        - '*apt-get install cccccc*'
    condition: keywords
falsepositives:
    - Unknown
level: high
""",
                }
            ),
            content_type="application/json",
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        self.assertIn("bbbbbb", response.json["objects"][0]["rule_uuid"])
        self.assertIn(
            "cccccc",
            response.json["objects"][0]["rule_yaml"],
        )

        self.assertIn(
            "cccccc",
            response.json["objects"][0]["description"],
        )
        self.assertIn(
            "cccccc",
            response.json["objects"][0]["title"],
        )


class SigmaRuleListResourceTest(BaseTest):
    """Test Sigma resource."""

    def test_get_sigma_rule_list(self):
        self.login()
        response = self.client.get("/api/v1/sigmarules/")
        self.assertIsNotNone(response)
        self.assertEqual(
            len(response.json["objects"]), response.json["meta"]["rules_count"]
        )
        rule = response.json["objects"][0]
        self.assertIn("5266a592-b793-11ea-b3de-0242ac", rule["rule_uuid"])
        self.assertIsNotNone(rule["created_at"])


class SigmaRuleByTextResourceTest(BaseTest):
    """Test SigmaRule by text resource."""

    correct_rule = """
        title: Installation of foobar
        id: bb1e0d1d-cd13-4b65-bf7e-69b4e740266b
        description: Detects suspicious installation of foobar
        references:
            - https://sample.com/foobar
        author: Alexander Jaeger
        date: 2020/12/10
        modified: 2020/12/10
        tags:
            - attack.discovery
            - attack.t1046
        logsource:
            product: linux
            service: shell
        detection:
            keywords:
                # Generic suspicious commands
                - '*apt-get install foobar*'
            condition: keywords
        falsepositives:
            - Unknown
        level: high
        """
    expected_response = {
        "meta": {"parsed": True},
        "objects": [
            {
                "title": "Installation of foobar",
                "id": "bb1e0d1d-cd13-4b65-bf7e-69b4e740266b",
                "description": "Detects suspicious installation of foobar",
                "references": ["https://sample.com/foobar"],
                "author": "Alexander Jaeger",
                "date": "2020/12/10",
                "modified": "2020/12/10",
                "tags": ["attack.discovery", "attack.t1046"],
                "logsource": {"product": "linux", "service": "shell"},
                "detection": {
                    "keywords": ["*apt-get install foobar*"],
                    "condition": "keywords",
                },
                "falsepositives": ["Unknown"],
                "level": "high",
                "search_query": '(data_type:("shell:zsh:history" OR "bash:history:command" OR "apt:history:line" OR "selinux:line") AND "apt-get install foobar")',  # pylint: disable=line-too-long
                "file_name": "N/A",
            }
        ],
    }

    def test_get_sigma_rule(self):
        """Authenticated request to get an sigma rule by text."""
        self.login()

        data = {"content": self.correct_rule}
        response = self.client.post(
            "/api/v1/sigmarules/text/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        self.assertEqual(response.json, response.json | self.expected_response)
        self.assert200(response)

    def test_get_non_existing_rule_by_text(self):
        """Authenticated request to get an sigma rule by text with non parseable
        yaml text."""
        self.login()
        data = {"content": "foobar: asd"}
        response = self.client.post(
            "/api/v1/sigmarules/text/",
            data=json.dumps(data, ensure_ascii=False),
            content_type="application/json",
        )
        data = json.loads(response.get_data(as_text=True))

        self.assertIn("Sigma parsing error generating rule", data["message"])
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)

    def test_get_rule_by_text_no_form_data(self):
        """Authenticated request to get an sigma rule by text with no form
        data"""
        self.login()
        response = self.client.post(
            "/api/v1/sigmarules/text/",
            data=json.dumps({"action": "post"}),
            content_type="application/json",
        )
        data = json.loads(response.get_data(as_text=True))

        self.assertIn("Missing value in the request", data["message"])
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)


class IntelligenceResourceTest(BaseTest):
    """Test Intelligence resource."""

    def test_get_intelligence_tag_metadata(self):
        """Authenticated request to get intelligence tag metadata."""
        expected_tag_metadata = {
            "malware": {"weight": 100, "type": "danger"},
            "bad": {"weight": 90, "type": "danger"},
            "suspicious": {"weight": 50, "type": "warning"},
            "good": {"weight": 10, "type": "legit"},
            "legit": {"weight": 10, "type": "legit"},
            "default": {"weight": 0, "type": "default"},
            "export": {"weight": 100, "type": "info"},
            "regexes": {
                "^GROUPNAME": {"type": "danger", "weight": 100},
                "^inv_": {"type": "warning", "weight": 80},
            },
        }
        self.login()
        response = self.client.get("/api/v1/intelligence/tagmetadata/")
        data = json.loads(response.get_data(as_text=True))
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        self.assertDictEqual(data, expected_tag_metadata)


class ContextLinksResourceTest(BaseTest):
    """Test Context Links resources."""

    maxDiff = None

    def test_get_context_links_config(self):
        """Authenticated request to get the context links configuration."""

        expected_configuration = {
            "hash": [
                {
                    "context_link": "https://lookupone.local/q=<ATTR_VALUE>",
                    "redirect_warning": True,
                    "short_name": "LookupOne",
                    "type": "linked_services",
                    "validation_regex": "/^[0-9a-f]{40}$|^[0-9a-f]{32}$/i",
                },
                {
                    "context_link": "https://lookuptwo.local/q=<ATTR_VALUE>",
                    "redirect_warning": False,
                    "short_name": "LookupTwo",
                    "type": "linked_services",
                    "validation_regex": "/^[0-9a-f]{64}$/i",
                },
            ],
            "original_url": [
                {
                    "module": "module_two",
                    "short_name": "ModuleTwo",
                    "type": "hardcoded_modules",
                }
            ],
            "sha256_hash": [
                {
                    "context_link": "https://lookuptwo.local/q=<ATTR_VALUE>",
                    "redirect_warning": False,
                    "short_name": "LookupTwo",
                    "type": "linked_services",
                    "validation_regex": "/^[0-9a-f]{64}$/i",
                }
            ],
            "uri": [
                {
                    "module": "module_two",
                    "short_name": "ModuleTwo",
                    "type": "hardcoded_modules",
                }
            ],
            "url": [
                {
                    "module": "module_two",
                    "short_name": "ModuleTwo",
                    "type": "hardcoded_modules",
                },
                {
                    "context_link": "https://lookupthree.local/q=<ATTR_VALUE>",
                    "redirect_warning": True,
                    "short_name": "LookupThree",
                    "type": "linked_services",
                },
            ],
            "xml": [
                {
                    "module": "module_one",
                    "short_name": "ModuleOne",
                    "type": "hardcoded_modules",
                    "validation_regex": "/^[0-9a-f]{64}$/i",
                }
            ],
            "xml_string": [
                {
                    "module": "module_one",
                    "short_name": "ModuleOne",
                    "type": "hardcoded_modules",
                    "validation_regex": "/^[0-9a-f]{64}$/i",
                }
            ],
        }

        self.login()
        response = self.client.get("/api/v1/contextlinks/")
        data = json.loads(response.get_data(as_text=True))
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        self.assertDictEqual(data, expected_configuration)


class UserListTest(BaseTest):
    """Test UserListResource."""

    def test_user_post_resource_admin(self):
        """Authenticated request (admin user) to create another user."""
        self.login_admin()

        data = {"username": "testuser", "password": "testpassword"}
        response = self.client.post(
            "/api/v1/users/",
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertIsNotNone(response)

    def test_user_post_resource_without_admin(self):
        """Authenticated request (no admin) to create another user,
        which should not work."""
        self.login()

        data = {"username": "testuser", "password": "testpassword"}
        response = self.client.post(
            "/api/v1/users/",
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_FORBIDDEN)

    def test_user_post_resource_missing_username(self):
        """Authenticated request (admin user) to create another user,
        but with missing username, which should not work."""
        self.login_admin()

        data = {"username": "", "password": "testpassword"}
        response = self.client.post(
            "/api/v1/users/",
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_NOT_FOUND)

    def test_user_post_resource_missing_password(self):
        """Authenticated request (admin user) to create another user,
        but with missing password, which should not work."""
        self.login_admin()

        data = {"username": "testuser", "password": ""}
        response = self.client.post(
            "/api/v1/users/",
            data=json.dumps(data),
            content_type="application/json",
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_NOT_FOUND)


class UserTest(BaseTest):
    """Test UserResource."""

    def test_user_get_resource_admin(self):
        """Authenticated request (admin user) to create another user."""
        self.login_admin()

        response = self.client.get("/api/v1/users/1/")
        data = json.loads(response.get_data(as_text=True))
        self.assertEqual(data["objects"][0]["username"], "test1")


class SystemSettingsResourceTest(BaseTest):
    """Test system settings resource."""

    resource_url = "/api/v1/settings/"

    def test_system_settings_resource(self):
        """Authenticated request to get system settings."""
        self.app.config["LLM_PROVIDER_CONFIGS"] = {"default": {"test": {}}}
        self.app.config["DFIQ_ENABLED"] = False
        self.app.config["LLM_LOG_ANALYZER_DEFAULT_PROMPT"] = "test prompt"
        self.login()
        response = self.client.get(self.resource_url)

        self.assertEqual(response.json["DFIQ_ENABLED"], False)
        self.assertEqual(response.json["SEARCH_PROCESSING_TIMELINES"], False)

        self.assertIn("LLM_FEATURES_AVAILABLE", response.json)
        self.assertIn("default", response.json["LLM_FEATURES_AVAILABLE"])
        self.assertEqual(response.json["LOG_ANALYZER_DEFAULT_PROMPT"], "test prompt")

    def test_system_settings_invalid_llm_config(self):
        """Test with invalid LLM configuration."""
        self.app.config["LLM_PROVIDER_CONFIGS"] = "invalid_config"
        self.app.config["LLM_LOG_ANALYZER_DEFAULT_PROMPT"] = "another prompt"
        self.login()
        response = self.client.get(self.resource_url)

        expected_response = {
            "DFIQ_ENABLED": False,
            "SEARCH_PROCESSING_TIMELINES": False,
            "ENABLE_V3_INVESTIGATION_VIEW": False,
            "LLM_FEATURES_AVAILABLE": {"default": False},
            "LOG_ANALYZER_DEFAULT_PROMPT": "another prompt",
        }

        self.assertDictEqual(response.json, expected_response)


class ScenariosResourceTest(BaseTest):
    """Tests the scenarios resource."""

    @mock.patch("timesketch.lib.analyzers.dfiq_plugins.manager.DFIQAnalyzerManager")
    def test_check_and_run_dfiq_analysis_steps(self, mock_analyzer_manager):
        """Test triggering analyzers for different DFIQ objects."""
        test_sketch = self.sketch1
        test_user = self.user1
        self.sketch1.set_status("ready")
        self._commit_to_database(test_sketch)

        # Load DFIQ objects
        dfiq_obj = DFIQCatalog("./tests/test_data/dfiq/")

        scenario = dfiq_obj.scenarios[0]
        scenario_sql = Scenario(
            dfiq_identifier=scenario.id,
            uuid=scenario.uuid,
            name=scenario.name,
            display_name=scenario.name,
            description=scenario.description,
            spec_json=scenario.to_json(),
            sketch=test_sketch,
            user=test_user,
        )

        facet = dfiq_obj.facets[0]
        facet_sql = Facet(
            dfiq_identifier=facet.id,
            uuid=facet.uuid,
            name=facet.name,
            display_name=facet.name,
            description=facet.description,
            spec_json=facet.to_json(),
            sketch=test_sketch,
            user=test_user,
        )
        scenario_sql.facets = [facet_sql]

        question = dfiq_obj.questions[0]
        question_sql = InvestigativeQuestion(
            dfiq_identifier=question.id,
            uuid=question.uuid,
            name=question.name,
            display_name=question.name,
            description=question.description,
            spec_json=question.to_json(),
            sketch=test_sketch,
            scenario=scenario_sql,
            user=test_user,
        )
        facet_sql.questions = [question_sql]

        approach = question.approaches[0]
        approach_sql = InvestigativeQuestionApproach(
            name=approach.name,
            display_name=approach.name,
            description=approach.description,
            spec_json=approach.to_json(),
            user=test_user,
        )
        question_sql.approaches = [approach_sql]

        self._commit_to_database(approach_sql)
        self._commit_to_database(question_sql)
        self._commit_to_database(facet_sql)
        self._commit_to_database(scenario_sql)

        # Test without analysis step
        result = scenarios.check_and_run_dfiq_analysis_steps(scenario_sql, test_sketch)
        self.assertFalse(result)

        result = scenarios.check_and_run_dfiq_analysis_steps(facet_sql, test_sketch)
        self.assertFalse(result)

        result = scenarios.check_and_run_dfiq_analysis_steps(approach_sql, test_sketch)
        self.assertFalse(result)

        # Add analysis step to approach
        approach.steps.append(
            {
                "stage": "analysis",
                "type": "timesketch-analyzer",
                "value": "test_analyzer",
            }
        )
        approach_sql.spec_json = approach.to_json()

        # Mocking analyzer manager response.
        mock_analyzer_manager.trigger_analyzers_for_approach.return_value = [
            mock.MagicMock()
        ]

        # Test with analysis step
        result = scenarios.check_and_run_dfiq_analysis_steps(
            scenario_sql, test_sketch, mock_analyzer_manager
        )
        self.assertEqual(result, [mock.ANY, mock.ANY])
        mock_analyzer_manager.trigger_analyzers_for_approach.assert_called_with(
            approach=approach_sql
        )

        result = scenarios.check_and_run_dfiq_analysis_steps(
            facet_sql, test_sketch, mock_analyzer_manager
        )
        self.assertEqual(result, [mock.ANY])
        mock_analyzer_manager.trigger_analyzers_for_approach.assert_called_with(
            approach=approach_sql
        )

        result = scenarios.check_and_run_dfiq_analysis_steps(
            question_sql, test_sketch, mock_analyzer_manager
        )
        self.assertEqual(result, [mock.ANY])
        mock_analyzer_manager.trigger_analyzers_for_approach.assert_called_with(
            approach=approach_sql
        )

        # Test with invalid object
        result = scenarios.check_and_run_dfiq_analysis_steps("invalid", test_sketch)
        self.assertFalse(result)


@mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
class LLMResourceTest(BaseTest):
    """Test LLMResource."""

    resource_url = "/api/v1/sketches/1/llm/"

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    @mock.patch(
        "timesketch.lib.llms.features.manager.FeatureManager.get_feature_instance"
    )
    @mock.patch("timesketch.lib.utils.get_validated_indices")
    @mock.patch("timesketch.api.v1.resources.llm.LLMResource._execute_llm_call")
    @mock.patch("timesketch.lib.llms.providers.manager.LLMManager.create_provider")
    def test_post_success(
        self,
        mock_create_provider,
        mock_execute_llm,
        mock_get_validated_indices,
        mock_get_feature,
        mock_get_with_acl,
    ):
        """Test a successful POST request to the LLM endpoint."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_sketch.id = 1
        mock_get_with_acl.return_value = mock_sketch

        mock_feature = mock.MagicMock()
        mock_feature.NAME = "test_feature"
        # Force legacy workflow path which this test is written for
        del mock_feature.execute
        mock_feature.generate_prompt.return_value = "test prompt"
        mock_feature.process_response.return_value = {"result": "test result"}
        mock_get_feature.return_value = mock_feature

        mock_get_validated_indices.return_value = (["index1"], [1])
        mock_execute_llm.return_value = {"response": "mock response"}
        mock_create_provider.return_value = mock.MagicMock()

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "test_feature", "filter": {}}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertEqual(response_data, {"result": "test result"})

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    def test_post_missing_data(self, mock_get_with_acl):
        """Test POST request with missing data."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_get_with_acl.return_value = mock_sketch

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"some_param": "some_value"}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("The 'feature' parameter is required", response_data["message"])

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    def test_post_missing_feature(self, mock_get_with_acl):
        """Test POST request with no feature parameter."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_get_with_acl.return_value = mock_sketch

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"filter": {}}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("The 'feature' parameter is required", response_data["message"])

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    def test_post_invalid_sketch(self, mock_get_with_acl):
        """Test POST request with an invalid sketch ID."""
        mock_get_with_acl.return_value = None

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "test_feature", "filter": {}}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_NOT_FOUND)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("No sketch found with this ID", response_data["message"])

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    def test_post_no_permission(self, mock_get_with_acl):
        """Test POST request when user lacks write permission."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = False
        mock_get_with_acl.return_value = mock_sketch

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "test_feature", "filter": {}}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_FORBIDDEN)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn(
            "User does not have sufficient access to modify the sketch.",
            response_data["message"],
        )

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    @mock.patch(
        "timesketch.lib.llms.features.manager.FeatureManager.get_feature_instance"
    )
    def test_post_invalid_feature(self, mock_get_feature, mock_get_with_acl):
        """Test POST request with an invalid feature name."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_get_with_acl.return_value = mock_sketch

        mock_get_feature.side_effect = KeyError("Invalid feature")

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "invalid_feature", "filter": {}}),
            content_type="application/json",
        )
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("Invalid LLM feature: invalid_feature", response_data["message"])

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    @mock.patch(
        "timesketch.lib.llms.features.manager.FeatureManager.get_feature_instance"
    )
    @mock.patch("timesketch.lib.utils.get_validated_indices")
    @mock.patch("timesketch.lib.llms.providers.manager.LLMManager.create_provider")
    def test_post_prompt_generation_error(
        self,
        mock_create_provider,
        mock_get_validated_indices,
        mock_get_feature,
        mock_get_with_acl,
    ):
        """Test handling of errors during prompt generation."""
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_sketch.id = 1
        mock_get_with_acl.return_value = mock_sketch

        mock_feature = mock.MagicMock()
        mock_feature.NAME = "test_feature"
        # Force legacy workflow path
        del mock_feature.execute
        mock_feature.generate_prompt.side_effect = ValueError(
            "Prompt generation failed"
        )
        mock_get_feature.return_value = mock_feature

        mock_get_validated_indices.return_value = (["index1"], [1])
        mock_create_provider.return_value = mock.MagicMock()

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "test_feature", "filter": {}}),
            content_type="application/json",
        )

        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("Prompt generation failed", response_data["message"])

        mock_feature.generate_prompt.assert_called_once()

    @mock.patch("timesketch.models.sketch.Sketch.get_with_acl")
    @mock.patch(
        "timesketch.lib.llms.features.manager.FeatureManager.get_feature_instance"
    )
    @mock.patch("timesketch.lib.utils.get_validated_indices")
    @mock.patch("multiprocessing.Manager")
    @mock.patch("multiprocessing.Process")
    @mock.patch("timesketch.lib.llms.providers.manager.LLMManager.create_provider")
    def test_post_llm_execution_timeout(
        self,
        mock_create_provider,
        mock_process,
        mock_manager,
        mock_get_validated_indices,
        mock_get_feature,
        mock_get_with_acl,
    ):
        """Test handling of LLM execution timeouts."""
        # Setup mocks
        mock_sketch = mock.MagicMock()
        mock_sketch.has_permission.return_value = True
        mock_sketch.id = 1
        mock_get_with_acl.return_value = mock_sketch

        # Setup Manager Mock
        mock_manager_instance = mock.MagicMock()
        mock_manager.return_value.__enter__.return_value = mock_manager_instance
        mock_manager_instance.dict.return_value = {}

        mock_feature = mock.MagicMock()
        mock_feature.NAME = "test_feature"
        del mock_feature.execute
        mock_feature.generate_prompt.return_value = "test prompt"
        mock_get_feature.return_value = mock_feature

        mock_get_validated_indices.return_value = (["index1"], [1])
        mock_create_provider.return_value = mock.MagicMock()

        process_instance = mock.MagicMock()
        process_instance.is_alive.return_value = True
        mock_process.return_value = process_instance

        self.login()
        response = self.client.post(
            self.resource_url,
            data=json.dumps({"feature": "test_feature", "filter": {}}),
            content_type="application/json",
        )

        self.assertEqual(response.status_code, HTTP_STATUS_CODE_GATEWAY_TIMEOUT)
        response_data = json.loads(response.get_data(as_text=True))
        self.assertIn("LLM call timed out", response_data["message"])

        process_instance.terminate.assert_called_once()


class ExportStreamListResourceTest(BaseTest):
    """Tests for the ExportStreamListResource."""

    def setUp(self):
        super().setUp()
        self.resource_url = "/api/v1/sketches/1/exportstream/"

    @mock.patch("timesketch.api.v1.resources.exportstream.OpenSearchDataStore")
    @mock.patch("timesketch.api.v1.resources.exportstream.utils.validate_indices")
    @mock.patch("timesketch.api.v1.resources.exportstream.utils.get_validated_indices")
    def test_post_export(self, mock_get_validated, mock_validate, mock_ds_cls):
        """Test the POST request for exporting events."""
        self.login()

        # Mock the datastore instance
        mock_ds = mock.Mock()
        mock_ds_cls.return_value = mock_ds

        # Mock index validation
        mock_get_validated.return_value = (["test_index"], [1])
        mock_validate.return_value = ["test_index"]

        # Mock the query builder
        mock_ds.build_query.return_value = {"query": {"match_all": {}}}

        # Mock the slicing export generator
        mock_events = [
            {"_id": "1", "_index": "test_index", "_source": {"message": "event1"}},
            {"_id": "2", "_index": "test_index", "_source": {"message": "event2"}},
        ]
        mock_ds.export_events_with_slicing.return_value = iter(mock_events)

        data = {"filter": {"indices": [1]}, "fields": "message,datetime"}

        response = self.client.post(self.resource_url, json=data)

        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)

        # Verify the datastore was initialized with the correct pool size
        mock_ds_cls.assert_called_with(pool_maxsize=60)

        # Verify the streamed response content
        response_text = response.get_data(as_text=True).strip()
        lines = response_text.split("\n")

        self.assertEqual(len(lines), 2)
        self.assertEqual(json.loads(lines[0]), mock_events[0])
        self.assertEqual(json.loads(lines[1]), mock_events[1])

        # Verify export_events_with_slicing arguments
        mock_ds.export_events_with_slicing.assert_called_once()
        call_args = mock_ds.export_events_with_slicing.call_args
        self.assertEqual(call_args.kwargs["indices_for_pit"], ["test_index"])
        # Check that source filtering was applied based on fields
        self.assertEqual(
            call_args.kwargs["base_query_body"]["_source"], ["message", "datetime"]
        )

    @mock.patch("timesketch.api.v1.resources.exportstream.OpenSearchDataStore")
    def test_export_no_indices(self, mock_ds_cls):  # pylint: disable=unused-argument
        """Test export with no valid indices found."""
        self.login()
        with mock.patch(
            "timesketch.api.v1.resources.exportstream.utils.get_validated_indices",
            return_value=([], []),
        ):
            response = self.client.post(self.resource_url, json={})
            self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)

    def test_export_forbidden(self):
        """Test export on a sketch the user doesn't have access to."""
        # Sketch 2 is not owned by the default test user
        self.login()
        response = self.client.post("/api/v1/sketches/2/exportstream/", json={})
        self.assertEqual(response.status_code, HTTP_STATUS_CODE_FORBIDDEN)

    @mock.patch("timesketch.api.v1.resources.exportstream.OpenSearchDataStore")
    @mock.patch("timesketch.api.v1.resources.exportstream.utils.validate_indices")
    @mock.patch("timesketch.api.v1.resources.exportstream.utils.get_validated_indices")
    def test_post_export_no_valid_backend_indices(
        self, mock_get_validated, mock_validate, mock_ds_cls
    ):
        """Test export when backend validation returns no indices."""
        # pylint: disable=unused-argument
        self.login()
        mock_ds_cls.return_value = mock.Mock()

        # DB says the timeline exists...
        mock_get_validated.return_value = (["test_index"], [1])
        # ...but OpenSearch validation says it does not.
        mock_validate.return_value = []

        data = {"filter": {"indices": [1]}}
        response = self.client.post(self.resource_url, json=data)

        self.assertEqual(response.status_code, HTTP_STATUS_CODE_BAD_REQUEST)
        # Verify we hit the specific abort for empty indices_for_pit
        self.assertIn("No valid search indices", response.json["message"])


class UploadFileResourceTest(BaseTest):
    """Test UploadFileResource."""

    def setUp(self):
        super().setUp()
        self.upload_folder = tempfile.mkdtemp()
        self.app.config["UPLOAD_FOLDER"] = self.upload_folder
        self.app.config["UPLOAD_ENABLED"] = True

    def tearDown(self):
        super().tearDown()
        if os.path.exists(self.upload_folder):
            shutil.rmtree(self.upload_folder)

    @mock.patch("timesketch.api.v1.resources.upload.utils.format_upload_path")
    @mock.patch("timesketch.api.v1.resources.upload.current_app")
    def test_out_of_order_chunks(self, mock_current_app, mock_format_upload_path):
        """Test chunks uploaded out of order are written to correct file offsets."""
        mock_current_app.config = self.app.config

        # Setup mock behavior
        chunk_index_name = "00000000000000000000000000000001"
        file_path = os.path.join(self.upload_folder, chunk_index_name)
        mock_format_upload_path.return_value = file_path

        # Simulate two chunks of a file: "Hello" and "World"
        # Total size = 10 bytes
        chunk1 = b"Hello"
        chunk2 = b"World"
        total_file_size = len(chunk1) + len(chunk2)

        # Create resource instance
        resource = upload.UploadFileResource()

        # Mock request file storage
        file_storage_mock = mock.MagicMock()

        # 1. Upload chunk 2 (World) first - offset 5
        file_storage_mock.read.return_value = chunk2
        form_data = {
            "chunk_index": "1",
            "chunk_byte_offset": "5",
            "chunk_total_chunks": "2",
            "total_file_size": str(total_file_size),
            "chunk_index_name": chunk_index_name,
            "name": "test_timeline",
            "sketch_id": "1",
        }

        # We need to mock request.files and request.form
        with mock.patch("timesketch.api.v1.resources.upload.request") as mock_request:
            mock_request.files = {"file": file_storage_mock}
            mock_request.form = form_data

            # Call the internal method directly or via post()
            sketch_mock = mock.MagicMock()
            sketch_mock.id = 1

            # Call _upload_file
            # Mock _upload_and_index for the first chunk as well to avoid DB calls
            # pylint: disable=protected-access, unused-variable
            with mock.patch.object(resource, "_upload_and_index") as mock_process:
                resource._upload_file(
                    file_storage=file_storage_mock,
                    form=form_data,
                    sketch=sketch_mock,
                    index_name="",
                    chunk_index_name=chunk_index_name,
                )

        # 2. Upload chunk 1 (Hello) second - offset 0
        file_storage_mock.read.return_value = chunk1
        form_data["chunk_index"] = "0"
        form_data["chunk_byte_offset"] = "0"

        # Mock successful upload of last chunk triggers processing
        # We mock _upload_and_index because size check passes now with r+b
        # pylint: disable=protected-access, unused-variable
        with mock.patch.object(resource, "_upload_and_index") as mock_process:
            resource._upload_file(
                file_storage=file_storage_mock,
                form=form_data,
                sketch=sketch_mock,
                index_name="",
                chunk_index_name=chunk_index_name,
            )

        # Verify final file content
        with open(file_path, "rb") as f:
            content = f.read()

        # Verify content matches full file
        self.assertEqual(
            content,
            b"HelloWorld",
            "Content mismatch: chunks not assembled correctly",
        )

    @mock.patch("timesketch.api.v1.resources.upload.utils.format_upload_path")
    @mock.patch("timesketch.api.v1.resources.upload.current_app")
    def test_chunk_retry_idempotency(self, mock_current_app, mock_format_upload_path):
        """Test that retrying a chunk upload overwrites idempotently."""
        mock_current_app.config = self.app.config

        # Setup mock behavior
        chunk_index_name = "00000000000000000000000000000002"
        file_path = os.path.join(self.upload_folder, chunk_index_name)
        mock_format_upload_path.return_value = file_path

        chunk1 = b"Hello"
        total_file_size = len(chunk1)

        resource = upload.UploadFileResource()
        file_storage_mock = mock.MagicMock()
        sketch_mock = mock.MagicMock()
        sketch_mock.id = 1

        # 1. Upload chunk 0
        file_storage_mock.read.return_value = chunk1
        form_data = {
            "chunk_index": "0",
            "chunk_byte_offset": "0",
            "chunk_total_chunks": "1",
            "total_file_size": str(total_file_size),
            "chunk_index_name": chunk_index_name,
            "name": "test_timeline",
            "sketch_id": "1",
        }

        # Mock _upload_and_index
        # pylint: disable=protected-access, unused-variable
        with mock.patch.object(resource, "_upload_and_index") as mock_process:
            # First upload
            resource._upload_file(
                file_storage=file_storage_mock,
                form=form_data,
                sketch=sketch_mock,
                index_name="",
                chunk_index_name=chunk_index_name,
            )

        # 2. Retry upload chunk 0 (simulate client retry)
        file_storage_mock.read.return_value = chunk1

        # Retry upload
        try:
            # pylint: disable=protected-access, unused-variable
            with mock.patch.object(resource, "_upload_and_index") as mock_process:
                resource._upload_file(
                    file_storage=file_storage_mock,
                    form=form_data,
                    sketch=sketch_mock,
                    index_name="",
                    chunk_index_name=chunk_index_name,
                )
        except Exception:  # pylint: disable=broad-exception-caught
            # We expect an abort(400) here usually if size was wrong
            pass

        # Verify file size
        final_size = os.path.getsize(file_path)

        # Verify size matches single chunk (not doubled)
        self.assertEqual(
            final_size,
            5,
            "File size mismatch: retry should overwrite, not append",
        )
