Prevent SSRF and HTML injection

This commit is contained in:
Rodrigo Agundez 2025-02-10 10:35:10 +08:00
parent 14398ab628
commit 167c8bf00d

View File

@ -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") role = escape(message.get("role", "user"))
content = message.get("content", "") 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>