Browse Source

Prevent SSRF and HTML injection

Rodrigo Agundez 2 months ago
parent
commit
167c8bf00d
1 changed files with 7 additions and 6 deletions
  1. 7 6
      backend/open_webui/utils/pdf_generator.py

+ 7 - 6
backend/open_webui/utils/pdf_generator.py

@@ -2,6 +2,7 @@ from datetime import datetime
 from io import BytesIO
 from io import BytesIO
 from pathlib import Path
 from pathlib import Path
 from typing import Dict, Any, List
 from typing import Dict, Any, List
+from html import escape
 
 
 from markdown import markdown
 from markdown import markdown
 
 
@@ -11,7 +12,6 @@ from fpdf import FPDF
 from open_webui.env import STATIC_DIR, FONTS_DIR
 from open_webui.env import STATIC_DIR, FONTS_DIR
 from open_webui.models.chats import ChatTitleMessagesForm
 from open_webui.models.chats import ChatTitleMessagesForm
 
 
-
 class PDFGenerator:
 class PDFGenerator:
     """
     """
     Description:
     Description:
@@ -41,13 +41,13 @@ class PDFGenerator:
 
 
     def _build_html_message(self, message: Dict[str, Any]) -> str:
     def _build_html_message(self, message: Dict[str, Any]) -> str:
         """Build HTML for a single message."""
         """Build HTML for a single message."""
-        role = message.get("role", "user")
-        content = message.get("content", "")
+        role = escape(message.get("role", "user"))
+        content = escape(message.get("content", ""))
         timestamp = message.get("timestamp")
         timestamp = message.get("timestamp")
 
 
-        model = message.get("model") if role == "assistant" else ""
+        model = escape(message.get("model") if role == "assistant" else "")
 
 
-        date_str = self.format_timestamp(timestamp) if timestamp else ""
+        date_str = escape(self.format_timestamp(timestamp) if timestamp else "")
 
 
         # extends pymdownx extension to convert markdown to html.
         # extends pymdownx extension to convert markdown to html.
         # - https://facelessuser.github.io/pymdown-extensions/usage_notes/
         # - https://facelessuser.github.io/pymdown-extensions/usage_notes/
@@ -76,6 +76,7 @@ class PDFGenerator:
 
 
     def _generate_html_body(self) -> str:
     def _generate_html_body(self) -> str:
         """Generate the full HTML body for the PDF."""
         """Generate the full HTML body for the PDF."""
+        escaped_title = escape(self.form_data.title)
         return f"""
         return f"""
         <html>
         <html>
             <head>
             <head>
@@ -84,7 +85,7 @@ class PDFGenerator:
             <body>
             <body>
             <div>
             <div>
                 <div>
                 <div>
-                    <h2>{self.form_data.title}</h2>
+                    <h2>{escaped_title}</h2>
                     {self.messages_html}
                     {self.messages_html}
                 </div>
                 </div>
             </div>
             </div>