test/test_malware_scan.py

"""Tests for GuardDuty S3 upload malware scanning."""

from __future__ import annotations

import os
from unittest.mock import MagicMock, patch

import gradio as gr
import pytest
from botocore.exceptions import ClientError

from tools import malware_scan as ms


@pytest.fixture(autouse=True)
def _reset_scan_cache():
    ms.clear_scan_cache()
    yield
    ms.clear_scan_cache()


def test_malware_scan_enabled_requires_all_flags(monkeypatch):
    import tools.config as config

    monkeypatch.setattr(config, "SCAN_UPLOADS_FOR_MALWARE", True)
    monkeypatch.setattr(config, "RUN_AWS_FUNCTIONS", True)
    monkeypatch.setattr(config, "MALWARE_SCAN_S3_BUCKET", "scan-bucket")
    assert ms.malware_scan_enabled() is True

    monkeypatch.setattr(config, "SCAN_UPLOADS_FOR_MALWARE", False)
    assert ms.malware_scan_enabled() is False

    monkeypatch.setattr(config, "SCAN_UPLOADS_FOR_MALWARE", True)
    monkeypatch.setattr(config, "MALWARE_SCAN_S3_BUCKET", "")
    assert ms.malware_scan_enabled() is False


def test_normalize_gradio_file_paths():
    assert ms.normalize_gradio_file_paths(None) == []
    assert ms.normalize_gradio_file_paths("/tmp/a.pdf") == [
        os.path.abspath("/tmp/a.pdf")
    ]
    assert ms.normalize_gradio_file_paths({"name": "/tmp/b.pdf"}) == [
        os.path.abspath("/tmp/b.pdf")
    ]

    class _FileObj:
        name = "/tmp/c.pdf"

    assert ms.normalize_gradio_file_paths(_FileObj()) == [os.path.abspath("/tmp/c.pdf")]


def test_scan_local_file_noop_when_disabled(monkeypatch, tmp_path):
    sample = tmp_path / "doc.pdf"
    sample.write_bytes(b"%PDF-1.4")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: False)
    ms.scan_local_file_for_malware(str(sample))


def test_scan_local_file_clean_deletes_staging_object(monkeypatch, tmp_path):
    sample = tmp_path / "clean.pdf"
    sample.write_bytes(b"clean")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [
            {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "NO_THREATS_FOUND"}
        ]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        ms.scan_local_file_for_malware(str(sample))

    s3.upload_file.assert_called_once()
    s3.delete_object.assert_called_once()
    assert ms.already_scanned_clean(str(sample))


def test_scan_local_file_threat_deletes_s3(monkeypatch, tmp_path):
    sample = tmp_path / "bad.pdf"
    sample.write_bytes(b"eicar")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [{"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "THREATS_FOUND"}]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(ms.MalwareScanRejectedError):
            ms.scan_local_file_for_malware(str(sample))

    s3.delete_object.assert_called_once()
    assert sample.exists()


def test_scan_gradio_file_upload_threat_deletes_local(monkeypatch, tmp_path):
    sample = tmp_path / "bad.pdf"
    sample.write_bytes(b"eicar")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [{"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "THREATS_FOUND"}]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(gr.Error):
            ms.scan_gradio_file_upload(str(sample))

    s3.delete_object.assert_called_once()
    assert not sample.exists()


def test_scan_local_file_failed_status_fail_closed(monkeypatch, tmp_path):
    sample = tmp_path / "failed.pdf"
    sample.write_bytes(b"x")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [{"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "FAILED"}]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(ms.MalwareScanRejectedError):
            ms.scan_local_file_for_malware(str(sample))

    s3.delete_object.assert_called_once()


def test_scan_local_file_timeout_deletes_s3(monkeypatch, tmp_path):
    sample = tmp_path / "slow.pdf"
    sample.write_bytes(b"x")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 0.02, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {"TagSet": []}

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(ms.MalwareScanRejectedError, match="timed out"):
            ms.scan_local_file_for_malware(str(sample))

    s3.delete_object.assert_called_once()


def test_scan_local_file_poll_exception_still_deletes_s3(monkeypatch, tmp_path):
    sample = tmp_path / "err.pdf"
    sample.write_bytes(b"x")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.side_effect = RuntimeError("aws down")

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(ms.MalwareScanRejectedError, match="timed out"):
            ms.scan_local_file_for_malware(str(sample))

    s3.delete_object.assert_called_once()


def test_multi_file_second_failure_cleans_first_staging(monkeypatch, tmp_path):
    first = tmp_path / "one.pdf"
    second = tmp_path / "two.pdf"
    first.write_bytes(b"1")
    second.write_bytes(b"2")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()

    def _tagging(*_args, **_kwargs):
        key = _kwargs.get("Key", "")
        if key.endswith("one.pdf") or "one.pdf" in key:
            return {
                "TagSet": [
                    {
                        "Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY,
                        "Value": "NO_THREATS_FOUND",
                    }
                ]
            }
        return {
            "TagSet": [
                {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "THREATS_FOUND"}
            ]
        }

    s3.get_object_tagging.side_effect = _tagging

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(gr.Error):
            ms.scan_gradio_file_upload([str(first), str(second)])

    assert s3.delete_object.call_count == 2
    assert not second.exists()
    assert first.exists()


def test_scan_gradio_file_upload_invalidates_stale_cache(monkeypatch, tmp_path):
    sample = tmp_path / "cached.pdf"
    sample.write_bytes(b"cached")
    ms.mark_scanned_clean(str(sample))

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [
            {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "NO_THREATS_FOUND"}
        ]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        ms.scan_gradio_file_upload(str(sample))

    s3.upload_file.assert_called_once()


def test_scan_gradio_file_upload_shows_checking_info(monkeypatch, tmp_path):
    sample = tmp_path / "scan.pdf"
    sample.write_bytes(b"scan")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_SHOW_CHECKING_INFO", True, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [
            {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "NO_THREATS_FOUND"}
        ]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with patch("tools.malware_scan.gr.Info") as mock_info:
            with patch("builtins.print") as mock_print:
                ms.scan_gradio_file_upload(str(sample))
                mock_print.assert_any_call(
                    ms.MALWARE_SCAN_CHECKING_INFO_MESSAGE, flush=True
                )
                mock_print.assert_any_call(
                    ms.MALWARE_SCAN_SUCCESS_INFO_MESSAGE, flush=True
                )
                mock_info.assert_any_call(ms.MALWARE_SCAN_CHECKING_INFO_MESSAGE)
                mock_info.assert_any_call(ms.MALWARE_SCAN_SUCCESS_INFO_MESSAGE)
                assert mock_info.call_count == 2


def test_scan_gradio_file_upload_shows_success_info(monkeypatch, tmp_path):
    sample = tmp_path / "scan.pdf"
    sample.write_bytes(b"scan")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_SHOW_CHECKING_INFO", True, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [
            {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "NO_THREATS_FOUND"}
        ]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with patch("tools.malware_scan.gr.Info") as mock_info:
            ms.scan_gradio_file_upload(str(sample))
            assert mock_info.call_args_list == [
                ((ms.MALWARE_SCAN_CHECKING_INFO_MESSAGE,),),
                ((ms.MALWARE_SCAN_SUCCESS_INFO_MESSAGE,),),
            ]


def test_scan_gradio_file_upload_hides_checking_info_when_disabled(
    monkeypatch, tmp_path
):
    sample = tmp_path / "scan.pdf"
    sample.write_bytes(b"scan")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_SHOW_CHECKING_INFO", False, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_POLL_INTERVAL_SEC", 0.01, raising=False)
    monkeypatch.setattr(ms, "MALWARE_SCAN_TIMEOUT_SEC", 1.0, raising=False)

    s3 = MagicMock()
    s3.get_object_tagging.return_value = {
        "TagSet": [
            {"Key": ms.GUARDDUTY_MALWARE_SCAN_TAG_KEY, "Value": "NO_THREATS_FOUND"}
        ]
    }

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with patch("tools.malware_scan.gr.Info") as mock_info:
            ms.scan_gradio_file_upload(str(sample))
            mock_info.assert_not_called()


def test_scan_gradio_file_upload_s3_upload_failure_shows_error(monkeypatch, tmp_path):
    sample = tmp_path / "denied.pdf"
    sample.write_bytes(b"data")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)

    s3 = MagicMock()
    s3.upload_file.side_effect = ClientError(
        {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}},
        "PutObject",
    )

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(gr.Error, match="configuration or permissions"):
            ms.scan_gradio_file_upload(str(sample))

    assert not sample.exists()


def test_require_files_malware_scanned_blocks_unclean(monkeypatch, tmp_path):
    sample = tmp_path / "unclean.pdf"
    sample.write_bytes(b"x")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    with patch("tools.malware_scan.scan_gradio_file_upload"):
        with pytest.raises(gr.Error, match="must pass malware scanning"):
            ms.require_files_malware_scanned(str(sample))


def test_require_files_malware_scanned_scans_unclean_paths(monkeypatch, tmp_path):
    sample = tmp_path / "unclean.pdf"
    sample.write_bytes(b"x")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    def _scan_and_mark_clean(file_input):
        for path in ms.normalize_gradio_file_paths(file_input):
            ms.mark_scanned_clean(path)

    with patch(
        "tools.malware_scan.scan_gradio_file_upload", side_effect=_scan_and_mark_clean
    ):
        ms.require_files_malware_scanned(str(sample))


def test_scan_local_file_s3_upload_failure_raises_service_error(monkeypatch, tmp_path):
    sample = tmp_path / "denied.pdf"
    sample.write_bytes(b"data")

    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    monkeypatch.setattr(ms, "MALWARE_SCAN_S3_BUCKET", "scan-bucket", raising=False)

    s3 = MagicMock()
    s3.upload_file.side_effect = ClientError(
        {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}},
        "PutObject",
    )

    with patch("tools.malware_scan.boto3.client", return_value=s3):
        with pytest.raises(
            ms.MalwareScanServiceError, match="configuration or permissions"
        ):
            ms.scan_local_file_for_malware(str(sample))

    s3.delete_object.assert_not_called()
    assert sample.exists()


def test_read_scan_status_access_denied_raises_service_error():
    s3 = MagicMock()
    s3.get_object_tagging.side_effect = ClientError(
        {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}},
        "GetObjectTagging",
    )
    with pytest.raises(
        ms.MalwareScanServiceError, match="configuration or permissions"
    ):
        ms._read_scan_status(s3, "scan-bucket", "key.pdf")


def test_make_malware_scan_disable_outputs_noop_when_disabled(monkeypatch):
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: False)
    disable = ms.make_malware_scan_disable_outputs(1)
    assert disable() == gr.update()
    disable_two = ms.make_malware_scan_disable_outputs(2)
    assert disable_two() == (gr.update(), gr.update())


def test_make_malware_scan_disable_outputs_disables_buttons(monkeypatch):
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)
    disable = ms.make_malware_scan_disable_outputs(1)
    assert disable() == gr.update(interactive=False)
    disable_two = ms.make_malware_scan_disable_outputs(2)
    assert disable_two() == (
        gr.update(interactive=False),
        gr.update(interactive=False),
    )


def test_make_malware_scan_enable_outputs_single_button():
    enable = ms.make_malware_scan_enable_outputs(1)
    assert enable() == gr.update(interactive=True)


def test_make_malware_scan_upload_failure_outputs_clears_file_and_enables():
    handler = ms.make_malware_scan_upload_failure_outputs(1)
    assert handler() == (gr.update(value=None), gr.update(interactive=True))


def test_trusted_bundled_example_skips_malware_scan(monkeypatch, tmp_path):
    example_dir = tmp_path / "example_data"
    example_dir.mkdir()
    sample = example_dir / "demo.pdf"
    sample.write_bytes(b"%PDF-1.4")

    monkeypatch.setattr(
        "tools.example_data_paths.resolve_example_data_dirs",
        lambda: [example_dir.resolve()],
    )
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    from tools.example_data_paths import (
        is_bundled_example_file,
        is_trusted_bundled_example_path,
    )

    assert is_trusted_bundled_example_path(str(sample)) is True
    assert is_bundled_example_file(str(sample)) is True
    assert ms.path_is_malware_clean(str(sample)) is True

    with patch("tools.malware_scan.boto3.client") as mock_client:
        ms.scan_gradio_file_upload(str(sample))
        mock_client.assert_not_called()


def test_gradio_temp_copy_of_bundled_example_skips_malware_scan(monkeypatch, tmp_path):
    example_dir = tmp_path / "example_data"
    example_dir.mkdir()
    bundled = example_dir / "demo.pdf"
    bundled.write_bytes(b"%PDF-1.4 example")
    gradio_copy = tmp_path / "gradio_tmp" / "abc123_demo.pdf"
    gradio_copy.parent.mkdir()
    gradio_copy.write_bytes(b"%PDF-1.4 example")

    monkeypatch.setattr(
        "tools.example_data_paths.resolve_example_data_dirs",
        lambda: [example_dir.resolve()],
    )
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    from tools.example_data_paths import is_bundled_example_file

    assert is_bundled_example_file(str(gradio_copy)) is True
    assert ms.path_is_malware_clean(str(gradio_copy)) is True

    with patch("tools.malware_scan.boto3.client") as mock_client:
        ms.scan_gradio_file_upload(str(gradio_copy))
        mock_client.assert_not_called()


def test_require_files_malware_scanned_allows_bundled_example(monkeypatch, tmp_path):
    example_dir = tmp_path / "example_data"
    example_dir.mkdir()
    sample = example_dir / "demo.pdf"
    sample.write_bytes(b"x")

    monkeypatch.setattr(
        "tools.example_data_paths.resolve_example_data_dirs",
        lambda: [example_dir.resolve()],
    )
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    ms.require_files_malware_scanned(str(sample))


def test_ensure_upload_scanned_for_malware_skips_when_cached(monkeypatch, tmp_path):
    sample = tmp_path / "cached.pdf"
    sample.write_bytes(b"cached")
    ms.mark_scanned_clean(str(sample))
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    with patch("tools.malware_scan.scan_local_file_for_malware") as mock_scan:
        ms.ensure_upload_scanned_for_malware(str(sample))
        mock_scan.assert_not_called()


def test_normalize_gradio_file_paths_listfiles(tmp_path):
    from gradio.data_classes import FileData, ListFiles

    sample = tmp_path / "doc.pdf"
    sample.write_bytes(b"%PDF-1.4")
    lf = ListFiles(root=[FileData(path=str(sample), orig_name="doc.pdf")])
    assert ms.normalize_gradio_file_paths(lf) == [os.path.abspath(str(sample))]


def test_clear_scan_cache_for_path_drops_all_mtimes(tmp_path):
    sample = tmp_path / "doc.pdf"
    sample.write_bytes(b"v1")
    ms.mark_scanned_clean(str(sample))
    assert ms.already_scanned_clean(str(sample))

    ms.clear_scan_cache_for_path(str(sample))
    assert not ms.already_scanned_clean(str(sample))


def test_handle_gradio_file_deleted_clears_cache(monkeypatch, tmp_path):
    sample = tmp_path / "removed.pdf"
    sample.write_bytes(b"removed")
    ms.mark_scanned_clean(str(sample))
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    class _Deleted:
        file = type("F", (), {"path": str(sample), "name": None})()

    ms.handle_gradio_file_deleted(_Deleted())
    assert not ms.already_scanned_clean(str(sample))


def test_make_malware_scan_upload_start_scans_then_disables(monkeypatch, tmp_path):
    sample = tmp_path / "upload.pdf"
    sample.write_bytes(b"upload")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    with patch("tools.malware_scan.scan_gradio_file_upload") as mock_scan:
        start = ms.make_malware_scan_upload_start(1)
        result = start(str(sample))
        mock_scan.assert_called_once_with(str(sample))
        assert result == gr.update(interactive=False)


def test_scan_gradio_file_upload_skips_bundled_example(monkeypatch, tmp_path, caplog):
    import logging

    example_dir = tmp_path / "example_data"
    example_dir.mkdir()
    sample = example_dir / "demo.pdf"
    sample.write_bytes(b"%PDF-1.4 example")

    monkeypatch.setattr(
        "tools.example_data_paths.resolve_example_data_dirs",
        lambda: [example_dir.resolve()],
    )
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    with caplog.at_level(logging.INFO):
        with patch("tools.malware_scan.boto3.client") as mock_client:
            ms.scan_gradio_file_upload(str(sample))
            mock_client.assert_not_called()

    assert "Skipping malware scan for 1 already-clean path(s)" in caplog.text


def test_make_malware_scan_upload_failure_outputs_zero_buttons_clears_file():
    handler = ms.make_malware_scan_upload_failure_outputs(0)
    assert handler() == gr.update(value=None)


def test_bind_malware_scan_upload_wires_chain():
    file_input = MagicMock()
    button = MagicMock()
    upload_event = MagicMock()
    success_event = MagicMock()
    file_input.upload.return_value = upload_event
    upload_event.success.return_value = success_event

    ms.bind_malware_scan_upload(file_input, button)

    file_input.upload.assert_called_once()
    upload_event.success.assert_called_once()
    success_event.failure.assert_called_once()
    failure_kwargs = success_event.failure.call_args.kwargs
    assert failure_kwargs["outputs"] == [file_input, button]


def test_bind_malware_scan_upload_requires_button():
    with pytest.raises(ValueError, match="at least one button"):
        ms.bind_malware_scan_upload(MagicMock(), [])


def test_mark_app_generated_files_malware_clean(monkeypatch, tmp_path):
    sample = tmp_path / "ocr_output.csv"
    sample.write_text("page,text\n1,hello\n")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: True)

    assert not ms.already_scanned_clean(str(sample))
    ms.mark_app_generated_files_malware_clean([str(sample)])
    assert ms.already_scanned_clean(str(sample))


def test_mark_app_generated_files_malware_clean_noop_when_disabled(
    monkeypatch, tmp_path
):
    sample = tmp_path / "ocr_output.csv"
    sample.write_text("page,text\n")
    monkeypatch.setattr(ms, "malware_scan_enabled", lambda: False)
    ms.mark_app_generated_files_malware_clean([str(sample)])
    assert not ms.already_scanned_clean(str(sample))


def test_custom_regex_load_requires_malware_scan(tmp_path):
    sample = tmp_path / "allow.csv"
    sample.write_text("term_one\n")

    class _File:
        name = str(sample)

    from tools.helper_functions import custom_regex_load

    with patch("tools.malware_scan.require_files_malware_scanned") as mock_req:
        text, values = custom_regex_load([_File()], "allow_list")
        mock_req.assert_called_once()
    assert "allow list file loaded" in text
    assert values == ["term_one"]


def test_custom_regex_load_skips_require_when_empty():
    from tools.helper_functions import custom_regex_load

    with patch("tools.malware_scan.require_files_malware_scanned") as mock_req:
        text, values = custom_regex_load([], "allow_list")
        mock_req.assert_not_called()
    assert text == "No file provided."
    assert values == []


def test_run_duplicate_analysis_requires_malware_scan(tmp_path):
    import pandas as pd

    sample = tmp_path / "ocr.csv"
    sample.write_text("page,text\n1,hello\n")
    try:
        from tools.find_duplicate_pages import run_duplicate_analysis
    except ModuleNotFoundError as exc:
        pytest.skip(f"NLP optional deps missing: {exc}")

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            run_duplicate_analysis(
                [str(sample)],
                0.95,
                10,
                1,
                True,
                pd.DataFrame(),
                [],
            )


def test_merge_csv_files_requires_malware_scan(tmp_path):
    sample = tmp_path / "review.csv"
    sample.write_text("page,label,color,xmin,ymin,xmax,ymax\n1,a,#000,0,0,1,1\n")
    from tools.helper_functions import merge_csv_files

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            merge_csv_files([str(sample)], output_folder=str(tmp_path) + os.sep)


def test_combine_review_pdf_files_requires_malware_scan(tmp_path):
    sample = tmp_path / "doc_redactions_for_review.pdf"
    sample.write_bytes(b"%PDF-1.4")
    from tools.file_conversion import combine_review_pdf_files

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            combine_review_pdf_files(
                [str(sample)], output_folder=str(tmp_path) + os.sep
            )


def test_anonymise_files_requires_malware_scan(tmp_path):
    sample = tmp_path / "data.csv"
    sample.write_text("col\nvalue\n")
    try:
        from tools.data_anonymise import anonymise_files_with_open_text
    except ModuleNotFoundError as exc:
        pytest.skip(f"NLP optional deps missing: {exc}")

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            anonymise_files_with_open_text(
                [str(sample)],
                "",
                "replace with 'REDACTED'",
                ["col"],
                ["PERSON"],
            )


def test_run_tabular_duplicate_detection_requires_malware_scan(tmp_path):
    sample = tmp_path / "data.csv"
    sample.write_text("col\nvalue\n")
    try:
        from tools.find_duplicate_tabular import run_tabular_duplicate_detection
    except ModuleNotFoundError as exc:
        pytest.skip(f"NLP optional deps missing: {exc}")

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            run_tabular_duplicate_detection(
                [str(sample)], 0.9, 1, ["col"], output_folder=str(tmp_path) + os.sep
            )


def test_summarise_document_wrapper_requires_malware_scan(tmp_path):
    import pandas as pd

    sample = tmp_path / "ocr.csv"
    sample.write_text("page,line,text\n1,1,hello\n")
    from tools.summaries import summarise_document_wrapper

    with patch(
        "tools.malware_scan.require_files_malware_scanned",
        side_effect=gr.Error("blocked"),
    ):
        with pytest.raises(gr.Error, match="blocked"):
            summarise_document_wrapper(
                pd.DataFrame(),
                str(tmp_path),
                "local",
                "",
                0.6,
                "doc",
                "",
                "",
                "",
                "",
                "",
                "bullets",
                "",
                10,
                [str(sample)],
            )