config loader

Author: Behnam Ousat

.pyrit_conf_example

@@ -125,6 +125,18 @@ operation: op_trash_panda
 # env_akv_ref:
 #   - https://my-vault.vault.azure.net/secrets/my-pyrit-env
 
+# Datasets
+# --------
+# List of seed dataset names to load into memory after initialization completes.
+# Each name must match a registered dataset (run `pyrit_scan --list-datasets` to
+# see available datasets). Loaded datasets are added to CentralMemory and are
+# available to scenarios and attacks.
+#
+# Example:
+# datasets:
+#   - airt_illegal
+#   - airt_malware
+
 # Max Concurrent Scenario Runs
 # ----------------------------
 # Maximum number of scenario runs that can execute concurrently in the backend.

pyrit/setup/configuration_loader.py

@@ -103,6 +103,7 @@ class ConfigurationLoader(YamlLoadable):
             None means "use defaults", [] means "load nothing".
         env_files: List of environment file paths to load.
             None means "use defaults (.env, .env.local)", [] means "load nothing".
+        datasets: List of seed dataset names to load into memory after initialization.
         silent: Whether to suppress initialization messages.
         operator: Name for the current operator, e.g. a team or username.
         operation: Name for the current operation.
@@ -123,6 +124,10 @@ class ConfigurationLoader(YamlLoadable):
           - /path/to/.env
           - /path/to/.env.local
 
+        datasets:
+          - airt_illegal
+          - airt_malware
+
         silent: false
 
         operator: my_team
@@ -141,6 +146,7 @@ class ConfigurationLoader(YamlLoadable):
     initialization_scripts: list[str] | None = None
     env_files: list[str] | None = None
     env_akv_ref: list[str] | None = None
+    datasets: list[str] = field(default_factory=list)
     silent: bool = False
     operator: str | None = None
     operation: str | None = None
@@ -379,6 +385,7 @@ def load_with_overrides(
                 config_data["initialization_scripts"] = default_config.initialization_scripts
                 config_data["env_files"] = default_config.env_files
                 config_data["env_akv_ref"] = default_config.env_akv_ref
+                config_data["datasets"] = default_config.datasets
                 config_data["silent"] = default_config.silent
                 if default_config.operator:
                     config_data["operator"] = default_config.operator
@@ -405,6 +412,7 @@ def load_with_overrides(
             config_data["initialization_scripts"] = explicit_config.initialization_scripts
             config_data["env_files"] = explicit_config.env_files
             config_data["env_akv_ref"] = explicit_config.env_akv_ref
+            config_data["datasets"] = explicit_config.datasets
             config_data["silent"] = explicit_config.silent
             if explicit_config.operator:
                 config_data["operator"] = explicit_config.operator
@@ -560,7 +568,8 @@ async def initialize_pyrit_async(self) -> None:
         Initialize PyRIT with the loaded configuration.
 
         This method resolves all initializer names to instances and calls
-        the core initialize_pyrit_async function.
+        the core initialize_pyrit_async function. Any datasets configured via
+        the ``datasets`` block are loaded into memory after initialization.
 
         Raises:
             ValueError: If configuration is invalid or initializers cannot be resolved.
@@ -581,6 +590,37 @@ async def initialize_pyrit_async(self) -> None:
             silent=self.silent,
         )
 
+        await self._load_datasets_async()
+
+    async def _load_datasets_async(self) -> None:
+        """
+        Load the configured seed datasets into memory.
+
+        Fetches each dataset named in the ``datasets`` block and adds its seeds
+        to ``CentralMemory``. This runs after PyRIT initialization so that memory
+        is available. No-op when no datasets are configured.
+
+        Raises:
+            ValueError: If any configured dataset name does not exist.
+        """
+        if not self.datasets:
+            return
+
+        import logging
+
+        from pyrit.datasets import SeedDatasetProvider
+        from pyrit.memory import CentralMemory
+
+        logger = logging.getLogger(__name__)
+        logger.info("Loading %d dataset(s) from configuration...", len(self.datasets))
+
+        datasets = await SeedDatasetProvider.fetch_datasets_async(dataset_names=self.datasets)
+
+        memory = CentralMemory.get_memory_instance()
+        await memory.add_seed_datasets_to_memory_async(datasets=datasets, added_by="ConfigurationLoader")
+
+        logger.info("Successfully loaded %d dataset(s) into memory", len(datasets))
+
 
 async def initialize_from_config_async(
     config_path: str | pathlib.Path | None = None,

tests/unit/setup/test_configuration_loader.py

@@ -397,6 +397,43 @@ async def test_initialize_pyrit_async_unknown_initializer_raises_error(self, moc
         with pytest.raises(ValueError, match="not found in registry"):
             await config.initialize_pyrit_async()
 
+    @mock.patch("pyrit.memory.CentralMemory")
+    @mock.patch("pyrit.datasets.SeedDatasetProvider")
+    @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async")
+    async def test_initialize_pyrit_async_loads_datasets(self, mock_init, mock_provider, mock_memory_cls):
+        """Test that configured datasets are fetched and added to memory after init."""
+        fetched = [mock.MagicMock()]
+        mock_provider.fetch_datasets_async = mock.AsyncMock(return_value=fetched)
+        mock_memory = mock.MagicMock()
+        mock_memory.add_seed_datasets_to_memory_async = mock.AsyncMock()
+        mock_memory_cls.get_memory_instance.return_value = mock_memory
+
+        config = ConfigurationLoader(
+            memory_db_type="in_memory",
+            datasets=["airt_illegal", "airt_malware"],
+        )
+        await config.initialize_pyrit_async()
+
+        mock_init.assert_called_once()
+        mock_provider.fetch_datasets_async.assert_awaited_once_with(
+            dataset_names=["airt_illegal", "airt_malware"]
+        )
+        mock_memory.add_seed_datasets_to_memory_async.assert_awaited_once_with(
+            datasets=fetched, added_by="ConfigurationLoader"
+        )
+
+    @mock.patch("pyrit.datasets.SeedDatasetProvider")
+    @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async")
+    async def test_initialize_pyrit_async_no_datasets_skips_loading(self, mock_init, mock_provider):
+        """Test that dataset loading is skipped when no datasets are configured."""
+        mock_provider.fetch_datasets_async = mock.AsyncMock()
+
+        config = ConfigurationLoader(memory_db_type="in_memory")
+        await config.initialize_pyrit_async()
+
+        mock_init.assert_called_once()
+        mock_provider.fetch_datasets_async.assert_not_called()
+
 
 @pytest.mark.usefixtures("patch_central_database")
 class TestInitializeFromConfigAsync: