scripts/run_all_adapters.py

#!/usr/bin/env python3
"""
ART-SafeBench v2.0.0 - Benchmark Augmentation Generator

This script orchestrates all external dataset adapters to augment the benchmark
with data from verified external sources.

Usage:
    python run_all_adapters.py --core-only         # Only Core (redistributable) datasets
    python run_all_adapters.py --extended-only     # Only Extended (research) datasets
    python run_all_adapters.py --all               # All datasets
    python run_all_adapters.py --list              # List available adapters

IMPORTANT LICENSING NOTE:
- Core datasets (MIT/BSD/Apache) can be redistributed with the benchmark
- Extended datasets (CC-BY-NC, MIT+Clause) are research-only and must be
  generated at runtime by the end user. Do NOT commit Extended outputs to git.

Copyright (c) 2025 Fujitsu Research of Europe Limited.
Licensed under CC-BY-4.0.
"""

import argparse
import json
import logging
import sys
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any, Optional

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


@dataclass
class AdapterInfo:
    """Information about an available adapter."""
    name: str
    adapter_class: str
    target_surface: str
    license: str
    can_redistribute: bool
    requires_download: bool
    source_url: str
    description: str


# Registry of all available adapters
ADAPTER_REGISTRY: Dict[str, AdapterInfo] = {
    "figstep": AdapterInfo(
        name="FigStep",
        adapter_class="b2_figstep_adapter.FigStepAdapter",
        target_surface="B2",
        license="MIT",
        can_redistribute=True,
        requires_download=True,
        source_url="https://github.com/ThuCCSLab/FigStep",
        description="SafeBench (500) + typographic image prompts for multimodal jailbreaking",
    ),
    "harmbench": AdapterInfo(
        name="HarmBench",
        adapter_class="b3_harmbench_adapter.HarmBenchAdapter",
        target_surface="B3",
        license="MIT",
        can_redistribute=True,
        requires_download=True,
        source_url="https://github.com/centerforaisafety/HarmBench",
        description="400 curated harmful behaviors for red teaming"
    ),
    "jailbreakbench": AdapterInfo(
        name="JailbreakBench",
        adapter_class="b3_jailbreakbench_adapter.JailbreakBenchAdapter",
        target_surface="B3",
        license="MIT",
        can_redistribute=True,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/JailbreakBench/JBB-Behaviors",
        description="100 harmful behaviors for jailbreak evaluation"
    ),
    "gandalf": AdapterInfo(
        name="Gandalf",
        adapter_class="b3_gandalf_adapter.GandalfAdapter",
        target_surface="B3",
        license="MIT",
        can_redistribute=True,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/Lakera/gandalf_ignore_instructions",
        description="Human-generated prompt injections from Gandalf game"
    ),
    "tensortrust": AdapterInfo(
        name="TensorTrust",
        adapter_class="b3_tensortrust_adapter.TensorTrustAdapter",
        target_surface="B3",
        license="BSD-2-Clause",
        can_redistribute=True,
        requires_download=True,
        source_url="https://github.com/HumanCompatibleAI/tensor-trust",
        description="PromptGame-style prompt extraction/hijacking attacks (requires exported JSONL)",
    ),
    "injecagent": AdapterInfo(
        name="InjecAgent",
        adapter_class="b4_injecagent_adapter.InjecAgentAdapter",
        target_surface="B4",
        license="MIT",
        can_redistribute=True,
        requires_download=True,
        source_url="https://github.com/uiuc-kang-lab/InjecAgent",
        description="1,054 indirect prompt injection test cases"
    ),
    "raguard": AdapterInfo(
        name="RAGuard",
        adapter_class="b1_raguard_adapter.RAGuardAdapter",
        target_surface="B1",
        license="MIT",
        can_redistribute=True,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/UCSC-IRKM/RAGuard",
        description="RAG robustness against misleading retrievals"
    ),
    "toolemu": AdapterInfo(
        name="ToolEmu",
        adapter_class="b4_toolemu_adapter.ToolEmuAdapter",
        target_surface="B4",
        license="Apache-2.0",
        can_redistribute=True,
        requires_download=True,
        source_url="https://github.com/ryoungj/ToolEmu",
        description="Tool-use risk benchmark cases for LM agents (144 cases)",
    ),
    "agentharm": AdapterInfo(
        name="AgentHarm",
        adapter_class="b4_agentharm_adapter.AgentHarmAdapter",
        target_surface="B4",
        license="MIT+Research-Clause",
        can_redistribute=False,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/ai-safety-institute/AgentHarm",
        description="110 malicious agent behaviors (RESEARCH ONLY)"
    ),
    "mindgard": AdapterInfo(
        name="Mindgard",
        adapter_class="b3_mindgard_adapter.MindgardAdapter",
        target_surface="B3",
        license="CC-BY-NC-4.0",
        can_redistribute=False,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/Mindgard/evaded-prompt-injection-and-jailbreak-samples",
        description="Evaded prompt injections (RESEARCH ONLY)"
    ),
    "mmsafetybench": AdapterInfo(
        name="MM-SafetyBench",
        adapter_class="b2_mmsafetybench_adapter.MMSafetyBenchAdapter",
        target_surface="B2",
        license="CC-BY-NC-4.0",
        can_redistribute=False,
        requires_download=False,  # Auto-download from HuggingFace
        source_url="https://huggingface.co/datasets/PKU-Alignment/MM-SafetyBench",
        description="Multimodal safety prompts with images (RESEARCH ONLY)",
    ),
}


def list_adapters(filter_type: str = "all") -> None:
    """List all available adapters."""
    print("\n" + "="*80)
    print("ART-SAFEBENCH v2.0.0 - Available Adapters")
    print("="*80)

    core_adapters = {k: v for k, v in ADAPTER_REGISTRY.items() if v.can_redistribute}
    extended_adapters = {k: v for k, v in ADAPTER_REGISTRY.items() if not v.can_redistribute}

    if filter_type in ["all", "core"]:
        print("\n[CORE DATASETS - Redistributable (MIT/BSD/Apache)]")
        print("-"*60)
        for key, info in core_adapters.items():
            dl_status = "Manual download" if info.requires_download else "Auto (HuggingFace)"
            print(f"  {key:15} | {info.target_surface} | {info.license:8} | {dl_status}")
            print(f"                  | {info.description}")
            print(f"                  | {info.source_url}")
            print()

    if filter_type in ["all", "extended"]:
        print("\n[EXTENDED DATASETS - Research Only (DO NOT REDISTRIBUTE)]")
        print("-"*60)
        for key, info in extended_adapters.items():
            dl_status = "Manual download" if info.requires_download else "Auto (HuggingFace)"
            print(f"  {key:15} | {info.target_surface} | {info.license:20} | {dl_status}")
            print(f"                  | {info.description}")
            print(f"                  | {info.source_url}")
            print()

    print("="*80)


def run_adapter(
    adapter_key: str,
    output_base: Path,
    max_records: Optional[int] = None,
    verbose: bool = False,
    **kwargs
) -> Dict[str, Any]:
    """
    Run a single adapter.

    Args:
        adapter_key: Key from ADAPTER_REGISTRY
        output_base: Base output directory
        max_records: Maximum records to process
        verbose: Enable verbose logging
        **kwargs: Additional adapter-specific arguments

    Returns:
        Result dictionary with success status and statistics
    """
    if adapter_key not in ADAPTER_REGISTRY:
        return {"success": False, "error": f"Unknown adapter: {adapter_key}"}

    info = ADAPTER_REGISTRY[adapter_key]

    # Determine output directory based on license
    if info.can_redistribute:
        output_dir = output_base / "core"
    else:
        output_dir = output_base / "extended"

    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(f"Running adapter: {info.name}")
    logger.info(f"Target surface: {info.target_surface}")
    logger.info(f"Output directory: {output_dir}")

    try:
        # Dynamic import of adapter class
        module_name = info.adapter_class.rsplit('.', 1)[0]
        class_name = info.adapter_class.rsplit('.', 1)[1]

        # Import from the adapters package
        import importlib
        module = importlib.import_module(f"adapters.{module_name}")
        adapter_class = getattr(module, class_name)

        # Instantiate and run adapter
        adapter = adapter_class(output_dir=output_dir, verbose=verbose, **kwargs)
        result = adapter.run(max_records=max_records)

        return {
            "success": result.success,
            "adapter": adapter_key,
            "name": info.name,
            "target_surface": info.target_surface,
            "records_processed": result.records_processed,
            "records_failed": result.records_failed,
            "output_file": result.output_file,
            "errors": result.errors,
            "warnings": result.warnings,
            "license": info.license,
            "can_redistribute": info.can_redistribute
        }

    except ImportError as e:
        return {
            "success": False,
            "adapter": adapter_key,
            "error": f"Import error: {e}. Make sure 'datasets' library is installed."
        }
    except Exception as e:
        logger.exception(f"Adapter {adapter_key} failed")
        return {
            "success": False,
            "adapter": adapter_key,
            "error": str(e)
        }


def run_all_adapters(
    filter_type: str = "all",
    output_base: Path = None,
    max_records: Optional[int] = None,
    verbose: bool = False,
    harmbench_csv_path: Optional[Path] = None,
    injecagent_data_dir: Optional[Path] = None,
    injecagent_attack_type: str = "both",
    figstep_repo_dir: Optional[Path] = None,
    figstep_subset: str = "safebench",
    figstep_copy_images: bool = False,
    toolemu_repo_dir: Optional[Path] = None,
    toolemu_cases_json_path: Optional[Path] = None,
    tensortrust_attacks_jsonl_path: Optional[Path] = None,
    mmsafetybench_configs: Optional[List[str]] = None,
    mmsafetybench_splits: Optional[List[str]] = None,
    mmsafetybench_include_text_only: bool = False,
    mmsafetybench_export_images: bool = False,
) -> Dict[str, Any]:
    """
    Run all adapters matching the filter.

    Args:
        filter_type: "all", "core", or "extended"
        output_base: Base output directory
        max_records: Maximum records per adapter
        verbose: Enable verbose logging

    Returns:
        Summary dictionary with all results
    """
    if output_base is None:
        output_base = Path("data/external_augmented")

    output_base.mkdir(parents=True, exist_ok=True)

    # Filter adapters
    if filter_type == "core":
        adapters_to_run = {k: v for k, v in ADAPTER_REGISTRY.items() if v.can_redistribute}
    elif filter_type == "extended":
        adapters_to_run = {k: v for k, v in ADAPTER_REGISTRY.items() if not v.can_redistribute}
    else:
        adapters_to_run = ADAPTER_REGISTRY

    results = []
    total_records = 0
    successful = 0
    failed = 0

    print("\n" + "="*80)
    print(f"RUNNING {len(adapters_to_run)} ADAPTERS ({filter_type.upper()})")
    print("="*80 + "\n")

    for adapter_key, info in adapters_to_run.items():
        adapter_kwargs = {}
        if adapter_key == "figstep" and figstep_repo_dir:
            adapter_kwargs["repo_dir"] = figstep_repo_dir
            adapter_kwargs["subset"] = figstep_subset
            adapter_kwargs["copy_images"] = figstep_copy_images
        if adapter_key == "harmbench" and harmbench_csv_path:
            adapter_kwargs["csv_path"] = harmbench_csv_path
        if adapter_key == "injecagent" and injecagent_data_dir:
            adapter_kwargs["data_dir"] = injecagent_data_dir
            adapter_kwargs["attack_type"] = injecagent_attack_type
        if adapter_key == "toolemu" and (toolemu_repo_dir or toolemu_cases_json_path):
            adapter_kwargs["repo_dir"] = toolemu_repo_dir
            adapter_kwargs["cases_json_path"] = toolemu_cases_json_path
        if adapter_key == "tensortrust" and tensortrust_attacks_jsonl_path:
            adapter_kwargs["attacks_jsonl_path"] = tensortrust_attacks_jsonl_path
        if adapter_key == "mmsafetybench":
            if mmsafetybench_configs:
                adapter_kwargs["configs"] = mmsafetybench_configs
            if mmsafetybench_splits:
                adapter_kwargs["splits"] = mmsafetybench_splits
            if mmsafetybench_include_text_only:
                adapter_kwargs["include_text_only"] = True
            if mmsafetybench_export_images:
                adapter_kwargs["export_images"] = True

        if info.requires_download and not adapter_kwargs:
            logger.warning(
                f"Skipping {info.name}: Requires manual download from {info.source_url}"
            )
            results.append({
                "adapter": adapter_key,
                "success": False,
                "skipped": True,
                "reason": f"Requires manual download from {info.source_url}"
            })
            continue

        result = run_adapter(
            adapter_key,
            output_base,
            max_records=max_records,
            verbose=verbose,
            **adapter_kwargs,
        )

        results.append(result)

        if result.get("success"):
            successful += 1
            total_records += result.get("records_processed", 0)
        else:
            failed += 1

    # Generate summary
    summary = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "filter_type": filter_type,
        "adapters_attempted": len(adapters_to_run),
        "adapters_successful": successful,
        "adapters_failed": failed,
        "total_records_generated": total_records,
        "output_base": str(output_base),
        "results": results
    }

    # Save summary
    summary_file = output_base / f"adapter_run_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(summary_file, 'w') as f:
        json.dump(summary, f, indent=2)

    # Print summary
    print("\n" + "="*80)
    print("EXECUTION SUMMARY")
    print("="*80)
    print(f"  Adapters attempted: {len(adapters_to_run)}")
    print(f"  Successful: {successful}")
    print(f"  Failed: {failed}")
    print(f"  Total records generated: {total_records}")
    print(f"  Summary saved to: {summary_file}")
    print("="*80 + "\n")

    return summary


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="ART-SafeBench v2.0.0 - Benchmark Augmentation Generator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # List all available adapters
    python run_all_adapters.py --list

    # Run only Core (redistributable) adapters
    python run_all_adapters.py --core-only

    # Run only Extended (research) adapters
    python run_all_adapters.py --extended-only

    # Run all adapters
    python run_all_adapters.py --all

    # Run specific adapter
    python run_all_adapters.py --adapter gandalf

    # Limit records for testing
    python run_all_adapters.py --core-only --max-records 100

IMPORTANT: Extended datasets are for research only. Do not redistribute.
        """
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--list", action="store_true", help="List available adapters")
    group.add_argument("--core-only", action="store_true", help="Run Core adapters only")
    group.add_argument("--extended-only", action="store_true", help="Run Extended adapters only")
    group.add_argument("--all", action="store_true", help="Run all adapters")
    group.add_argument("--adapter", type=str, help="Run a specific adapter by key")

    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented"),
        help="Base output directory"
    )
    parser.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Maximum records per adapter"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose output"
    )

    parser.add_argument(
        "--harmbench-csv-path",
        type=Path,
        default=None,
        help="Path to HarmBench harmbench_behaviors_text_all.csv (enables running the HarmBench adapter)"
    )
    parser.add_argument(
        "--injecagent-data-dir",
        type=Path,
        default=None,
        help="Path to InjecAgent data directory (enables running the InjecAgent adapter)"
    )
    parser.add_argument(
        "--injecagent-attack-type",
        type=str,
        default="both",
        choices=["dh", "ds", "both"],
        help="InjecAgent attack type to process (only used with --injecagent-data-dir)"
    )
    parser.add_argument(
        "--figstep-repo-dir",
        type=Path,
        default=None,
        help="Path to FigStep repo root (enables running the FigStep adapter)",
    )
    parser.add_argument(
        "--figstep-subset",
        type=str,
        default="safebench",
        choices=["safebench", "tiny"],
        help="Which FigStep subset to use (only used with --figstep-repo-dir)",
    )
    parser.add_argument(
        "--figstep-copy-images",
        action="store_true",
        help="Copy FigStep images into the output directory (only used with --figstep-repo-dir)",
    )
    parser.add_argument(
        "--toolemu-repo-dir",
        type=Path,
        default=None,
        help="Path to ToolEmu repo root (enables running the ToolEmu adapter)",
    )
    parser.add_argument(
        "--toolemu-cases-json-path",
        type=Path,
        default=None,
        help="Direct path to ToolEmu all_cases.json (enables running the ToolEmu adapter)",
    )
    parser.add_argument(
        "--tensortrust-attacks-jsonl-path",
        type=Path,
        default=None,
        help="Path to TensorTrust exported attacks JSONL (enables running the TensorTrust adapter)",
    )
    parser.add_argument(
        "--mmsafetybench-configs",
        type=str,
        default=None,
        help="Comma-separated MM-SafetyBench configs to include (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-splits",
        type=str,
        default=None,
        help="Comma-separated MM-SafetyBench splits to include (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-include-text-only",
        action="store_true",
        help="Include MM-SafetyBench Text_only split (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-export-images",
        action="store_true",
        help="Export MM-SafetyBench images into output directory (Extended; do not redistribute)",
    )

    args = parser.parse_args()

    if args.list:
        list_adapters("all")
        return

    if args.adapter:
        if args.adapter not in ADAPTER_REGISTRY:
            print(f"Error: Unknown adapter '{args.adapter}'")
            print("Use --list to see available adapters")
            sys.exit(1)

        adapter_kwargs = {}
        if args.adapter == "harmbench":
            if not args.harmbench_csv_path:
                print("Error: --harmbench-csv-path is required to run the HarmBench adapter")
                sys.exit(1)
            adapter_kwargs["csv_path"] = args.harmbench_csv_path
        if args.adapter == "injecagent":
            if not args.injecagent_data_dir:
                print("Error: --injecagent-data-dir is required to run the InjecAgent adapter")
                sys.exit(1)
            adapter_kwargs["data_dir"] = args.injecagent_data_dir
            adapter_kwargs["attack_type"] = args.injecagent_attack_type
        if args.adapter == "figstep":
            if not args.figstep_repo_dir:
                print("Error: --figstep-repo-dir is required to run the FigStep adapter")
                sys.exit(1)
            adapter_kwargs["repo_dir"] = args.figstep_repo_dir
            adapter_kwargs["subset"] = args.figstep_subset
            adapter_kwargs["copy_images"] = args.figstep_copy_images
        if args.adapter == "toolemu":
            if not (args.toolemu_repo_dir or args.toolemu_cases_json_path):
                print("Error: --toolemu-repo-dir or --toolemu-cases-json-path is required to run the ToolEmu adapter")
                sys.exit(1)
            adapter_kwargs["repo_dir"] = args.toolemu_repo_dir
            adapter_kwargs["cases_json_path"] = args.toolemu_cases_json_path
        if args.adapter == "tensortrust":
            if not args.tensortrust_attacks_jsonl_path:
                print("Error: --tensortrust-attacks-jsonl-path is required to run the TensorTrust adapter")
                sys.exit(1)
            adapter_kwargs["attacks_jsonl_path"] = args.tensortrust_attacks_jsonl_path
        if args.adapter == "mmsafetybench":
            if args.mmsafetybench_configs:
                adapter_kwargs["configs"] = [c.strip() for c in args.mmsafetybench_configs.split(",") if c.strip()]
            if args.mmsafetybench_splits:
                adapter_kwargs["splits"] = [s.strip() for s in args.mmsafetybench_splits.split(",") if s.strip()]
            if args.mmsafetybench_include_text_only:
                adapter_kwargs["include_text_only"] = True
            if args.mmsafetybench_export_images:
                adapter_kwargs["export_images"] = True

        result = run_adapter(
            args.adapter,
            args.output_dir,
            max_records=args.max_records,
            verbose=args.verbose,
            **adapter_kwargs,
        )

        if result.get("success"):
            print(f"\nSuccess! Processed {result.get('records_processed', 0)} records")
            print(f"Output: {result.get('output_file', 'N/A')}")
        else:
            print(f"\nFailed: {result.get('error', 'Unknown error')}")
            sys.exit(1)
    else:
        if args.core_only:
            filter_type = "core"
        elif args.extended_only:
            filter_type = "extended"
            print("\n" + "!"*60)
            print("WARNING: Extended datasets are for RESEARCH ONLY.")
            print("Do NOT redistribute the generated data.")
            print("!"*60 + "\n")
        else:
            filter_type = "all"

        summary = run_all_adapters(
            filter_type=filter_type,
            output_base=args.output_dir,
            max_records=args.max_records,
            verbose=args.verbose,
            harmbench_csv_path=args.harmbench_csv_path,
            injecagent_data_dir=args.injecagent_data_dir,
            injecagent_attack_type=args.injecagent_attack_type,
            figstep_repo_dir=args.figstep_repo_dir,
            figstep_subset=args.figstep_subset,
            figstep_copy_images=args.figstep_copy_images,
            toolemu_repo_dir=args.toolemu_repo_dir,
            toolemu_cases_json_path=args.toolemu_cases_json_path,
            tensortrust_attacks_jsonl_path=args.tensortrust_attacks_jsonl_path,
            mmsafetybench_configs=[c.strip() for c in args.mmsafetybench_configs.split(",") if c.strip()] if args.mmsafetybench_configs else None,
            mmsafetybench_splits=[s.strip() for s in args.mmsafetybench_splits.split(",") if s.strip()] if args.mmsafetybench_splits else None,
            mmsafetybench_include_text_only=args.mmsafetybench_include_text_only,
            mmsafetybench_export_images=args.mmsafetybench_export_images,
        )

        if summary["adapters_failed"] > 0:
            sys.exit(1)


if __name__ == "__main__":
    main()