scripts/adapters/b4_injecagent_adapter.py

"""
InjecAgent Dataset Adapter for B4 (Orchestrator) Schema

InjecAgent: Benchmarking Indirect Prompt Injections in Tool-Integrated LLM Agents
License: MIT
Source: https://github.com/uiuc-kang-lab/InjecAgent
Paper: https://arxiv.org/abs/2403.02691 (ACL 2024 Findings)

Maps InjecAgent test cases to the B4 Orchestrator Attack schema.
"""

import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional, Generator, List

from .base_adapter import BaseAdapter, DatasetMetadata

logger = logging.getLogger(__name__)


class InjecAgentAdapter(BaseAdapter):
    """
    Adapter for InjecAgent dataset.

    InjecAgent evaluates LLM agent vulnerability to Indirect Prompt Injection (IPI)
    attacks with 1,054 test cases spanning 17 user tools and 62 attacker tools.
    """

    ATTACK_TYPE_MAPPING = {
        'dh': 'Direct Harm',
        'ds': 'Data Stealing'
    }

    def __init__(
        self,
        output_dir: Path,
        data_dir: Path = None,
        attack_type: str = "both",
        verbose: bool = False
    ):
        """
        Initialize InjecAgent adapter.

        Args:
            output_dir: Directory for output files
            data_dir: Path to InjecAgent data directory (containing attacker_cases_*.jsonl)
            attack_type: Type of attacks to load ("dh", "ds", or "both")
            verbose: Enable verbose logging
        """
        super().__init__(output_dir, verbose)
        self.data_dir = Path(data_dir) if data_dir else None
        self.attack_type = attack_type

    @property
    def metadata(self) -> DatasetMetadata:
        return DatasetMetadata(
            name="InjecAgent",
            version="1.0",
            license="MIT",
            license_url="https://github.com/uiuc-kang-lab/InjecAgent/blob/main/LICENSE",
            source_url="https://github.com/uiuc-kang-lab/InjecAgent",
            paper_url="https://arxiv.org/abs/2403.02691",
            can_redistribute=True,
            target_surface="B4",
            description="1,054 test cases for indirect prompt injection in tool-integrated LLM agents"
        )

    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        """Load InjecAgent data from local files.

        Supports two formats:
        1. test_cases_*.json files (combined user + attacker data) - PREFERRED
        2. attacker_cases_*.jsonl + user_cases.jsonl (separate files)
        """
        if not self.data_dir:
            raise ValueError(
                "Data directory not provided. Clone InjecAgent repo from "
                "https://github.com/uiuc-kang-lab/InjecAgent and provide path to data/"
            )

        data_dir = Path(self.data_dir)
        if not data_dir.exists():
            raise FileNotFoundError(f"InjecAgent data directory not found: {data_dir}")

        files_to_load = []

        # First try to load from combined test_cases JSON files (preferred)
        if self.attack_type in ["dh", "both"]:
            # Try enhanced first, then base
            dh_enhanced = data_dir / "test_cases_dh_enhanced.json"
            dh_base = data_dir / "test_cases_dh_base.json"
            if dh_enhanced.exists():
                files_to_load.append(("dh", dh_enhanced, "json"))
            elif dh_base.exists():
                files_to_load.append(("dh", dh_base, "json"))
            else:
                # Fallback to JSONL
                dh_jsonl = data_dir / "attacker_cases_dh.jsonl"
                if dh_jsonl.exists():
                    files_to_load.append(("dh", dh_jsonl, "jsonl"))
                else:
                    logger.warning(f"No direct harm files found in {data_dir}")

        if self.attack_type in ["ds", "both"]:
            ds_enhanced = data_dir / "test_cases_ds_enhanced.json"
            ds_base = data_dir / "test_cases_ds_base.json"
            if ds_enhanced.exists():
                files_to_load.append(("ds", ds_enhanced, "json"))
            elif ds_base.exists():
                files_to_load.append(("ds", ds_base, "json"))
            else:
                ds_jsonl = data_dir / "attacker_cases_ds.jsonl"
                if ds_jsonl.exists():
                    files_to_load.append(("ds", ds_jsonl, "jsonl"))
                else:
                    logger.warning(f"No data stealing files found in {data_dir}")

        # Load user cases for context (for JSONL fallback)
        user_cases_file = data_dir / "user_cases.jsonl"
        user_cases = {}
        if user_cases_file.exists():
            with open(user_cases_file, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        case = json.loads(line)
                        user_id = case.get('User Tool', case.get('user_tool', ''))
                        user_cases[user_id] = case

        for attack_type, file_path, file_format in files_to_load:
            logger.info(f"Loading InjecAgent {attack_type} cases from {file_path}")

            if file_format == "json":
                with open(file_path, 'r', encoding='utf-8') as f:
                    records = json.load(f)
                    for record in records:
                        record['_attack_type'] = attack_type
                        yield record
            else:  # jsonl
                with open(file_path, 'r', encoding='utf-8') as f:
                    for line in f:
                        if line.strip():
                            record = json.loads(line)
                            record['_attack_type'] = attack_type
                            record['_user_cases'] = user_cases
                            yield record

    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Transform InjecAgent record to B4 schema."""
        attack_type = record.get('_attack_type', 'unknown')

        # Extract key fields from test_cases JSON format
        user_tool = record.get('User Tool', '')
        attacker_tools = record.get('Attacker Tools', [])
        attacker_tool = attacker_tools[0] if attacker_tools else ''
        user_instruction = record.get('User Instruction', '')
        injected_content = record.get('Attacker Instruction', '')
        expected_achievements = record.get('Expected Achievements', '')
        attack_type_label = record.get('Attack Type', '')
        tool_response = record.get('Tool Response', '')

        if not user_instruction or not injected_content:
            logger.warning("Missing user instruction or injected content")
            return None

        # Construct combined query (simulating IPI attack)
        combined_query = f"{user_instruction}\n\n[RETRIEVED CONTEXT]\n{injected_content}"

        record_id = self.generate_id("B4-InjecAgent", combined_query)

        attack_category = self.ATTACK_TYPE_MAPPING.get(attack_type, 'Indirect Prompt Injection')

        return {
            "attack_id": record_id,
            "category": "LLM04",  # Indirect Prompt Injection maps to OWASP LLM04
            "subtype": f"InjecAgent_{attack_type.upper()}",
            "benign_query": user_instruction,
            "malicious_injection": injected_content,
            "combined_query": combined_query,
            "expected_tool": user_tool,
            "simulated_tool": attacker_tool,
            "all_attacker_tools": attacker_tools,
            "success": True,  # InjecAgent provides attack scenarios
            "judge_note": f"InjecAgent test case - Attack type: {attack_type_label}, Expected: {expected_achievements}",
            "models_used": {
                "crafter": "InjecAgent_Benchmark",
                "simulator": "External_Evaluation"
            },
            "script_version": self.BENCHMARK_VERSION,
            "record_id": record_id,
            "timestamp": self.get_timestamp(),
            "source_dataset": "InjecAgent",
            "source_license": "MIT",
            "attack_type_detail": attack_type_label,
            "expected_achievements": expected_achievements,
            "tool_response_sample": tool_response[:500] if tool_response else None
        }


def main():
    """CLI entry point for InjecAgent adapter."""
    import argparse

    parser = argparse.ArgumentParser(description="InjecAgent to B4 Schema Adapter")
    parser.add_argument(
        "--data-dir",
        type=Path,
        required=True,
        help="Path to InjecAgent data directory"
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented/core"),
        help="Output directory"
    )
    parser.add_argument(
        "--attack-type",
        type=str,
        default="both",
        choices=["dh", "ds", "both"],
        help="Attack type to process (dh=direct harm, ds=data stealing)"
    )
    parser.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Maximum records to process"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose output"
    )

    args = parser.parse_args()

    adapter = InjecAgentAdapter(
        output_dir=args.output_dir,
        data_dir=args.data_dir,
        attack_type=args.attack_type,
        verbose=args.verbose
    )

    result = adapter.run(max_records=args.max_records)

    print(f"\nAdapter Results:")
    print(f"  Success: {result.success}")
    print(f"  Records processed: {result.records_processed}")
    print(f"  Records failed: {result.records_failed}")
    if result.output_file:
        print(f"  Output file: {result.output_file}")


if __name__ == "__main__":
    main()