make payload names configurable

leondz · Oct 2, 2024 · 5b3edda · 5b3edda
1 parent cd8aa84
commit 5b3edda
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/garak/probes/latentinjection.py b/garak/probes/latentinjection.py
@@ -95,7 +95,7 @@ class LatentInjectionTranslationEnFr(Probe, LatentInjectionMixin):
     payloads = [f"{TRIGGER_MARKER}"]
     payload_triggers = ["Haha pwned!"]
 
-    payload_name = "rude_french"
+    DEFAULT_PARAMS = Probe.DEFAULT_PARAMS | {"payload_name": "rude_french"}
 
     def __init__(self, config_root=_config):
         super().__init__(config_root)
@@ -258,10 +258,14 @@ class LatentInjectionReport(Probe, LatentInjectionMixin):
     ]
     payload_triggers = []
 
+    DEFAULT_PARAMS = Probe.DEFAULT_PARAMS | {
+        "domains_payload_name": "domains_latentinjection"
+    }
+
     def __init__(self, config_root=_config):
         super().__init__(config_root)
         self.payload_triggers = (
-            garak.payloads.Director().load("domains_latentinjection").payloads
+            garak.payloads.Director().load(self.domains_payload_name).payloads
         )
         self._build_prompts_triggers()
 

diff --git a/garak/probes/xss.py b/garak/probes/xss.py
@@ -47,10 +47,16 @@ class MarkdownImageExfil(Probe):
         "quality:Security:ExtractionInversion",
     ]
 
+    DEFAULT_PARAMS = Probe.DEFAULT_PARAMS | {
+        "domains_payload_name": "example_domains_xss"
+    }
+
     def __init__(self, config_root=_config):
         super().__init__(config_root=config_root)
         self.prompts = []
-        for domain in garak.payloads.Director().load("example_domains_xss").payloads:
+        for domain in (
+            garak.payloads.Director().load(self.domains_payload_name).payloads
+        ):
             for encoding in self.encodings:
                 self.prompts.append(
                     self.template.replace("{domain}", domain).replace(