from unittest.mock import patch, MagicMock
import io
import zipfile

import CreateSigmaRuleIndicator
import demistomock as demisto  # noqa: F401
import pytest
from CommonServerPython import *  # noqa: F401
from CreateSigmaRuleIndicator import (
    create_indicator_relationships,
    get_mitre_technique_name,
    parse_and_create_indicator,
    parse_detection_field,
    parse_tags,
    create_relationship,
    main,
    tim_create_indicators,
    extract_rules_from_zip,
)
from sigma.rule import SigmaRule, SigmaRuleTag


def load_file(path: str) -> dict[str, Any]:
    with open(path) as f:
        return json.load(f)


def test_create_relationship():
    indicator = "Sigma Rule"
    entity_b = "Command and Scripting Interpreter"
    entity_b_type = "Attack Pattern"
    relation_type = "detects"
    result = EntityRelationship(
        entity_a="Sigma Rule",
        entity_a_type="Sigma Rule Indicator",
        name="detects",
        reverse_name="detected-by",
        entity_b="Command and Scripting Interpreter",
        entity_b_type="Attack Pattern",
    )
    assert create_relationship(indicator, entity_b, entity_b_type, relation_type).to_context() == result.to_context()


@pytest.mark.parametrize(
    "input, expected_result",
    [
        pytest.param(
            [SigmaRuleTag(namespace="attack", name="t1059", source=None)],
            (
                [{"value": "Command and Scripting Interpreter", "type": "Attack Pattern"}],
                ["T1059 - Command and Scripting Interpreter"],
                "CLEAR",
            ),
            id="Tag Creation - MITRE technique",
        ),
        pytest.param(
            [SigmaRuleTag(namespace="attack", name="resource-development"), SigmaRuleTag(namespace="tlp", name="RED")],
            ([], ["Resource Development"], "RED"),
            id="Tag Creation - MITRE tactic",
        ),
        pytest.param(
            [SigmaRuleTag(namespace="cve", name="2024-3400")],
            ([{"value": "CVE-2024-3400", "type": "CVE"}], ["CVE-2024-3400"], "CLEAR"),
            id="Tag Creation - CVEs",
        ),
    ],
)
@patch.object(CreateSigmaRuleIndicator, "get_mitre_technique_name")
def test_parse_tags(mock_get_mitre_technique_name, input, expected_result):
    mock_get_mitre_technique_name.return_value = "Command and Scripting Interpreter"
    assert parse_tags(input) == expected_result


@patch.object(CreateSigmaRuleIndicator, "execute_command")
def test_get_mitre_technique_name(mock_execute_command):
    mock_execute_command.return_value = True, {"value": "Command and Scripting Interpreter"}
    mitre_id = "T1059"
    indicator_type = "Attack Pattern"
    get_mitre_technique_name(mitre_id, indicator_type)
    mock_execute_command.assert_called_with(
        command="SearchIndicator", args={"query": f'type:"Attack Pattern" and {mitre_id}'}, fail_on_error=False
    )


@patch.object(CreateSigmaRuleIndicator, "create_relationship")
@patch.object(CreateSigmaRuleIndicator, "return_results")
def test_create_indicator_relationships(mock_return_results, mock_create_relationship):
    mock_create_relationship.return_value("relationship")
    indicator = "Sigma Rule Test"
    product = "Windows"
    relationships = [{"value": "Some technique", "type": "Attack Pattern"}, {"value": "CVE-2024-111", "type": "CVE"}]
    create_indicator_relationships(indicator, product, relationships)
    assert mock_create_relationship.call_count == 3


def test_parse_detection_field():
    with open("test_data/sigma_rule.yml") as f:
        sigma_rule = SigmaRule.from_yaml(f.read())

    result = [{"selection": "selection", "key": "displaymessage", "modifiers": "", "values": "(1)Max sign in attempts exceeded"}]
    assert parse_detection_field(sigma_rule) == result


def test_parse_and_create_indicator():
    with open("test_data/sigma_rule.yml") as f:
        rule = f.read()
    expected_indicator = load_file("test_data/expected_indicator.json")
    result = parse_and_create_indicator(SigmaRule.from_yaml(rule), raw_rule=rule)
    assert result["indicator"] == expected_indicator["indicator"]


@patch.object(demisto, "args")
@patch.object(CreateSigmaRuleIndicator, "return_results")
@patch.object(CreateSigmaRuleIndicator, "execute_command")
def test_main(mock_executeCommand, mock_return_results, mock_args):
    with open("test_data/sigma_rule.yml") as f:
        rule = f.read()

    mock_args.return_value = {"sigma_rule_str": rule, "entry_id": "", "create_indicators": "True"}
    main()
    mock_return_results.assert_called_once()
    args, kwargs = mock_return_results.call_args
    assert args[0].readable_output == "1 Sigma Rule(s) Created.\n1 Relationship(s) Created."


@patch("CreateSigmaRuleIndicator.execute_command")
@patch("CreateSigmaRuleIndicator.create_indicator_relationships")
@patch("CreateSigmaRuleIndicator.demisto.debug")
@patch("CreateSigmaRuleIndicator.time.time")
def test_tim_create_indicators(mock_time, mock_debug, mock_create_relationships, mock_execute_command):
    sample_indicators = [
        {
            "indicator": {
                "value": "Test Sigma Rule 1",
                "type": "Sigma Rule",
                "product": "windows",
            },
            "relationships": [
                {"value": "T1562 - Impair Defenses", "type": "Attack Pattern"},
                {"value": "CVE-2021-44228", "type": "CVE"},
            ],
        },
        {
            "indicator": {
                "value": "Test Sigma Rule 2",
                "type": "Sigma Rule",
                "product": "linux",
            },
            "relationships": [{"value": "S0601 - Cobalt Strike", "type": "Tool"}],
        },
    ]

    # Setup mocks
    mock_time.side_effect = [100, 105]  # Start and end times

    # Mock relationships that would be created
    mock_create_relationships.side_effect = [
        [
            EntityRelationship(
                entity_a="Test Sigma Rule 1",
                entity_a_type="Sigma Rule Indicator",
                name="related-to",
                entity_b="Windows",
                entity_b_type="Software",
            ),
            EntityRelationship(
                entity_a="Test Sigma Rule 1",
                entity_a_type="Sigma Rule Indicator",
                name="detects",
                entity_b="T1562 - Impair Defenses",
                entity_b_type="Attack Pattern",
            ),
            EntityRelationship(
                entity_a="Test Sigma Rule 1",
                entity_a_type="Sigma Rule Indicator",
                name="detects",
                entity_b="CVE-2021-44228",
                entity_b_type="CVE",
            ),
        ],
        [
            EntityRelationship(
                entity_a="Test Sigma Rule 2",
                entity_a_type="Sigma Rule Indicator",
                name="related-to",
                entity_b="Linux",
                entity_b_type="Software",
            ),
            EntityRelationship(
                entity_a="Test Sigma Rule 2",
                entity_a_type="Sigma Rule Indicator",
                name="detects",
                entity_b="S0601 - Cobalt Strike",
                entity_b_type="Tool",
            ),
        ],
    ]

    # Execute the function
    result = tim_create_indicators(sample_indicators)

    # Verify results
    assert mock_execute_command.call_count == 2
    mock_execute_command.assert_any_call("createNewIndicator", sample_indicators[0]["indicator"])
    mock_execute_command.assert_any_call("createNewIndicator", sample_indicators[1]["indicator"])

    assert mock_create_relationships.call_count == 2
    mock_create_relationships.assert_any_call("Test Sigma Rule 1", "windows", sample_indicators[0]["relationships"])
    mock_create_relationships.assert_any_call("Test Sigma Rule 2", "linux", sample_indicators[1]["relationships"])

    mock_debug.assert_called_once_with("2 indicators created. in 5 seconds")

    # Verify returned CommandResults
    assert isinstance(result, CommandResults)
    assert result.readable_output == "2 Sigma Rule(s) Created.\n5 Relationship(s) Created."
    assert len(result.relationships) == 5  # Total relationships from both indicators


def create_zip_mock(file_contents, file_names):
    """Helper to create a mock zipfile"""
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zip_file:
        for name, content in zip(file_names, file_contents):
            zip_file.writestr(name, content)
    return zip_buffer.getvalue()


@pytest.fixture
def valid_sigma_rule():
    return """
title: Test Rule
id: 12345678-1234-1234-1234-123456789012
status: test
description: Test rule for unit testing
author: Test Author
date: 2023/01/01
logsource:
    category: process_creation
    product: windows
detection:
    selection:
        CommandLine|contains: suspicious.exe
    condition: selection
level: high
tags:
    - attack.t1055
    - attack.s0001
    - cve.2023.1234
"""


@pytest.fixture
def invalid_sigma_rule():
    return """
title: Invalid Rule
status: test
# Missing required fields like id, detection, etc.
"""


@patch("zipfile.ZipFile")
def test_extract_rules_from_zip_successful(mock_zipfile, valid_sigma_rule):
    """Test extracting valid rules from a zip file"""

    # Setup mock zipfile with valid rules
    mock_zip_instance = MagicMock()
    mock_zipfile.return_value.__enter__.return_value = mock_zip_instance
    mock_zip_instance.namelist.return_value = ["rule1.yml", "rule2.yml", "__pycache__/ignored.yml", ".hidden.yml"]

    # Mock file content reading
    mock_file = MagicMock()
    mock_file.read.return_value = valid_sigma_rule.encode("utf-8")
    mock_zip_instance.open.return_value.__enter__.return_value = mock_file

    # Mock parse_and_create_indicator to return a predetermined value
    expected_indicator = {"indicator": {"value": "Test Rule"}, "relationships": []}
    with patch("CreateSigmaRuleIndicator.parse_and_create_indicator", return_value=expected_indicator):
        result = extract_rules_from_zip("test.zip")

    # Assertions
    assert len(result) == 2  # Should extract 2 valid files
    assert result[0] == expected_indicator
    assert result[1] == expected_indicator

    # Verify correct files were processed
    mock_zip_instance.open.assert_any_call("rule1.yml")
    mock_zip_instance.open.assert_any_call("rule2.yml")
    assert mock_zip_instance.open.call_count == 2


@patch("zipfile.ZipFile")
def test_extract_rules_from_zip_with_errors(mock_zipfile, valid_sigma_rule, invalid_sigma_rule, capfd):
    """Test handling errors during rule extraction"""

    # Setup mock zipfile with mixed valid and invalid rules
    mock_zip_instance = MagicMock()
    mock_zipfile.return_value.__enter__.return_value = mock_zip_instance
    mock_zip_instance.namelist.return_value = ["valid.yml", "invalid.yml"]

    # Setup mock file reading to return different content based on filename
    def mock_open_file(filename):
        from unittest.mock import MagicMock

        mock_file = MagicMock()
        if filename == "valid.yml":
            mock_file.read.return_value = valid_sigma_rule.encode("utf-8")
        else:
            mock_file.read.return_value = invalid_sigma_rule.encode("utf-8")
        return mock_file

    mock_zip_instance.open.side_effect = lambda filename: MagicMock(__enter__=lambda x: mock_open_file(filename))

    # Mock SigmaRule.from_yaml to raise an exception for the invalid rule
    original_from_yaml = SigmaRule.from_yaml

    def mock_from_yaml(yaml_str):
        if "Invalid Rule" in yaml_str:
            raise Exception("Invalid rule format")
        return original_from_yaml(yaml_str)

    # Mock parse_and_create_indicator for the valid rule
    expected_indicator = {"indicator": {"value": "Test Rule"}, "relationships": []}

    with (
        patch("sigma.rule.SigmaRule.from_yaml", side_effect=mock_from_yaml),
        patch("CreateSigmaRuleIndicator.parse_and_create_indicator", return_value=expected_indicator),
    ):
        result = extract_rules_from_zip("test.zip")

    # Assertions
    out, err = capfd.readouterr()
    assert len(result) == 1  # Only the valid rule should be processed
    assert result[0] == expected_indicator
    assert out == 'SGM: Error parsing Sigma rule from file "invalid.yml": Invalid rule format\n'


@patch("zipfile.ZipFile")
def test_extract_rules_from_zip_empty(mock_zipfile):
    """Test extracting rules from a zip file with no valid rules"""
    # Setup mock zipfile with no valid rules
    mock_zip_instance = MagicMock()
    mock_zipfile.return_value.__enter__.return_value = mock_zip_instance
    mock_zip_instance.namelist.return_value = ["__pycache__/ignored.yml", ".hidden.yml", "not_a_rule.txt"]

    result = extract_rules_from_zip("test.zip")

    # Assertions
    assert len(result) == 0  # No valid rules should be found
    assert not mock_zip_instance.open.called  # No files should be opened