Tianyang
13789cc0b4
- main.py: remove duplicate /health endpoint definition - text_extract.py: HTML-escape extracted plain text before indexing so raw markup in uploaded documents can't render as live HTML via search snippets (stored XSS fix) - document_service.py: base64-encode content when UTF-8 decoding fails instead of raising UnicodeDecodeError (500) on binary files; add upload validation for file extension (allowlist) and size (10MB max) - routers/documents.py: map new validation errors to HTTP 400 with a clear detail message instead of letting them 500 - add/extend tests covering escaping, binary get_document, and upload validation rejection + acceptance paths
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import base64
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from app.text_extract import html_to_text, markdown_to_text
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ALLOWED_UPLOAD_EXTENSIONS = {".md", ".html", ".htm", ".png", ".jpg", ".jpeg", ".gif", ".webp"}
|
|
MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
|
|
|
|
|
|
class DocumentNotFoundError(Exception):
|
|
pass
|
|
|
|
|
|
class UnsupportedFileTypeError(Exception):
|
|
def __init__(self, extension: str):
|
|
self.extension = extension
|
|
super().__init__(f"unsupported file type: {extension}")
|
|
|
|
|
|
class FileTooLargeError(Exception):
|
|
def __init__(self, max_size_bytes: int):
|
|
self.max_size_bytes = max_size_bytes
|
|
super().__init__(f"file too large: max {max_size_bytes // (1024 * 1024)}MB")
|
|
|
|
|
|
def _extension_of(key: str) -> str:
|
|
name = key.rsplit("/", 1)[-1]
|
|
if "." not in name:
|
|
return ""
|
|
return "." + name.rsplit(".", 1)[-1].lower()
|
|
|
|
|
|
def _guess_content_type(key: str) -> str:
|
|
if key.endswith(".md"):
|
|
return "text/markdown"
|
|
if key.endswith(".html") or key.endswith(".htm"):
|
|
return "text/html"
|
|
return "application/octet-stream"
|
|
|
|
|
|
class DocumentService:
|
|
def __init__(self, minio_client, search_client):
|
|
self._minio_client = minio_client
|
|
self._search_client = search_client
|
|
|
|
def list_documents(self, path: str) -> dict:
|
|
prefix = path if not path or path.endswith("/") else f"{path}/"
|
|
entries = self._minio_client.list_objects(prefix)
|
|
folders = []
|
|
files = []
|
|
for entry in entries:
|
|
key = entry["key"]
|
|
if entry.get("is_dir"):
|
|
name = key[len(prefix):].rstrip("/")
|
|
if name:
|
|
folders.append(name)
|
|
else:
|
|
name = key[len(prefix):]
|
|
files.append({
|
|
"key": key,
|
|
"name": name,
|
|
"size": entry["size"],
|
|
"modifiedAt": entry["modifiedAt"],
|
|
})
|
|
return {"folders": folders, "files": files}
|
|
|
|
def get_document(self, key: str) -> dict:
|
|
try:
|
|
data = self._minio_client.get_object(key)
|
|
except FileNotFoundError as exc:
|
|
raise DocumentNotFoundError(key) from exc
|
|
try:
|
|
content = data.decode("utf-8")
|
|
encoding = "utf-8"
|
|
except UnicodeDecodeError:
|
|
content = base64.b64encode(data).decode("ascii")
|
|
encoding = "base64"
|
|
return {
|
|
"key": key,
|
|
"content": content,
|
|
"contentType": _guess_content_type(key),
|
|
"encoding": encoding,
|
|
}
|
|
|
|
def _extract_text(self, content: bytes, content_type: str) -> str:
|
|
text = content.decode("utf-8", errors="ignore")
|
|
if content_type == "text/html":
|
|
return html_to_text(text)
|
|
if content_type == "text/markdown":
|
|
return markdown_to_text(text)
|
|
return text
|
|
|
|
def save_document(self, key: str, content: bytes, content_type: str) -> dict:
|
|
self._minio_client.put_object(key, content, content_type)
|
|
try:
|
|
text = self._extract_text(content, content_type)
|
|
title = key.rsplit("/", 1)[-1]
|
|
path = key.rsplit("/", 1)[0] if "/" in key else ""
|
|
self._search_client.index_document(key, path, title, text)
|
|
except Exception:
|
|
logger.exception("Failed to update search index for %s", key)
|
|
return {"key": key, "updatedAt": datetime.now(timezone.utc).isoformat()}
|
|
|
|
def delete_document(self, key: str) -> None:
|
|
self._minio_client.delete_object(key)
|
|
try:
|
|
self._search_client.delete_document(key)
|
|
except Exception:
|
|
logger.exception("Failed to delete search index entry for %s", key)
|
|
|
|
def upload_document(self, key: str, data: bytes, content_type: str) -> dict:
|
|
extension = _extension_of(key)
|
|
if extension not in ALLOWED_UPLOAD_EXTENSIONS:
|
|
raise UnsupportedFileTypeError(extension)
|
|
if len(data) > MAX_UPLOAD_SIZE_BYTES:
|
|
raise FileTooLargeError(MAX_UPLOAD_SIZE_BYTES)
|
|
self.save_document(key, data, content_type)
|
|
name = key.rsplit("/", 1)[-1]
|
|
return {"key": key, "name": name, "size": len(data)}
|