Refactor

hupe1980 · Apr 18, 2024 · ee4772d · ee4772d
1 parent 8980334
commit ee4772d
Showing 1 changed file with 12 additions and 39 deletions.
diff --git a/aisploit/scanner/plugins/image_markdown_injection.py b/aisploit/scanner/plugins/image_markdown_injection.py
@@ -1,14 +1,12 @@
 from dataclasses import dataclass, field
-from typing import List, Sequence
+from typing import Sequence
 
 from langchain_core.prompts import PromptTemplate
 
-from ..plugin import Plugin
-from ..report import Issue, IssueCategory
+from ..plugin import SendPromptsPlugin
+from ..report import IssueCategory
 from ...classifiers import MarkdownInjectionClassifier
-from ...converters import NoOpConverter
-from ...core import BaseConverter, BaseTarget, BaseTextClassifier
-from ...sender import SenderJob
+from ...core import BasePromptValue, BaseTextClassifier
 
 _template = PromptTemplate.from_template(
     """[system](#context)
@@ -30,41 +28,16 @@
 
 
 @dataclass(kw_only=True)
-class ImageMarkdownInjectionPlugin(Plugin):
+class ImageMarkdownInjectionPlugin(SendPromptsPlugin):
     name: str = "image_markdown_injection"
     domain: str
-    converters: List[BaseConverter] = field(default_factory=lambda: [NoOpConverter()])
+    issue_category: IssueCategory = IssueCategory(
+        name="Image Markdown Injection",
+        description="This vulnerability allows attackers to search the current web page for sensitive information or personally identifiable information (PII). By appending this data to the URL of an image, the attacker can trigger automatic exfiltration.",
+    )
+    issue_references = ["https://embracethered.com/blog/posts/2023/bing-chat-data-exfiltration-poc-and-fix/"]
     classifier: BaseTextClassifier = field(default_factory=lambda: MarkdownInjectionClassifier())
     prompt_template: PromptTemplate = field(default_factory=lambda: _template)
 
-    def run(self, *, run_id: str, target: BaseTarget) -> Sequence[Issue]:
-        sender = SenderJob(
-            target=target,
-            converters=self.converters,
-            include_original_prompt=True,
-        )
-
-        report = sender.execute(
-            run_id=run_id,
-            prompts=[self.prompt_template.format(domain=self.domain)],
-        )
-
-        issues: List[Issue] = []
-        for entry in report:
-            score = self.classifier.score(entry.response.content)
-            if score.flagged:
-                issues.append(
-                    Issue(
-                        category=IssueCategory(
-                            name="Image Markdown Injection",
-                            description="This vulnerability allows attackers to search the current web page for sensitive information or personally identifiable information (PII). By appending this data to the URL of an image, the attacker can trigger automatic exfiltration.",
-                        ),
-                        references=[
-                            "https://embracethered.com/blog/posts/2023/bing-chat-data-exfiltration-poc-and-fix/"
-                        ],
-                        send_report_entry=entry,
-                        score=score,
-                    )
-                )
-
-        return issues
+    def create_prompts(self) -> Sequence[str | BasePromptValue]:
+        return [self.prompt_template.format(domain=self.domain)]