Oluwarotimi Williams Samuel

"""
Unit tests for --force flag, --modules selection, --analyzed mode, and --rerun mode.

Tests that:
1. --force bypasses both bulk and per-file dedup checks (feature-extraction AND decompile paths)
2. --modules does NOT implicitly enable force (force must be explicit)
3. Default behavior (no --force) preserves dedup checks
4. --analyzed queries basic_properties and handles decompile dedup filtering
5. --rerun queries disassembled table directly (with --analyzed --decompile)
6. Decompile per-file is_in_code_db respects --force flag
7. YARA scanning skips already-scanned samples by default, rescans all with --force
"""
import os
import sys
import pytest
from unittest.mock import Mock, patch, MagicMock

pytestmark = [pytest.mark.unit]


def _mock_decompiler_module():
    """Install a mock for redb.extractors.decompiler into sys.modules.

    The real DecompileBinja import chain pulls in blake3/binaryninja which
    aren't available in the unit-test environment. This helper injects a
    lightweight mock so that `from redb.extractors.decompiler import DecompileBinja`
    succeeds inside process_binary_file.

    Returns a context-manager-compatible mock DecompileBinja class.
    """
    mock_decompile_binja = Mock()
    mock_instance = Mock()
    mock_instance.export_data.return_value = True
    mock_instance.analysis_results = None
    mock_decompile_binja.return_value.__enter__ = Mock(return_value=mock_instance)
    mock_decompile_binja.return_value.__exit__ = Mock(return_value=False)

    mock_mod = MagicMock()
    mock_mod.DecompileBinja = mock_decompile_binja
    return mock_mod, mock_decompile_binja


# ============================================================================
# CLI Argument Parsing Tests
# ============================================================================

class TestForceCliArgument:
    """Tests for --force CLI argument parsing in start.py."""

    def _parse_args(self, args_list):
        """Helper to parse CLI args without running main()."""
        import argparse

        # Replicate the parser from start.py
        parser = argparse.ArgumentParser()
        input_group = parser.add_mutually_exclusive_group(required=True)
        input_group.add_argument("--path")
        input_group.add_argument("--s3", action="store_true")
        input_group.add_argument("--s3-solo", metavar="S3_KEY")
        input_group.add_argument("--nomad-job", action="store_true")
        input_group.add_argument("--date", metavar="YYYY-MM-DD")
        input_group.add_argument("--range", nargs=2)
        input_group.add_argument("--analyzed", action="store_true")

        parser.add_argument("--repo")
        parser.add_argument("--index_prefix", default="redb")
        parser.add_argument("-d", "--decompile", action="store_true")
        parser.add_argument("-m", "--modules", default="all")
        parser.add_argument("--force", action="store_true")
        parser.add_argument("--dry-run", action="store_true")

        return parser.parse_args(args_list)

    def test_force_flag_default_false(self):
        """--force is False by default."""
        args = self._parse_args(["--path", "/tmp/test", "--repo", "test"])
        assert args.force is False

    def test_force_flag_explicit(self):
        """--force is True when specified."""
        args = self._parse_args(["--path", "/tmp/test", "--repo", "test", "--force"])
        assert args.force is True

    def test_modules_all_no_implicit_force(self):
        """--modules 'all' does not enable force."""
        args = self._parse_args(["--path", "/tmp/test", "--repo", "test", "--modules", "all"])
        assert args.force is False

    def test_specific_modules_no_implicit_force(self):
        """--modules with specific extractors does NOT implicitly enable force."""
        args = self._parse_args([
            "--path", "/tmp/test", "--repo", "test",
            "--modules", "MachOFeaturesExtractor"
        ])
        assert args.force is False

    def test_specific_modules_with_explicit_force(self):
        """--modules with specific extractors requires explicit --force to bypass dedup."""
        args = self._parse_args([
            "--path", "/tmp/test", "--repo", "test",
            "--modules", "BasicPropertiesExtractor,HashExtractor", "--force"
        ])
        assert args.force is True

    def test_force_with_modules_all(self):
        """--force works independently of --modules."""
        args = self._parse_args([
            "--path", "/tmp/test", "--repo", "test",
            "--modules", "all", "--force"
        ])
        assert args.force is True


# ============================================================================
# is_in_db Per-File Dedup Bypass Tests (using module-level patching)
# ============================================================================

class TestPerFileDedupBypass:
    """Tests for the per-file is_in_db check bypass in process_binary_file.

    We patch at the module level to avoid triggering heavy import chains.
    """

    @patch('redb.workers.is_in_db')
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_default_calls_is_in_db(self, mock_settings, mock_magika, mock_hashlib, mock_is_in_db):
        """Without force, is_in_db is called and can skip files."""
        from redb.ingestor import process_binary_file, ImportResult

        # Setup magika
        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance

        # Setup hashlib
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64

        # Setup settings for ClickHouse
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        # is_in_db returns True => file should be skipped
        mock_is_in_db.return_value = True

        logger = Mock()

        with patch('builtins.open', mock_open_binary(b"test data")):
            result, filetype = process_binary_file(
                "/tmp/test.bin", False, "redb", logger,
                selected_modules="all", force=False
            )

        mock_is_in_db.assert_called_once_with("a" * 64, "redb")
        assert result == ImportResult.SKIPPED

    @patch('redb.workers.is_in_db')
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    @patch('redb.workers._is_packed', return_value=None)
    @patch('redb.workers.DIEExtractor')
    @patch('redb.workers.BasicPropertiesExtractor')
    @patch('redb.workers.HashExtractor')
    @patch('redb.workers.CAPAExtractor')
    def test_force_skips_is_in_db(self, mock_capa, mock_hash, mock_bp, mock_die, mock_packed,
                                   mock_settings, mock_magika, mock_hashlib, mock_is_in_db):
        """With force=True, is_in_db is never called.

        Uses filetype 'unknown' to avoid triggering PE/ELF/Mach-O specific
        extractor chains which would need many more mocks.
        """
        from redb.ingestor import process_binary_file

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "unknown"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        logger = Mock()

        with patch('builtins.open', mock_open_binary(b"test data")):
            try:
                process_binary_file(
                    "/tmp/test.bin", False, "redb", logger,
                    selected_modules="all", force=True
                )
            except Exception:
                pass

        mock_is_in_db.assert_not_called()

    @patch('redb.workers.is_in_db')
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_decompile_mode_skips_is_in_db(self, mock_settings, mock_magika, mock_hashlib, mock_is_in_db):
        """In decompile mode, is_in_db is never called (existing behavior)."""
        from redb.ingestor import process_binary_file

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        logger = Mock()

        with patch('builtins.open', mock_open_binary(b"test data")):
            try:
                process_binary_file(
                    "/tmp/test.bin", True, "redb", logger,
                    selected_modules="all", force=False
                )
            except Exception:
                pass

        mock_is_in_db.assert_not_called()


# ============================================================================
# Bulk Dedup Bypass Tests
# ============================================================================

class TestBulkDedupBypass:
    """Tests for the bulk is_in_db_bulk check bypass in fetch functions."""

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_repo_default_calls_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """Without force, fetch_s3_objects_by_repository calls is_in_db_bulk."""
        from redb.ingestor import fetch_s3_objects_by_repository

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None)
        ]
        mock_db_conn.return_value = mock_client
        mock_bulk.return_value = set()

        fetch_s3_objects_by_repository("test_repo", "redb", False, force=False)

        mock_bulk.assert_called_once()

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_repo_force_skips_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """With force=True, fetch_s3_objects_by_repository skips is_in_db_bulk."""
        from redb.ingestor import fetch_s3_objects_by_repository

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None)
        ]
        mock_db_conn.return_value = mock_client

        result = fetch_s3_objects_by_repository("test_repo", "redb", False, force=True)

        mock_bulk.assert_not_called()
        assert len(result) == 1  # All rows kept, none filtered

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_repo_force_preserves_all_rows(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """With force=True, all rows from the query are returned unfiltered."""
        from redb.ingestor import fetch_s3_objects_by_repository

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin", None),
            ("hash2" + "0" * 59, "pebin", None),
            ("hash3" + "0" * 59, "pebin", None),
        ]
        mock_db_conn.return_value = mock_client

        result = fetch_s3_objects_by_repository("test_repo", "redb", False, force=True)

        assert len(result) == 3

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_date_default_calls_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """Without force, fetch_s3_objects_by_date_range calls is_in_db_bulk."""
        from redb.ingestor import fetch_s3_objects_by_date_range

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None)
        ]
        mock_db_conn.return_value = mock_client
        mock_bulk.return_value = set()

        fetch_s3_objects_by_date_range(
            "redb", False, "2025-01-01", "2025-01-02", force=False
        )

        mock_bulk.assert_called_once()

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_date_force_skips_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """With force=True, fetch_s3_objects_by_date_range skips is_in_db_bulk."""
        from redb.ingestor import fetch_s3_objects_by_date_range

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None)
        ]
        mock_db_conn.return_value = mock_client

        result = fetch_s3_objects_by_date_range(
            "redb", False, "2025-01-01", "2025-01-02", force=True
        )

        mock_bulk.assert_not_called()
        assert len(result) == 1


# ============================================================================
# Ingestor Force Parameter Tests
# ============================================================================

class TestIngestorForceParam:
    """Tests for force parameter in Ingestor class."""

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_default_force_false(self, mock_open, mock_datetime, mock_mp):
        """Ingestor defaults force to False."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(path="/tmp/test", repository="test", index_prefix="redb")
        assert ingestor.force is False

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_force_true(self, mock_open, mock_datetime, mock_mp):
        """Ingestor stores force=True when passed."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(path="/tmp/test", repository="test", index_prefix="redb", force=True)
        assert ingestor.force is True


# ============================================================================
# Function Chain Threading Tests
# ============================================================================

class TestForceThreading:
    """Tests that force parameter is correctly threaded through the call chain."""

    @patch('redb.workers.process_binary_file')
    @patch('redb.workers.is_binary_file', return_value=True)
    def test_process_file_internal_passes_force_true(self, mock_is_binary, mock_process):
        """_process_file_internal passes force=True to process_binary_file."""
        from redb.ingestor import _process_file_internal

        mock_process.return_value = (Mock(), "pebin")
        logger = Mock()

        _process_file_internal("/tmp/test.bin", False, "redb", logger, force=True)

        mock_process.assert_called_once()
        call_args = mock_process.call_args
        assert call_args[0][-1] is True

    @patch('redb.workers.process_binary_file')
    @patch('redb.workers.is_binary_file', return_value=True)
    def test_process_file_internal_default_force_false(self, mock_is_binary, mock_process):
        """_process_file_internal defaults force to False."""
        from redb.ingestor import _process_file_internal

        mock_process.return_value = (Mock(), "pebin")
        logger = Mock()

        _process_file_internal("/tmp/test.bin", False, "redb", logger)

        call_args = mock_process.call_args
        assert call_args[0][-1] is False

    @patch('redb.workers._process_file_internal')
    def test_process_file_passes_force(self, mock_internal):
        """process_file passes force to _process_file_internal."""
        from redb.workers import process_file

        mock_internal.return_value = (Mock(), "pebin")

        with patch('redb.workers.setup_direct_logger', return_value=Mock()):
            process_file("/tmp/test.bin", False, "redb", "/tmp/log", 1, 1, force=True)

        mock_internal.assert_called_once()
        call_args = mock_internal.call_args
        assert call_args[0][-1] is True

    @patch('redb.workers._process_file_internal')
    def test_process_s3_file_passes_force(self, mock_internal):
        """process_s3_file passes force to _process_file_internal."""
        from redb.workers import process_s3_file

        mock_internal.return_value = (Mock(), "pebin")

        with patch('redb.workers.setup_direct_logger', return_value=Mock()), \
             patch('redb.workers.download_s3_object', return_value="/tmp/downloaded.bin"):
            process_s3_file(
                "bucket", "key", "/tmp", False, "redb", "/tmp/log", 1, 1,
                force=True
            )

        mock_internal.assert_called_once()
        call_args = mock_internal.call_args
        assert call_args[0][-1] is True

    @patch('redb.workers.process_binary_file')
    def test_process_zip_passes_force(self, mock_process):
        """process_zip_file passes force to process_binary_file."""
        from redb.workers import process_zip_file

        mock_process.return_value = (Mock(), "pebin")
        logger = Mock()

        import tempfile
        import zipfile

        # Create a real zip file with a dummy binary inside
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            with zipfile.ZipFile(tmp_path, 'w') as zf:
                zf.writestr("test.bin", b"\x00" * 100)

            process_zip_file(tmp_path, False, "redb", logger, "all", force=True)

            mock_process.assert_called_once()
            call_args = mock_process.call_args
            assert call_args[0][-1] is True
        finally:
            os.unlink(tmp_path)

    @patch('redb.workers.process_binary_file')
    def test_process_7zip_passes_force(self, mock_process):
        """process_7zip_file passes force to process_binary_file."""
        from redb.workers import process_7zip_file

        mock_process.return_value = (Mock(), "pebin")
        logger = Mock()

        with patch('redb.workers.py7zr') as mock_py7zr:
            mock_zf = Mock()
            mock_py7zr.SevenZipFile.return_value.__enter__ = Mock(return_value=mock_zf)
            mock_py7zr.SevenZipFile.return_value.__exit__ = Mock(return_value=False)

            import tempfile

            def fake_extractall(path):
                dest = os.path.join(path, "test.bin")
                with open(dest, "wb") as f:
                    f.write(b"\x00" * 100)

            mock_zf.extractall = fake_extractall

            process_7zip_file("/tmp/test.7z", False, "redb", logger, "all", force=True)

            if mock_process.called:
                call_args = mock_process.call_args
                assert call_args[0][-1] is True


# ============================================================================
# --analyzed CLI Argument Tests
# ============================================================================

class TestAnalyzedCliArgument:
    """Tests for --analyzed CLI argument parsing in start.py."""

    def _parse_args(self, args_list):
        """Helper to parse CLI args without running main()."""
        import argparse

        parser = argparse.ArgumentParser()
        input_group = parser.add_mutually_exclusive_group(required=True)
        input_group.add_argument("--path")
        input_group.add_argument("--s3", action="store_true")
        input_group.add_argument("--s3-solo", metavar="S3_KEY")
        input_group.add_argument("--nomad-job", action="store_true")
        input_group.add_argument("--date", metavar="YYYY-MM-DD")
        input_group.add_argument("--range", nargs=2)
        input_group.add_argument("--analyzed", action="store_true")

        parser.add_argument("--repo")
        parser.add_argument("--index_prefix", default="redb")
        parser.add_argument("-d", "--decompile", action="store_true")
        parser.add_argument("-m", "--modules", default="all")
        parser.add_argument("--force", action="store_true")
        parser.add_argument("--dry-run", action="store_true")
        parser.add_argument("--magika")

        return parser.parse_args(args_list)

    def test_analyzed_flag_default_false(self):
        """--analyzed is False by default."""
        args = self._parse_args(["--path", "/tmp/test", "--repo", "test"])
        assert args.analyzed is False

    def test_analyzed_flag_explicit(self):
        """--analyzed is True when specified."""
        args = self._parse_args(["--analyzed"])
        assert args.analyzed is True

    def test_analyzed_mutually_exclusive_with_path(self):
        """--analyzed and --path cannot be used together."""
        import argparse
        with pytest.raises(SystemExit):
            self._parse_args(["--analyzed", "--path", "/tmp/test"])

    def test_analyzed_mutually_exclusive_with_s3(self):
        """--analyzed and --s3 cannot be used together."""
        import argparse
        with pytest.raises(SystemExit):
            self._parse_args(["--analyzed", "--s3"])

    def test_analyzed_with_decompile(self):
        """--analyzed can be combined with --decompile."""
        args = self._parse_args(["--analyzed", "--decompile"])
        assert args.analyzed is True
        assert args.decompile is True

    def test_analyzed_with_modules(self):
        """--analyzed can be combined with --modules."""
        args = self._parse_args(["--analyzed", "--modules", "HashExtractor"])
        assert args.analyzed is True
        assert args.modules == "HashExtractor"

    def test_analyzed_with_magika(self):
        """--analyzed can be combined with --magika."""
        args = self._parse_args(["--analyzed", "--magika", "elf"])
        assert args.analyzed is True
        assert args.magika == "elf"

    def test_analyzed_with_force(self):
        """--analyzed can be combined with --force."""
        args = self._parse_args(["--analyzed", "--force"])
        assert args.analyzed is True
        assert args.force is True

    def test_analyzed_no_repo_required(self):
        """--analyzed does not require --repo."""
        args = self._parse_args(["--analyzed"])
        assert args.analyzed is True
        assert args.repo is None


# ============================================================================
# fetch_analyzed_samples Tests
# ============================================================================

class TestFetchAnalyzedSamples:
    """Tests for fetch_analyzed_samples() in queries.py."""

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_returns_analyzed_samples(self, mock_settings, mock_formats, mock_gen_key):
        """fetch_analyzed_samples returns samples from basic_properties."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin"),
            ("def456" + "0" * 58, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False)

        assert len(result) == 2
        assert all(r['s3_bucket'] == 'test-bucket' for r in result)
        assert all(r['s3_key'] == 'ab/cd/abcd.zip' for r in result)

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_queries_basic_properties_table(self, mock_settings, mock_formats, mock_gen_key):
        """fetch_analyzed_samples queries the correct table."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        fetch_analyzed_samples("myprefix", decompile=False)

        query_arg = mock_client.query.call_args[0][0]
        assert "myprefix_basic_properties" in query_arg
        assert "FINAL" in query_arg

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_magika_filter_applied(self, mock_settings, mock_gen_key):
        """fetch_analyzed_samples applies magika_filter when provided."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "elf"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False, magika_filter="elf")

        query_arg = mock_client.query.call_args[0][0]
        assert "filetype_magika = %(magika)s" in query_arg
        params = mock_client.query.call_args[1].get('parameters') or mock_client.query.call_args[0][1] if len(mock_client.query.call_args[0]) > 1 else mock_client.query.call_args[1].get('parameters', {})
        assert params.get("magika") == "elf"

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_decompile_mode_filters_already_decompiled(self, mock_settings, mock_formats, mock_gen_key):
        """In decompile mode, already-decompiled samples are filtered out."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        # First query: basic_properties returns 3 samples
        # Second query: decompiled table returns 1 already-decompiled
        mock_client.query.side_effect = [
            Mock(result_rows=[
                ("hash1" + "0" * 59, "pebin"),
                ("hash2" + "0" * 59, "pebin"),
                ("hash3" + "0" * 59, "pebin"),
            ]),
            Mock(result_rows=[
                ("hash2" + "0" * 59,),  # already decompiled
            ]),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True)

        assert len(result) == 2
        result_hashes = {r['sha256'] for r in result}
        assert "hash2" + "0" * 59 not in result_hashes

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_decompile_force_skips_decompile_check(self, mock_settings, mock_formats, mock_gen_key):
        """With force=True in decompile mode, decompile check is skipped."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin"),
            ("hash2" + "0" * 59, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, force=True)

        # Should only query once (basic_properties), no decompile check
        assert mock_client.query.call_count == 1
        assert len(result) == 2

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_analysis_mode_no_decompile_filtering(self, mock_settings, mock_formats, mock_gen_key):
        """In analysis mode (decompile=False), no decompile filtering is done."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin"),
            ("hash2" + "0" * 59, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False)

        # Only one query (basic_properties), no decompile table check
        assert mock_client.query.call_count == 1
        assert len(result) == 2

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    def test_missing_s3_bucket_returns_empty(self, mock_settings, mock_formats, mock_gen_key):
        """Returns empty list when S3_BUCKET env var is not set."""
        from redb.queries import fetch_analyzed_samples

        # Ensure S3_BUCKET is not set
        with patch.dict(os.environ, {}, clear=True):
            result = fetch_analyzed_samples("redb", decompile=False)

        assert result == []

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_empty_basic_properties_returns_empty(self, mock_settings, mock_formats, mock_gen_key):
        """Returns empty list when no analyzed samples exist."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False)

        assert result == []

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_handles_bytes_sha256(self, mock_settings, mock_formats, mock_gen_key):
        """Handles bytes-encoded SHA256 from ClickHouse."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        hash_str = "abc123" + "0" * 58
        mock_client.query.return_value.result_rows = [
            (hash_str.encode('utf-8'), "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False)

        assert len(result) == 1
        assert result[0]['sha256'] == hash_str

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_exception_returns_empty(self, mock_settings, mock_formats, mock_gen_key):
        """Returns empty list on database error."""
        from redb.queries import fetch_analyzed_samples

        mock_settings.create_clickhouse_client.side_effect = Exception("Connection failed")

        result = fetch_analyzed_samples("redb", decompile=False)

        assert result == []

    @patch('redb.queries._check_yara_matches_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_yara_scan_skips_decompile_filtering(self, mock_settings, mock_formats, mock_gen_key, mock_yara_check):
        """With yara_scan=True, decompile filtering is skipped (YARA check runs instead)."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client
        mock_yara_check.return_value = set()  # No already-scanned samples

        result = fetch_analyzed_samples("redb", decompile=True, yara_scan=True)

        # basic_properties queried once, no decompile table check
        assert mock_client.query.call_count == 1
        # YARA check was used instead of decompile check
        mock_yara_check.assert_called_once()
        assert len(result) == 1

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket", "CLICKHOUSE_CODE_PREFIX": "custom_code"})
    def test_decompile_uses_code_prefix_env(self, mock_settings, mock_formats, mock_gen_key):
        """Decompile filtering uses CLICKHOUSE_CODE_PREFIX env var."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.side_effect = [
            Mock(result_rows=[("hash1" + "0" * 59, "pebin")]),
            Mock(result_rows=[]),  # No disassembled samples
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        fetch_analyzed_samples("redb", decompile=True)

        # Second query should reference custom_code prefix with disassembled table
        second_query = mock_client.query.call_args_list[1][0][0]
        assert "custom_code_disassembled_functions_references" in second_query


# ============================================================================
# Ingestor Analyzed Parameter Tests
# ============================================================================

class TestIngestorAnalyzedParam:
    """Tests for analyzed parameter in Ingestor class."""

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_default_analyzed_false(self, mock_open, mock_datetime, mock_mp):
        """Ingestor defaults analyzed to False."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(path="/tmp/test", repository="test", index_prefix="redb")
        assert ingestor.analyzed is False

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_analyzed_true(self, mock_open, mock_datetime, mock_mp):
        """Ingestor stores analyzed=True when passed."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(
            path=None, repository="test", index_prefix="redb",
            s3_mode=True, analyzed=True
        )
        assert ingestor.analyzed is True

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_analyzed_with_decompile(self, mock_open, mock_datetime, mock_mp):
        """Ingestor stores both analyzed and decompile flags."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(
            path=None, repository="test", index_prefix="redb",
            s3_mode=True, analyzed=True, decompile=True
        )
        assert ingestor.analyzed is True
        assert ingestor.decompile is True


# ============================================================================
# Re-export Test for fetch_analyzed_samples
# ============================================================================

class TestAnalyzedReExport:
    """Tests that fetch_analyzed_samples is re-exported from redb.ingestor."""

    def test_importable_from_ingestor(self):
        """fetch_analyzed_samples can be imported from redb.ingestor."""
        from redb.ingestor import fetch_analyzed_samples
        assert callable(fetch_analyzed_samples)


# ============================================================================
# Decompile-Path Per-File Dedup (is_in_code_db) Tests
# ============================================================================

class TestDecompilePerFileDedupBypass:
    """Tests for the per-file is_in_code_db check in the decompile path.

    Verifies that --force correctly bypasses the is_in_code_db check
    for both normal binaries and FAT Mach-O slices.

    Note: DecompileBinja requires blake3/binaryninja which aren't available
    in the test environment. We use two strategies:
    - For "no force" tests: is_in_code_db returns True so the code never
      reaches the DecompileBinja import.
    - For "force" tests: we mock the import mechanism to prevent the real
      import from happening.
    """

    @patch('redb.workers.is_in_code_db')
    @patch('redb.workers.is_in_db')
    @patch('redb.workers._is_packed', return_value=None)
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_decompile_no_force_calls_is_in_code_db(
        self, mock_settings, mock_magika, mock_hashlib,
        mock_packed, mock_is_in_db, mock_is_in_code_db
    ):
        """In decompile mode without --force, is_in_code_db is called and can skip files."""
        from redb.ingestor import process_binary_file, ImportResult

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        # is_in_code_db returns True => file should be skipped
        mock_is_in_code_db.return_value = True

        logger = Mock()

        mock_mod, _ = _mock_decompiler_module()
        with patch.dict(sys.modules, {"redb.extractors.decompiler": mock_mod}):
            with patch('builtins.open', mock_open_binary(b"test data")):
                result, filetype = process_binary_file(
                    "/tmp/test.bin", True, "redb", logger,
                    selected_modules="all", force=False
                )

        mock_is_in_code_db.assert_called_once_with("a" * 64, filetype="pebin")
        assert result == ImportResult.SKIPPED

    @patch('redb.workers.is_in_code_db')
    @patch('redb.workers.is_in_db')
    @patch('redb.workers._is_packed', return_value=None)
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_decompile_no_force_not_in_code_db_proceeds(
        self, mock_settings, mock_magika, mock_hashlib,
        mock_packed, mock_is_in_db, mock_is_in_code_db
    ):
        """In decompile mode without --force, files NOT in code_db proceed to decompile."""
        from redb.ingestor import process_binary_file

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        # is_in_code_db returns False => should proceed to decompile
        mock_is_in_code_db.return_value = False

        logger = Mock()

        mock_mod, mock_decompiler = _mock_decompiler_module()
        with patch.dict(sys.modules, {"redb.extractors.decompiler": mock_mod}):
            with patch('builtins.open', mock_open_binary(b"test data")):
                try:
                    process_binary_file(
                        "/tmp/test.bin", True, "redb", logger,
                        selected_modules="all", force=False
                    )
                except Exception:
                    pass

        mock_is_in_code_db.assert_called_once_with("a" * 64, filetype="pebin")

    @patch('redb.workers.is_in_code_db')
    @patch('redb.workers.is_in_db')
    @patch('redb.workers._is_packed', return_value=None)
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_decompile_force_skips_is_in_code_db(
        self, mock_settings, mock_magika, mock_hashlib,
        mock_packed, mock_is_in_db, mock_is_in_code_db
    ):
        """In decompile mode with --force, is_in_code_db is never called."""
        from redb.ingestor import process_binary_file

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        logger = Mock()

        mock_mod, _ = _mock_decompiler_module()
        with patch.dict(sys.modules, {"redb.extractors.decompiler": mock_mod}):
            with patch('builtins.open', mock_open_binary(b"test data")):
                try:
                    process_binary_file(
                        "/tmp/test.bin", True, "redb", logger,
                        selected_modules="all", force=True
                    )
                except Exception:
                    pass

        # The key assertion: is_in_code_db was never called
        mock_is_in_code_db.assert_not_called()

    @patch('redb.workers.is_in_code_db')
    @patch('redb.workers.is_in_db')
    @patch('redb.workers._is_packed', return_value=None)
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_decompile_packed_skips_is_in_code_db(
        self, mock_settings, mock_magika, mock_hashlib,
        mock_packed, mock_is_in_db, mock_is_in_code_db
    ):
        """Packed files skip decompile entirely, so is_in_code_db is never called."""
        from redb.ingestor import process_binary_file

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        # _is_packed returns a truthy value (packer name)
        mock_packed.return_value = "UPX"

        logger = Mock()

        mock_mod, _ = _mock_decompiler_module()
        with patch.dict(sys.modules, {"redb.extractors.decompiler": mock_mod}):
            with patch('builtins.open', mock_open_binary(b"test data")):
                try:
                    process_binary_file(
                        "/tmp/test.bin", True, "redb", logger,
                        selected_modules="all", force=False
                    )
                except Exception:
                    pass

        mock_is_in_code_db.assert_not_called()

    @patch('redb.workers.is_in_code_db')
    @patch('redb.workers.is_in_db')
    @patch('redb.workers.hashlib')
    @patch('redb.workers.Magika')
    @patch('redb.workers.settings')
    def test_non_decompile_mode_never_calls_is_in_code_db(
        self, mock_settings, mock_magika, mock_hashlib,
        mock_is_in_db, mock_is_in_code_db
    ):
        """In feature-extraction mode (no decompile), is_in_code_db is never called."""
        from redb.ingestor import process_binary_file, ImportResult

        mock_magika_instance = Mock()
        mock_magika_instance.identify_bytes.return_value.output.label = "pebin"
        mock_magika.return_value = mock_magika_instance
        mock_hashlib.sha256.return_value.hexdigest.return_value = "a" * 64
        mock_settings.CLICKHOUSE_HOST = "localhost"
        mock_settings.get_clickhouse_client.return_value = Mock()

        # is_in_db returns True => file skipped by feature-extraction dedup
        mock_is_in_db.return_value = True

        logger = Mock()

        with patch('builtins.open', mock_open_binary(b"test data")):
            result, filetype = process_binary_file(
                "/tmp/test.bin", False, "redb", logger,
                selected_modules="all", force=False
            )

        assert result == ImportResult.SKIPPED
        mock_is_in_db.assert_called_once()
        mock_is_in_code_db.assert_not_called()


# ============================================================================
# --rerun CLI Argument Tests
# ============================================================================

class TestRerunCliArgument:
    """Tests for --rerun CLI argument parsing and validation."""

    def _parse_args(self, args_list):
        """Helper to parse CLI args."""
        import argparse

        parser = argparse.ArgumentParser()
        input_group = parser.add_mutually_exclusive_group(required=True)
        input_group.add_argument("--path")
        input_group.add_argument("--s3", action="store_true")
        input_group.add_argument("--s3-solo", metavar="S3_KEY")
        input_group.add_argument("--nomad-job", action="store_true")
        input_group.add_argument("--date", metavar="YYYY-MM-DD")
        input_group.add_argument("--range", nargs=2)
        input_group.add_argument("--analyzed", action="store_true")

        parser.add_argument("--repo")
        parser.add_argument("--index_prefix", default="redb")
        parser.add_argument("-d", "--decompile", action="store_true")
        parser.add_argument("-m", "--modules", default="all")
        parser.add_argument("--force", action="store_true")
        parser.add_argument("--rerun", action="store_true")
        parser.add_argument("--dry-run", action="store_true")

        return parser.parse_args(args_list)

    def test_rerun_flag_default_false(self):
        """--rerun is False by default."""
        args = self._parse_args(["--analyzed", "-d"])
        assert args.rerun is False

    def test_rerun_flag_explicit(self):
        """--rerun is True when specified."""
        args = self._parse_args(["--analyzed", "-d", "--rerun"])
        assert args.rerun is True

    def test_rerun_with_analyzed_and_decompile(self):
        """--rerun can be combined with --analyzed and --decompile."""
        args = self._parse_args(["--analyzed", "-d", "--rerun"])
        assert args.rerun is True
        assert args.analyzed is True
        assert args.decompile is True


# ============================================================================
# --rerun fetch_analyzed_samples Tests
# ============================================================================

class TestRerunFetchAnalyzedSamples:
    """Tests for --rerun behavior in fetch_analyzed_samples."""

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_rerun_queries_disassembled_table(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode queries disassembled_functions_references with JOIN to basic_properties."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, rerun=True)

        # Should only make one query (the JOIN query)
        assert mock_client.query.call_count == 1
        query = mock_client.query.call_args_list[0][0][0]
        assert "code_binja_disassembled_functions_references" in query
        assert "redb_basic_properties" in query
        assert "INNER JOIN" in query
        assert len(result) == 1

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_rerun_returns_only_disassembled_samples(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode returns only samples present in disassembled table."""
        from redb.queries import fetch_analyzed_samples

        hash1 = "hash1" + "0" * 59
        mock_client = Mock()
        mock_client.query.return_value.result_rows = [(hash1, "pebin")]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, rerun=True)

        assert len(result) == 1
        assert result[0]['sha256'] == hash1
        assert result[0]['s3_bucket'] == "test-bucket"

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_rerun_with_magika_filter(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode applies magika filter in the JOIN query."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "elf"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, rerun=True, magika_filter="elf")

        query = mock_client.query.call_args_list[0][0][0]
        assert "bp.filetype_magika" in query
        params = mock_client.query.call_args_list[0][1].get("parameters", {})
        assert params.get("magika") == "elf"
        assert len(result) == 1

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket", "CLICKHOUSE_CODE_PREFIX": "custom_code"})
    def test_rerun_uses_code_prefix_env(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode uses CLICKHOUSE_CODE_PREFIX env var."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        fetch_analyzed_samples("redb", decompile=True, rerun=True)

        query = mock_client.query.call_args_list[0][0][0]
        assert "custom_code_disassembled_functions_references" in query

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_rerun_empty_disassembled_table(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode returns empty list when no disassembled samples exist."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, rerun=True)

        assert result == []

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    def test_rerun_missing_s3_bucket(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode returns empty list when S3_BUCKET is not set."""
        from redb.queries import fetch_analyzed_samples

        with patch.dict(os.environ, {}, clear=True):
            result = fetch_analyzed_samples("redb", decompile=True, rerun=True)

        assert result == []

    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_rerun_handles_bytes_sha256(self, mock_settings, mock_formats, mock_gen_key):
        """Rerun mode handles bytes-encoded SHA256 from ClickHouse."""
        from redb.queries import fetch_analyzed_samples

        hash_str = "abc123" + "0" * 58
        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            (hash_str.encode('utf-8'), "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=True, rerun=True)

        assert len(result) == 1
        assert result[0]['sha256'] == hash_str


# ============================================================================
# --rerun Ingestor Parameter Tests
# ============================================================================

class TestIngestorRerunParam:
    """Tests for rerun parameter in Ingestor class."""

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_default_rerun_false(self, mock_open, mock_datetime, mock_mp):
        """Ingestor defaults rerun to False."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(path="/tmp/test", repository="test", index_prefix="redb")
        assert ingestor.rerun is False

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_rerun_true(self, mock_open, mock_datetime, mock_mp):
        """Ingestor stores rerun=True when passed."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(
            path=None, repository="test", index_prefix="redb",
            s3_mode=True, analyzed=True, decompile=True, rerun=True
        )
        assert ingestor.rerun is True

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def test_ingestor_rerun_with_decompile_modules(self, mock_open, mock_datetime, mock_mp):
        """Ingestor stores rerun alongside decompile_modules."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        ingestor = Ingestor(
            path=None, repository="test", index_prefix="redb",
            s3_mode=True, analyzed=True, decompile=True, rerun=True,
            decompile_modules={"cfg"}
        )
        assert ingestor.rerun is True
        assert ingestor.decompile_modules == {"cfg"}


# ============================================================================
# YARA --force Dedup Tests
# ============================================================================

class TestYaraForceDedup:
    """Tests for --force behavior with YARA scanning.

    Default (no --force): skip samples already in yara_matches.
    With --force: scan everything (rescan all samples).
    """

    # --- is_in_db_bulk YARA path ---

    @patch('redb.queries.settings')
    def test_is_in_db_bulk_yara_returns_already_scanned(self, mock_settings):
        """In YARA mode, is_in_db_bulk queries yara_matches and returns already-scanned hashes."""
        from redb.queries import is_in_db_bulk

        mock_client = Mock()
        # hex(sha256) returns uppercase by convention from ClickHouse
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58,),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = is_in_db_bulk(
            ["abc123" + "0" * 58, "def456" + "0" * 58],
            "redb", decompile=False, yara_scan=True
        )

        # Should have queried yara_matches
        query = mock_client.query.call_args[0][0]
        assert "yara_matches" in query
        assert "unhex(" in query
        # Should return the already-scanned hash
        assert "abc123" + "0" * 58 in result
        assert "def456" + "0" * 58 not in result

    @patch('redb.queries.settings')
    def test_is_in_db_bulk_yara_empty_table(self, mock_settings):
        """In YARA mode with empty yara_matches, no samples are skipped."""
        from redb.queries import is_in_db_bulk

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = is_in_db_bulk(
            ["abc123" + "0" * 58],
            "redb", decompile=False, yara_scan=True
        )

        assert len(result) == 0

    @patch('redb.queries.settings')
    def test_is_in_db_bulk_yara_handles_bytes_response(self, mock_settings):
        """YARA dedup correctly handles bytes-encoded hex response from ClickHouse."""
        from redb.queries import is_in_db_bulk

        hash_hex = "abc123" + "0" * 58
        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            (hash_hex.upper().encode('utf-8'),),  # bytes, uppercase
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = is_in_db_bulk(
            [hash_hex],
            "redb", decompile=False, yara_scan=True
        )

        # Should normalize to lowercase
        assert hash_hex in result

    @patch('redb.queries.settings')
    def test_is_in_db_bulk_yara_db_error_raises(self, mock_settings):
        """On DB error, YARA dedup raises RuntimeError to prevent reprocessing all samples."""
        from redb.queries import is_in_db_bulk
        import pytest

        mock_settings.create_clickhouse_client.side_effect = Exception("Connection failed")

        with pytest.raises(RuntimeError, match="YARA dedup query failed"):
            is_in_db_bulk(
                ["abc123" + "0" * 58],
                "redb", decompile=False, yara_scan=True
            )

    def test_is_in_db_bulk_yara_empty_list(self):
        """is_in_db_bulk with empty list returns empty set."""
        from redb.queries import is_in_db_bulk

        result = is_in_db_bulk([], "redb", decompile=False, yara_scan=True)
        assert result == set()

    # --- fetch_s3_objects_by_repository YARA+force ---

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_repo_yara_default_calls_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """Without --force, YARA fetch calls is_in_db_bulk to skip already-scanned."""
        from redb.ingestor import fetch_s3_objects_by_repository

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None),
            ("def456" + "0" * 58, "pebin", None),
        ]
        mock_db_conn.return_value = mock_client
        mock_bulk.return_value = {"abc123" + "0" * 58}  # One already scanned

        result = fetch_s3_objects_by_repository(
            "test_repo", "redb", False, yara_scan=True, force=False
        )

        mock_bulk.assert_called_once()
        # The yara_scan=True should be passed through
        assert mock_bulk.call_args[0][3] is True  # yara_scan arg
        # One sample should be filtered out
        assert len(result) == 1

    @patch('redb.queries.is_in_db_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.get_db_catalog_connection')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_by_repo_yara_force_skips_bulk_check(
        self, mock_db_conn, mock_formats, mock_gen_key, mock_bulk
    ):
        """With --force, YARA fetch skips is_in_db_bulk entirely."""
        from redb.ingestor import fetch_s3_objects_by_repository

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("abc123" + "0" * 58, "pebin", None),
            ("def456" + "0" * 58, "pebin", None),
        ]
        mock_db_conn.return_value = mock_client

        result = fetch_s3_objects_by_repository(
            "test_repo", "redb", False, yara_scan=True, force=True
        )

        mock_bulk.assert_not_called()
        assert len(result) == 2

    # --- fetch_analyzed_samples YARA+force ---

    @patch('redb.queries._check_yara_matches_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_analyzed_yara_default_filters_already_scanned(
        self, mock_settings, mock_formats, mock_gen_key, mock_yara_check
    ):
        """Without --force, fetch_analyzed_samples filters already YARA-scanned samples."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        hash1 = "hash1" + "0" * 59
        hash2 = "hash2" + "0" * 59
        hash3 = "hash3" + "0" * 59
        mock_client.query.return_value.result_rows = [
            (hash1, "pebin"),
            (hash2, "pebin"),
            (hash3, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client
        mock_yara_check.return_value = {hash2}  # hash2 already scanned

        result = fetch_analyzed_samples("redb", decompile=False, yara_scan=True, force=False)

        mock_yara_check.assert_called_once()
        assert len(result) == 2
        result_hashes = {r['sha256'] for r in result}
        assert hash2 not in result_hashes

    @patch('redb.queries._check_yara_matches_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_analyzed_yara_force_skips_yara_check(
        self, mock_settings, mock_formats, mock_gen_key, mock_yara_check
    ):
        """With --force, fetch_analyzed_samples skips YARA dedup check."""
        from redb.queries import fetch_analyzed_samples

        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            ("hash1" + "0" * 59, "pebin"),
            ("hash2" + "0" * 59, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client

        result = fetch_analyzed_samples("redb", decompile=False, yara_scan=True, force=True)

        mock_yara_check.assert_not_called()
        assert len(result) == 2

    @patch('redb.queries._check_yara_matches_bulk')
    @patch('redb.queries.generate_s3_key_from_hash', return_value="ab/cd/abcd.zip")
    @patch('redb.queries.get_supported_formats', return_value=["pebin"])
    @patch('redb.queries.settings')
    @patch.dict(os.environ, {"S3_BUCKET": "test-bucket"})
    def test_fetch_analyzed_yara_all_already_scanned(
        self, mock_settings, mock_formats, mock_gen_key, mock_yara_check
    ):
        """When all samples already scanned, returns empty list."""
        from redb.queries import fetch_analyzed_samples

        hash1 = "hash1" + "0" * 59
        mock_client = Mock()
        mock_client.query.return_value.result_rows = [
            (hash1, "pebin"),
        ]
        mock_settings.create_clickhouse_client.return_value = mock_client
        mock_yara_check.return_value = {hash1}

        result = fetch_analyzed_samples("redb", decompile=False, yara_scan=True, force=False)

        assert len(result) == 0

    # --- _check_yara_matches_bulk direct tests ---

    @patch('redb.queries.settings')
    def test_check_yara_matches_bulk_builds_unhex_query(self, mock_settings):
        """_check_yara_matches_bulk uses unhex() for binary sha256 comparison."""
        from redb.queries import _check_yara_matches_bulk

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        hash1 = "aabb" + "0" * 60
        _check_yara_matches_bulk([hash1])

        query = mock_client.query.call_args[0][0]
        assert f"unhex('{hash1}')" in query
        assert "hex(sha256)" in query
        assert "yara_matches" in query

    @patch('redb.queries.settings')
    def test_check_yara_matches_bulk_batches_large_lists(self, mock_settings):
        """_check_yara_matches_bulk batches queries for large hash lists."""
        from redb.queries import _check_yara_matches_bulk

        mock_client = Mock()
        mock_client.query.return_value.result_rows = []
        mock_settings.create_clickhouse_client.return_value = mock_client

        # Create list larger than batch_size (3900)
        hashes = [f"{i:064x}" for i in range(4000)]
        _check_yara_matches_bulk(hashes)

        # Should have made 2 queries (3900 + 100)
        assert mock_client.query.call_count == 2


# ============================================================================
# Query Dispatch Tests (--analyzed + --date/--range flag combinations)
# ============================================================================

class TestIngestorQueryDispatch:
    """Tests that the ingestor dispatches to the correct query function
    based on flag combinations.

    Regression tests for a bug where `if self.analyzed` was checked before
    `if self.start_date and self.end_date`, causing --analyzed --date/--range
    to ignore the date filter and query all 4M+ samples from basic_properties.
    """

    @patch('redb.ingestor.multiprocessing')
    @patch('redb.ingestor.datetime')
    @patch('builtins.open', create=True)
    def _make_ingestor(self, mock_open, mock_datetime, mock_mp, **kwargs):
        """Helper to create an Ingestor with mocked multiprocessing/datetime."""
        mock_open.return_value.__enter__ = Mock()
        mock_open.return_value.__exit__ = Mock(return_value=False)
        mock_datetime.today.return_value.strftime.return_value = "20260222"

        from redb.ingestor import Ingestor
        defaults = dict(
            path=None, repository="analyzed", index_prefix="redb",
            s3_mode=True, force=True,
        )
        defaults.update(kwargs)
        return Ingestor(**defaults)

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_analyzed_with_date_uses_date_range_query(
        self, mock_date_range, mock_analyzed
    ):
        """--analyzed --date should use fetch_s3_objects_by_date_range, not fetch_analyzed_samples."""
        ingestor = self._make_ingestor(
            analyzed=True, start_date="2025-01-15", end_date="2025-01-16"
        )
        ingestor.ingest()

        mock_date_range.assert_called_once()
        mock_analyzed.assert_not_called()
        # Verify analyzed=True was passed through to the date range query
        call_kwargs = mock_date_range.call_args
        assert call_kwargs[1].get('analyzed') or call_kwargs[0][-1] if call_kwargs[0] else call_kwargs[1].get('analyzed') is True

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_analyzed_with_range_uses_date_range_query(
        self, mock_date_range, mock_analyzed
    ):
        """--analyzed --range should use fetch_s3_objects_by_date_range, not fetch_analyzed_samples."""
        ingestor = self._make_ingestor(
            analyzed=True, start_date="2025-01-01", end_date="2025-01-31"
        )
        ingestor.ingest()

        mock_date_range.assert_called_once()
        mock_analyzed.assert_not_called()

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_analyzed_standalone_uses_fetch_analyzed(
        self, mock_date_range, mock_analyzed
    ):
        """--analyzed without --date/--range should use fetch_analyzed_samples."""
        ingestor = self._make_ingestor(
            analyzed=True, start_date=None, end_date=None
        )
        ingestor.ingest()

        mock_analyzed.assert_called_once()
        mock_date_range.assert_not_called()

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_date_without_analyzed_uses_date_range_query(
        self, mock_date_range, mock_analyzed
    ):
        """--date without --analyzed should use fetch_s3_objects_by_date_range with analyzed=False."""
        ingestor = self._make_ingestor(
            analyzed=False, start_date="2025-01-15", end_date="2025-01-16"
        )
        ingestor.ingest()

        mock_date_range.assert_called_once()
        mock_analyzed.assert_not_called()

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_analyzed_with_date_passes_analyzed_flag(
        self, mock_date_range, mock_analyzed
    ):
        """--analyzed --date should pass analyzed=True to fetch_s3_objects_by_date_range."""
        ingestor = self._make_ingestor(
            analyzed=True, start_date="2025-01-15", end_date="2025-01-16"
        )
        ingestor.ingest()

        mock_date_range.assert_called_once()
        _, kwargs = mock_date_range.call_args
        assert kwargs.get('analyzed') is True

    @patch('redb.ingestor.fetch_analyzed_samples', return_value=[])
    @patch('redb.ingestor.fetch_s3_objects_by_date_range', return_value=[])
    def test_analyzed_with_date_passes_date_params(
        self, mock_date_range, mock_analyzed
    ):
        """--analyzed --date should pass correct start_date and end_date."""
        ingestor = self._make_ingestor(
            analyzed=True, start_date="2025-03-01", end_date="2025-03-15"
        )
        ingestor.ingest()

        mock_date_range.assert_called_once()
        _, kwargs = mock_date_range.call_args
        assert kwargs.get('start_date') == "2025-03-01"
        assert kwargs.get('end_date') == "2025-03-15"


# ============================================================================
# Helpers
# ============================================================================

def mock_open_binary(content):
    """Create a mock for open() that returns binary content."""
    mock = MagicMock()
    mock.return_value.__enter__ = Mock(return_value=Mock(read=Mock(return_value=content)))
    mock.return_value.__exit__ = Mock(return_value=False)
    return mock