file-enrichment

import json

import demistomock as demisto
from FileEnrichment import file_enrichment_script


def util_load_json(path: str):
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def test_file_enrichment_script_end_to_end_with_files(mocker):
    """
    Given:
        - Two file hashes (both SHA256).
        - TIM file results from test_data/mock_file_tim_results.json.
        - Batch results from test_data/mock_file_batch_results.json (create + enrich + core-get-hash-analytics-prevalence).
    When:
        - file_enrichment_script runs end-to-end (external_enrichment=True).
    Then:
        - FileEnrichmentV2 contains both hashes.
        - For file1:
            * Results has 2 entries (TIM + WildFire-v2).
            * Hashes contain MD5 + SHA256.
            * TIMScore=3, MaxScore=3, MaxVerdict=Malicious.
            * Top-level Status == "Manual" (due to manuallyEditedFields.Score).
            * TIM row in Results has NO Status/ModifiedTime (popped to top-level).
        - For file2:
            * Results has 2 entries (TIM + WildFire-v2), reliability Low.
    """
    # ---------- Load fixtures ----------
    tim_pages = util_load_json("test_data/mock_file_tim_results.json")["pages"]
    batch_blob = util_load_json("test_data/mock_file_batch_results.json")

    file_list = [
        "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",  # file1 SHA256
        "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",  # file2 SHA256
    ]

    mocker.patch.object(demisto, "args", return_value={"file_hash": ",".join(file_list)})

    # ---------- Mock execute_command ONLY for extractIndicators ----------
    def extractIndicators_side_effect(cmd, args=None, extract_contents=False, fail_on_error=True):
        if cmd == "extractIndicators":
            return [{"EntryContext": {"ExtractedIndicators": {"File": file_list}}}]
        return []

    mocker.patch("AggregatedCommandApiModule.execute_command", side_effect=extractIndicators_side_effect)

    class _MockSearcher:
        def __init__(self, pages):
            self.pages = pages

        def __iter__(self):
            return iter(self.pages)

    mocker.patch("AggregatedCommandApiModule.IndicatorsSearcher", return_value=_MockSearcher(tim_pages))

    # ---------- Enabled modules/brands (BrandManager) ----------
    mocker.patch.object(
        demisto,
        "getModules",
        return_value={
            "wf": {"state": "active", "brand": "WildFire-v2"},
            "core": {"state": "active", "brand": "Cortex Core - IR"},
        },
    )
    mocker.patch("FileEnrichment.is_platform", return_value=False)

    # ---------- Mock BatchExecutor.execute_list_of_batches using JSON ----------
    def _fake_execute_list_of_batches(self, list_of_batches, brands_to_run=None, verbose=False):
        out = []

        create_items = list(batch_blob.get("createNewIndicator", []))
        enrich_items = list(batch_blob.get("enrichIndicators", []))
        core_items = list(batch_blob.get("coreGetHashAnalyticsPrevalence", []))

        # Batch 0: CreateNewIndicatorsOnly
        batch0_cmds = list_of_batches[0]
        batch0_results = []
        for _ in batch0_cmds:
            item = create_items.pop(0) if create_items else {"Type": 1, "EntryContext": {}}
            batch0_results.append([(item, "", "")])
        out.append(batch0_results)

        # Batch 1: enrichIndicators + core-get-hash-analytics-prevalence
        # (a single prevalence command carries a list of SHA256 values)
        batch1_cmds = list_of_batches[1]
        batch1_results = []
        for cmd in batch1_cmds:
            if cmd.name == "enrichIndicators":
                items = enrich_items or [{"Type": 1, "EntryContext": {}, "Metadata": {"brand": "WildFire-v2"}}]
                batch1_results.append([(e, "", "") for e in items])
            elif cmd.name == "core-get-hash-analytics-prevalence":
                items = core_items or [
                    {
                        "Type": 1,
                        "EntryContext": {},
                        "Metadata": {"brand": "Cortex Core - IR"},
                    }
                ]
                batch1_results.append([(item, "", "") for item in items])
            else:
                batch1_results.append([({"Type": 1, "EntryContext": {}}, "", "")])
        out.append(batch1_results)

        return out

    mocker.patch("AggregatedCommandApiModule.BatchExecutor.execute_list_of_batches", _fake_execute_list_of_batches)

    # ---------- Act ----------
    command_results = file_enrichment_script(
        file_list=file_list,
        external_enrichment=True,
        verbose=True,
        enrichment_brands=["WildFire-v2"],
        additional_fields=False,
    )
    outputs = command_results.outputs

    # ---------- Assert: FileEnrichmentV2 indicators ----------
    enrichment_key = "FileEnrichment(val.Value && val.Value == obj.Value)"
    enrichment_list = outputs.get(enrichment_key, [])
    assert len(enrichment_list) == 2

    enrichment_map = {item["Value"]: item for item in enrichment_list}
    # In this scenario, Value will be the canonical TIM "value" (sha256) → same as file_list
    assert set(enrichment_map.keys()) == set(file_list)

    # ---- file1 assertions ----
    f1 = enrichment_map[file_list[0]]
    brands_present_f1 = {r.get("Brand") for r in f1["Results"]}
    assert brands_present_f1 == {"TIM", "WildFire-v2"}
    assert len(f1["Results"]) == 2

    # Hashes aggregated from TIM indicator
    hashes1 = f1.get("Hashes", {})
    assert hashes1.get("MD5") is not None
    assert hashes1.get("SHA256") == file_list[0]

    # TIM row present but without Status/ModifiedTime (popped to top-level)
    tim_row_f1 = next(r for r in f1["Results"] if r["Brand"] == "TIM")
    assert "Status" not in tim_row_f1
    assert "ModifiedTime" not in tim_row_f1

    # Top-level scores & status
    assert f1["TIMScore"] == 3
    assert f1["MaxScore"] == 3
    assert f1["MaxVerdict"] == "Malicious"
    assert f1["Status"] == "Manual"  # due to manuallyEditedFields.Score in TIM IOC

    # ---- file2 assertions ----
    f2 = enrichment_map[file_list[1]]
    brands_present_f2 = {r.get("Brand") for r in f2["Results"]}
    assert brands_present_f2 == {"TIM", "WildFire-v2"}
    assert len(f2["Results"]) == 2

    wf2 = next(r for r in f2["Results"] if r["Brand"] == "WildFire-v2")
    assert wf2["Score"] == 1
    assert wf2.get("Reliability") == "Low"


def _capture_built_commands(mocker, file_list, is_platform=False):
    """Helper: runs file_enrichment_script with mocks and returns the command batches exactly as
    they were built by file_enrichment_script, captured before any brand/type filtering is applied."""
    mocker.patch.object(demisto, "args", return_value={"file_hash": ",".join(file_list)})
    mocker.patch("FileEnrichment.is_platform", return_value=is_platform)

    def extractIndicators_side_effect(cmd, args=None, extract_contents=False, fail_on_error=True):
        if cmd == "extractIndicators":
            return [{"EntryContext": {"ExtractedIndicators": {"File": file_list}}}]
        return []

    mocker.patch("AggregatedCommandApiModule.execute_command", side_effect=extractIndicators_side_effect)

    # Build a minimal TIM page so each input hash is "found" in TIM. This keeps the
    # enrichment flow on the success path (instead of the all-failed path which now
    # raises), while these tests only care about the commands that were built.
    def _build_ioc(file_hash):
        custom_fields = {"sha256": file_hash} if len(file_hash) == 64 else {"md5": file_hash}
        return {
            "value": file_hash,
            "score": 1,
            "modifiedTime": "2025-09-01T00:00:00Z",
            "CustomFields": custom_fields,
        }

    tim_pages = [{"iocs": [_build_ioc(file_hash) for file_hash in file_list]}]

    class _MockSearcher:
        def __init__(self, pages):
            self.pages = pages

        def __iter__(self):
            return iter(self.pages)

    mocker.patch("AggregatedCommandApiModule.IndicatorsSearcher", return_value=_MockSearcher(tim_pages))
    mocker.patch.object(
        demisto,
        "getModules",
        return_value={
            "wf": {"state": "active", "brand": "WildFire-v2"},
            "core": {"state": "active", "brand": "Cortex Core - IR"},
        },
    )

    captured = {}

    def _fake_prepare_commands_batches(self, external_enrichment=False):
        # Capture the commands as built (pre-filter) so the assertions are independent of
        # the brand/type filtering performed by the real prepare_commands_batches.
        captured["batches"] = self.commands
        return self.commands

    mocker.patch(
        "AggregatedCommandApiModule.ReputationAggregatedCommand.prepare_commands_batches",
        _fake_prepare_commands_batches,
    )

    def _fake_execute_list_of_batches(self, list_of_batches, brands_to_run=None, verbose=False):
        return [[[({"Type": 1, "EntryContext": {}}, "", "")] for _ in batch] for batch in list_of_batches]

    mocker.patch("AggregatedCommandApiModule.BatchExecutor.execute_list_of_batches", _fake_execute_list_of_batches)

    file_enrichment_script(
        file_list=file_list,
        external_enrichment=True,
        verbose=False,
        enrichment_brands=["WildFire-v2"],
        additional_fields=False,
    )
    return captured["batches"]


def test_prevalence_command_omitted_when_no_sha256(mocker):
    """
    Given:
        - A single MD5 hash (no SHA256) as input.
    When:
        - file_enrichment_script builds the command batches.
    Then:
        - The core-get-hash-analytics-prevalence command is NOT included.
    """
    md5_hash = "d41d8cd98f00b204e9800998ecf8427e"
    batches = _capture_built_commands(mocker, [md5_hash])

    all_command_names = [cmd.name for batch in batches for cmd in batch]
    assert "core-get-hash-analytics-prevalence" not in all_command_names


def test_prevalence_command_included_when_sha256_present(mocker):
    """
    Given:
        - A SHA256 hash as input.
    When:
        - file_enrichment_script builds the command batches.
    Then:
        - A single core-get-hash-analytics-prevalence command is included with the SHA256 value.
    """
    sha256_hash = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    batches = _capture_built_commands(mocker, [sha256_hash])

    prevalence_cmds = [cmd for batch in batches for cmd in batch if cmd.name == "core-get-hash-analytics-prevalence"]
    assert len(prevalence_cmds) == 1
    assert prevalence_cmds[0].args == {"sha256": [sha256_hash]}


def test_file_enrichment_uses_builtin_command_on_platform(mocker):
    """
    Given:
        - A SHA256 hash as input, running on the unified Cortex platform (is_platform True).
    When:
        - file_enrichment_script builds the command batches.
    Then:
        - The prevalence command is the built-in "getHashAnalyticsPrevalence".
        - Its command_type is CommandType.BUILTIN (not the legacy INTERNAL core command).
    """
    from AggregatedCommandApiModule import CommandType

    sha256_hash = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    batches = _capture_built_commands(mocker, [sha256_hash], is_platform=True)

    prevalence_cmds = [cmd for batch in batches for cmd in batch if cmd.name == "getHashAnalyticsPrevalence"]
    assert len(prevalence_cmds) == 1
    assert prevalence_cmds[0].command_type == CommandType.BUILTIN
    assert "core-get-hash-analytics-prevalence" not in [cmd.name for batch in batches for cmd in batch]