Refactor

Author: hupe1980

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)]