"""Tests for ThreatintelPlugin."""

import json
from unittest import mock

from flask import current_app

from timesketch.lib.analyzers import yetiindicators
from timesketch.lib.testlib import BaseTest
from timesketch.lib.testlib import MockDataStore

MOCK_YETI_ENTITIES = [
    {
        "type": "malware",
        "name": "xmrig",
        "description": "coin miner",
        "created": "2024-02-16T12:10:48.670039Z",
        "modified": "2024-02-16T12:10:48.670040Z",
        "kill_chain_phases": [],
        "aliases": [],
        "family": "miner",
        "id": "2152646",
        "tags": [],
        "root_type": "entity",
    }
]

MOCK_YETI_NEIGHBORS_RESPONSE = {
    "vertices": {
        "indicators/2152802": {
            "name": "typo'd dhcpd",
            "type": "regex",
            "description": "Random description",
            "created": "2024-02-16T12:12:14.564723Z",
            "modified": "2024-02-16T12:12:14.564729Z",
            "valid_from": "2024-02-16T12:12:14.564730Z",
            "valid_until": "2024-03-17T12:12:14.564758Z",
            "pattern": "/usr/bin/dhpcd",
            "location": "filesystem",
            "diamond": "victim",
            "kill_chain_phases": [],
            "relevant_tags": ["xmrig", "malware"],
            "id": "2152802",
            "root_type": "indicator",
        },
        "indicators/2152803": {
            "name": "noisy indicator",
            "type": "regex",
            "description": "Random description",
            "created": "2024-02-16T12:12:14.564723Z",
            "modified": "2024-02-16T12:12:14.564729Z",
            "valid_from": "2024-02-16T12:12:14.564730Z",
            "valid_until": "2024-03-17T12:12:14.564758Z",
            "pattern": "ba",
            "location": "filesystem",
            "diamond": "victim",
            "kill_chain_phases": [],
            "relevant_tags": ["noise"],
            "id": "2152803",
            "root_type": "indicator",
            "tags": ["timesketch:mute"],
        },
    },
    "paths": [
        [
            {
                "source": "entities/2152646",
                "target": "indicators/2152802",
                "type": "dropped in",
                "description": "",
                "created": "2024-02-16T12:28:52.731740Z",
                "modified": "2024-02-16T12:28:52.731747Z",
                "id": "2153330",
            },
            {
                "source": "entities/2152646",
                "target": "indicators/2152803",
                "type": "uses",
                "description": "",
                "created": "2024-02-16T12:28:52.731740Z",
                "modified": "2024-02-16T12:28:52.731747Z",
                "id": "2153331",
            },
        ]
    ],
    "total": 1,
}


MATCHING_PATH_MESSAGE = {
    "__ts_timeline_id": 1,
    "es_index": "",
    "es_id": "",
    "label": "",
    "timestamp": 1410895419859714,
    "timestamp_desc": "",
    "datetime": "2014-09-16T19:23:40+00:00",
    "source_short": "",
    "source_long": "",
    "message": "/usr/bin/dhpcd",
}

MATCHING_BLOOM_HASH = {
    "__ts_timeline_id": 1,
    "es_index": "",
    "es_id": "",
    "label": "",
    "timestamp": 1410895419859714,
    "timestamp_desc": "",
    "datetime": "2014-09-16T19:23:40+00:00",
    "sha256_hash": "testhash",
    "message": "testhash",
}

FALSE_POSITIVE_PATH_MESSAGE = {
    "__ts_timeline_id": 1,
    "es_index": "",
    "es_id": "",
    "label": "",
    "timestamp": 1410895419859714,
    "timestamp_desc": "",
    "datetime": "2014-09-16T19:23:40+00:00",
    "source_short": "",
    "source_long": "",
    "message": "/bin/bash",
}


class YetiTestAnalyzer(yetiindicators.YetiGraphAnalyzer):
    NAME = "yetitest"
    DISPLAY_NAME = "Test for yeti analyzer"
    DESCRIPTION = "Just an analyzer for teting"

    _TYPE_SELECTOR = ["investigation:tag1,tag2", "malware"]
    _TARGET_NEIGHBOR_TYPE = ["sigma", "query", "regex", "observable"]
    _SAVE_INTELLIGENCE = True
    _DIRECTION = "any"
    _MAX_HOPS = 1


class TestYetiIndicators(BaseTest):
    """Tests the functionality of the YetiIndicators analyzer."""

    def setUp(self):
        super().setUp()
        current_app.config["YETI_API_ROOT"] = "blah"
        current_app.config["YETI_API_KEY"] = "blah"

    # Mock the OpenSearch datastore and the YetiApi
    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_api_query(self, mock_yeti_api_class):
        """Tests that queries to the API are well-formed."""
        mock_api = mock_yeti_api_class.return_value
        mock_api.search_entities.return_value = MOCK_YETI_ENTITIES
        mock_api.search_graph.return_value = MOCK_YETI_NEIGHBORS_RESPONSE

        analyzer = YetiTestAnalyzer("test_index", 1, 123)
        analyzer.datastore.client = mock.Mock()
        analyzer.run()

        mock_api.search_entities.assert_any_call(
            entity_type="investigation", tags=["tag1", "tag2"], count=0
        )
        mock_api.search_entities.assert_any_call(
            entity_type="malware", tags=[], count=0
        )

    # Mock the OpenSearch datastore and the YetiApi
    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_indicator_match(self, mock_yeti_api_class):
        """Test that ES queries for indicators are correctly built."""
        mock_api = mock_yeti_api_class.return_value
        mock_api.search_entities.return_value = MOCK_YETI_ENTITIES
        mock_api.search_graph.return_value = MOCK_YETI_NEIGHBORS_RESPONSE

        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        analyzer.datastore.client = mock.Mock()
        analyzer.datastore.import_event("test_index", MATCHING_PATH_MESSAGE, "0")

        message = json.loads(analyzer.run())
        self.assertEqual(
            message["result_summary"],
            (
                "1 events matched 1/1 indicators (0 failed).\n\n"
                "Entities found: xmrig:malware"
            ),
        )
        mock_api.search_entities.assert_called()
        mock_api.search_graph.assert_called()
        self.assertEqual(
            sorted(analyzer.tagged_events["0"]["tags"]), sorted(["malware", "xmrig"])
        )

    # Mock the OpenSearch datastore and the YetiApi
    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_indicator_nomatch(self, mock_yeti_api_class):
        """Test that ES queries for indicators are correctly built."""
        mock_api = mock_yeti_api_class.return_value
        mock_api.search_entities.return_value = MOCK_YETI_ENTITIES
        mock_api.search_graph.return_value = MOCK_YETI_NEIGHBORS_RESPONSE

        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        analyzer.datastore.client = mock.Mock()

        message = json.loads(analyzer.run())
        self.assertEqual(
            message["result_summary"],
            "0/1 indicators were found in the timeline (0 failed)",
        )
        mock_api.search_entities.assert_called()
        mock_api.search_graph.assert_called()

    # Mock the OpenSearch datastore and the YetiApi
    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_nomatch_muted_indicator(self, mock_yeti_api_class):
        """Test that ES queries for indicators are correctly built."""
        mock_api = mock_yeti_api_class.return_value
        mock_api.search_entities.return_value = MOCK_YETI_ENTITIES
        mock_api.search_graph.return_value = MOCK_YETI_NEIGHBORS_RESPONSE

        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        analyzer.datastore.client = mock.Mock()
        analyzer.datastore.import_event("test_index", FALSE_POSITIVE_PATH_MESSAGE, "0")

        message = json.loads(analyzer.run())
        self.assertEqual(
            message["result_summary"],
            (
                "1 events matched 1/1 indicators (0 failed).\n\n"
                "Entities found: xmrig:malware"
            ),
        )
        mock_api.search_entities.assert_called()
        mock_api.search_graph.assert_called()

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_slug(self, _):
        """Tests that slugs are formed correctly."""
        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        mock_event = mock.Mock()
        mock_event.get_comments.return_value = []
        analyzer.mark_event(
            MOCK_YETI_NEIGHBORS_RESPONSE["vertices"]["indicators/2152802"],
            mock_event,
            MOCK_YETI_ENTITIES,
        )
        mock_event.add_tags.assert_called_once()
        self.assertIn(
            sorted(["xmrig", "malware"]),
            [sorted(x) for x in mock_event.add_tags.call_args[0]],
        )

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_build_query_from_regexp(self, _):
        """Tests that that queries are correctly built from regex indicators."""
        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        query = analyzer.build_query_from_regexp(
            {
                "name": "random regex",
                "type": "regex",
                "description": "Random description",
                "created": "2024-02-16T12:12:14.564723Z",
                "modified": "2024-02-16T12:12:14.564729Z",
                "valid_from": "2024-02-16T12:12:14.564730Z",
                "valid_until": "2024-03-17T12:12:14.564758Z",
                "pattern": r"this_is_my_{2,3}regex[0-9]",
                "location": "filesystem",
                "diamond": "victim",
                "kill_chain_phases": [],
                "relevant_tags": ["regex"],
                "id": "2152802",
                "root_type": "indicator",
            }
        )
        self.assertEqual(
            query,
            {
                "query": {
                    "bool": {
                        "should": [
                            {
                                "regexp": {
                                    "filename.keyword": {
                                        "value": r".*this_is_my_{2,3}regex[0-9].*",
                                        "case_insensitive": True,
                                    }
                                }
                            },
                            {
                                "regexp": {
                                    "display_name.keyword": {
                                        "value": r".*this_is_my_{2,3}regex[0-9].*",
                                        "case_insensitive": True,
                                    }
                                }
                            },
                        ]
                    }
                }
            },
        )

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_build_query_from_sigma(self, _):
        """Tests that that queries are correctly built from sigma indicators."""
        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        sigma_pattern = """id: asd
title: test
description: test
status: experimental
date: 2024/01/01
author: 'tomchop'
references:
  - blah
logsource:
    category: winprefetch
    service: winprefetch
detection:
    selection1:
      Image|endswith:
        - '\\rundll32.exe'
    condition: selection1
falsepositives:
    - Unknown
level: medium
tags:
    - blah
"""
        query = analyzer.build_query_from_sigma(
            {
                "name": "random sigma",
                "type": "sigma",
                "description": "Random description",
                "created": "2024-02-16T12:12:14.564723Z",
                "modified": "2024-02-16T12:12:14.564729Z",
                "valid_from": "2024-02-16T12:12:14.564730Z",
                "valid_until": "2024-03-17T12:12:14.564758Z",
                "pattern": sigma_pattern,
                "location": "not_used",
                "diamond": "victim",
                "kill_chain_phases": [],
                "relevant_tags": ["sigma"],
                "id": "2152802",
                "root_type": "indicator",
            }
        )
        self.assertEqual(
            query,
            {
                "query": {
                    "query_string": {
                        "query": (
                            '(data_type:"windows\\:prefetch\\:execution" '
                            'AND message:("\\\\rundll32.exe"))'
                        )
                    }
                }
            },
        )

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_build_query_from_observable(self, _):
        """Tests that that queries are correctly built from regex indicators."""
        analyzer = yetiindicators.YetiBadnessIndicators("test_index", 1, 123)
        query = analyzer.build_query_from_observable(
            {
                "value": "C:\\ProgramFiles\\mimi.exe",
                "type": "url",
                "created": "2024-04-18T08:42:11.330182Z",
                "context": [],
                "last_analysis": {},
                "id": "46833442",
                "tags": [
                    {
                        "name": "mimikatz",
                        "last_seen": "2024-04-18T08:42:11.370806Z",
                        "expires": "2024-05-18T08:42:11.370811Z",
                        "fresh": True,
                        "id": None,
                    }
                ],
                "root_type": "observable",
            }
        )
        self.assertEqual(
            query,
            {
                "query": {
                    "wildcard": {
                        "message.keyword": {
                            "value": "*C:\\\\ProgramFiles\\\\mimi.exe*",
                            "case_insensitive": True,
                        }
                    }
                }
            },
        )

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore")
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_run_composite_aggregation(
        self, _, mock_opensearch_datastore
    ):  # pylint: disable=invalid-name
        """Tests that the composite aggregation is correctly built."""
        analyzer = yetiindicators.YetiBloomChecker("test_index", 1, 123)
        mock_opensearch_datastore.return_value.search.return_value = {
            "aggregations": {
                "my_composite_agg_test_index": {
                    "buckets": [{"key": {"sha256_hash": "testhash"}}],
                }
            }
        }
        hashset = set()
        analyzer.run_composite_aggregation(hashset)
        self.assertEqual(hashset, {"testhash"})

    @mock.patch("timesketch.lib.analyzers.interface.OpenSearchDataStore", MockDataStore)
    @mock.patch("timesketch.lib.analyzers.yetiindicators.YetiApi")
    def test_bloomanalyzer_run(self, mock_yeti_api_class):
        """Tests the YetiBloomAnalyzer run function."""
        mock_api = mock_yeti_api_class.return_value
        mock_api.search_bloom.return_value = [
            {"value": "testhash", "hits": ["hitsource1"]}
        ]

        analyzer = yetiindicators.YetiBloomChecker("test_index", 1, 123)

        # pylint: disable=unused-argument
        def agg_mock(hashset, after_key=None):
            """Mock for the composite aggregation."""
            hashset.add("testhash")
            return hashset, None

        analyzer.run_composite_aggregation = agg_mock
        analyzer.datastore.import_event("test_index", MATCHING_BLOOM_HASH, "0")

        message = json.loads(analyzer.run())

        self.assertEqual(
            message["result_summary"],
            (
                "Bloom filter check completed. 1 hashes checked, 1 hits "
                "found, 1 events tagged."
            ),
        )
        self.assertEqual(
            message["platform_meta_data"]["created_tags"], ["bloom:hitsource1"]
        )
