WAInjectBench

WAInjectBench è un benchmark completo per il rilevamento di prompt injection negli agenti web.
Copre 6 tipi di attacchi, attraverso due modalità: testo e immagine.

---

📂 Struttura del Dataset

• data/
  • text/
    • benign/ → 4 categorie, memorizzate come file JSONL
    • malicious/ → 8 tipi di attacco, memorizzati come file JSONL
  • image/
    • benign/ → 2 categorie, memorizzate in sottocartelle
    • malicious/ → 7 tipi di attacco, memorizzati in sottocartelle

---

⚙️ Installazione

Clona il repository e crea l'ambiente:

git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench

🚀 Valutazione

WAInjectBench supporta due pipeline di valutazione: rilevamento basato su testo e rilevamento basato su immagini.

🔹 Rilevamento basato su testo

python main_text.py \
  --data_dir [path to text dataset] \
  --detector [detector name] \
  --result_dir [output path] \
  --gpu [gpu id]

Detector disponibili: ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]

PromptArmor → richiede OPENAI_API_KEY come variabile d'ambiente.

DataSentinel →

git clone https://github.com/liu00222/Open-Prompt-Injection.git

Scarica il modello pre-addestrato in: WAInjectBench/Open-Prompt-Injection/DataSentinel_Models Imposta la directory e il percorso del modello in detector_text/datasentinel.py.

🔹 Rilevamento basato su immagini

python main_image.py \
  --data_dir [path to image dataset] \
  --detector [detector name] \
  --result_dir [output path] \
  --gpu [gpu id]

Detector disponibili: ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]

GPT-4o-Prompt → richiede OPENAI_API_KEY come variabile d'ambiente.

JailGuard →

git clone https://github.com/shiningrain/JailGuard.git

Segui il suo README per configurare MiniGPT4.

LLaVA-1.5-7B-FT → richiede il download del nostro modello fine-tuned e l'impostazione del suo percorso in detector_image/llava.py.

🔹 Generalizzazione in-domain

Forniamo anche versioni addestrate in-domain dei modelli Embedding-T e Embedding-I, disponibili in model/embedding-t/in-domain e model/embedding-i/in-domain. Per utilizzarli, segui la stessa procedura di valutazione degli esperimenti principali, ma aggiorna il percorso del modello in detector_text/embedding-t.py e detector_image/embedding-i.py.

🏋️ Addestramento

Forniamo il codice per addestrare classificatori binari basati su embedding sia per testo che per immagini.

Classificatore di embedding testuale

python train/embedding-t.py \
  --input_dir [dir with training text jsonl files] \
  --output_dir [model output path]

Formato JSONL:

{"text": "example", "label": 1}   # 1 for malicious, 0 for benign

Classificatore di embedding per immagini

python train/embedding-i.py \
  --input_dir [dir with training image jsonl files] \
  --output_dir [model output path]

Formato JSONL:

{"path": "path/to/image.png", "label": 1}

Fine-tuning di LLaVA-1.5-7B

python train.py \
  --train_jsonl train.jsonl \
  --val_jsonl val.jsonl \
  --use_lora \
  --amp_dtype bf16 \
  --device_mode single \
  --gpu_id 0

I file JSONL devono contenere i percorsi delle immagini e le etichette (1 = malicious, 0 = benign). Gli esperimenti nel nostro paper utilizzano gli iperparametri predefiniti.