mirror of https://github.com/langgenius/dify.git
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import os
|
|
from email.message import Message
|
|
from urllib.parse import quote
|
|
|
|
from flask import Response
|
|
|
|
HTML_MIME_TYPES = frozenset({"text/html", "application/xhtml+xml"})
|
|
HTML_EXTENSIONS = frozenset({"html", "htm"})
|
|
|
|
|
|
def _normalize_mime_type(mime_type: str | None) -> str:
|
|
if not mime_type:
|
|
return ""
|
|
message = Message()
|
|
message["Content-Type"] = mime_type
|
|
return message.get_content_type().strip().lower()
|
|
|
|
|
|
def _is_html_extension(extension: str | None) -> bool:
|
|
if not extension:
|
|
return False
|
|
return extension.lstrip(".").lower() in HTML_EXTENSIONS
|
|
|
|
|
|
def is_html_content(mime_type: str | None, filename: str | None, extension: str | None = None) -> bool:
|
|
normalized_mime_type = _normalize_mime_type(mime_type)
|
|
if normalized_mime_type in HTML_MIME_TYPES:
|
|
return True
|
|
|
|
if _is_html_extension(extension):
|
|
return True
|
|
|
|
if filename:
|
|
return _is_html_extension(os.path.splitext(filename)[1])
|
|
|
|
return False
|
|
|
|
|
|
def enforce_download_for_html(
|
|
response: Response,
|
|
*,
|
|
mime_type: str | None,
|
|
filename: str | None,
|
|
extension: str | None = None,
|
|
) -> bool:
|
|
if not is_html_content(mime_type, filename, extension):
|
|
return False
|
|
|
|
if filename:
|
|
encoded_filename = quote(filename)
|
|
response.headers["Content-Disposition"] = f"attachment; filename*=UTF-8''{encoded_filename}"
|
|
else:
|
|
response.headers["Content-Disposition"] = "attachment"
|
|
|
|
response.headers["Content-Type"] = "application/octet-stream"
|
|
response.headers["X-Content-Type-Options"] = "nosniff"
|
|
return True
|