Upload html_sanitizer.py

src/services/html_sanitizer.py

@@ -1,79 +1,119 @@
-"""Small allow-list sanitizer for answer fragments returned to the browser."""
-
-from html import escape
+"""Small allow-list sanitizer for answer fragments returned to the browser."""
+
+from html import escape, unescape
 from html.parser import HTMLParser
+import re
+from typing import Collection
 from urllib.parse import urlparse
+
+
+ALLOWED_TAGS = {"a", "b", "blockquote", "br", "code", "details", "div", "em", "h3", "h4", "hr", "i", "li", "ol", "p", "pre", "span", "strong", "summary", "table", "tbody", "td", "th", "thead", "tr", "ul"}
+VOID_TAGS = {"br", "hr"}
+DROP_CONTENT_TAGS = {"button", "form", "iframe", "input", "math", "object", "script", "style", "svg"}
+GLOBAL_ATTRIBUTES = {"aria-controls", "aria-expanded", "aria-label", "class", "role"}
+TAG_ATTRIBUTES = {"a": {"href", "rel", "target", "title"}, "td": {"colspan", "rowspan"}, "th": {"colspan", "rowspan", "scope"}}
+
+
+_VISIBLE_URL_RE = re.compile(r"(?:https?://|mailto:|tel:)[^\s<>\"']+")
+_MARKDOWN_LINK_RE = re.compile(
+    r"\[([^\]\n]+)\]\(\s*((?:https?://|mailto:|tel:)[^\s)]+)\s*\)"
+)
 
 
-ALLOWED_TAGS = {"a", "b", "blockquote", "br", "code", "details", "div", "em", "h3", "h4", "hr", "i", "li", "ol", "p", "pre", "span", "strong", "summary", "table", "tbody", "td", "th", "thead", "tr", "ul"}
-VOID_TAGS = {"br", "hr"}
-DROP_CONTENT_TAGS = {"button", "form", "iframe", "input", "math", "object", "script", "style", "svg"}
-GLOBAL_ATTRIBUTES = {"aria-controls", "aria-expanded", "aria-label", "class", "role"}
-TAG_ATTRIBUTES = {"a": {"href", "rel", "target", "title"}, "td": {"colspan", "rowspan"}, "th": {"colspan", "rowspan", "scope"}}
+def _normalise_url(value: str) -> str:
+    return unescape(value).strip().rstrip(".,;:!?)]}")
 
 
-def _safe_href(value: str) -> str | None:
-    candidate = value.strip()
+def _safe_href(value: str, allowed_urls: set[str] | None = None) -> str | None:
+    candidate = _normalise_url(value)
     if not candidate:
         return None
     parsed = urlparse(candidate)
     if parsed.scheme.lower() not in {"https", "mailto", "tel"}:
         return None
+    if allowed_urls is not None and candidate not in allowed_urls:
+        return None
     return candidate
 
 
+def markdown_to_safe_html(value: str) -> str:
+    """Convert the small Markdown subset commonly emitted by an LLM.
+
+    LibBee's API contract is sanitized HTML. This prevents Markdown markers such
+    as ``**bold**`` and ``[label](url)`` from leaking into the browser as text.
+    """
+    value = value or ""
+    value = _MARKDOWN_LINK_RE.sub(
+        lambda m: f'<a href="{escape(m.group(2), quote=True)}" target="_blank">{escape(m.group(1))}</a>',
+        value,
+    )
+    value = re.sub(r"\*\*([^*\n]+)\*\*", r"<strong>\1</strong>", value)
+    value = re.sub(r"(?m)^\s{0,3}#{1,6}\s+(.+)$", r"<strong>\1</strong><br>", value)
+    return value
+
+
 class _Sanitizer(HTMLParser):
-    def __init__(self) -> None:
+    def __init__(self, allowed_urls: Collection[str] | None = None) -> None:
         super().__init__(convert_charrefs=True)
         self.output: list[str] = []
         self.drop_depth = 0
-
-    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
-        tag = tag.lower()
-        if tag in DROP_CONTENT_TAGS:
-            self.drop_depth += 1
-            return
-        if self.drop_depth or tag not in ALLOWED_TAGS:
-            return
-        rendered: list[str] = []
-        for raw_name, raw_value in attrs:
-            name = raw_name.lower()
-            if name.startswith("on") or name == "style":
-                continue
-            if name not in GLOBAL_ATTRIBUTES and name not in TAG_ATTRIBUTES.get(tag, set()):
-                continue
+        self.allowed_urls = (
+            {_normalise_url(url) for url in allowed_urls if _normalise_url(url)}
+            if allowed_urls is not None
+            else None
+        )
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        tag = tag.lower()
+        if tag in DROP_CONTENT_TAGS:
+            self.drop_depth += 1
+            return
+        if self.drop_depth or tag not in ALLOWED_TAGS:
+            return
+        rendered: list[str] = []
+        for raw_name, raw_value in attrs:
+            name = raw_name.lower()
+            if name.startswith("on") or name == "style":
+                continue
+            if name not in GLOBAL_ATTRIBUTES and name not in TAG_ATTRIBUTES.get(tag, set()):
+                continue
             value = raw_value or ""
             if tag == "a" and name == "href":
-                value = _safe_href(value) or ""
-                if not value:
-                    continue
-            if tag == "a" and name == "target" and value != "_blank":
-                continue
-            rendered.append(f' {name}="{escape(value, quote=True)}"')
-        if tag == "a":
-            rendered = [item for item in rendered if not item.startswith(" rel=")]
-            rendered.append(' rel="noopener noreferrer"')
-        self.output.append(f"<{tag}{''.join(rendered)}>")
-
-    def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
-        self.handle_starttag(tag, attrs)
-
-    def handle_endtag(self, tag: str) -> None:
-        tag = tag.lower()
-        if tag in DROP_CONTENT_TAGS:
-            if self.drop_depth:
-                self.drop_depth -= 1
-            return
-        if not self.drop_depth and tag in ALLOWED_TAGS and tag not in VOID_TAGS:
-            self.output.append(f"</{tag}>")
-
+                value = _safe_href(value, self.allowed_urls) or ""
+                if not value:
+                    continue
+            if tag == "a" and name == "target" and value != "_blank":
+                continue
+            rendered.append(f' {name}="{escape(value, quote=True)}"')
+        if tag == "a":
+            rendered = [item for item in rendered if not item.startswith(" rel=")]
+            rendered.append(' rel="noopener noreferrer"')
+        self.output.append(f"<{tag}{''.join(rendered)}>")
+
+    def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        self.handle_starttag(tag, attrs)
+
+    def handle_endtag(self, tag: str) -> None:
+        tag = tag.lower()
+        if tag in DROP_CONTENT_TAGS:
+            if self.drop_depth:
+                self.drop_depth -= 1
+            return
+        if not self.drop_depth and tag in ALLOWED_TAGS and tag not in VOID_TAGS:
+            self.output.append(f"</{tag}>")
+
     def handle_data(self, data: str) -> None:
         if not self.drop_depth:
+            if self.allowed_urls is not None:
+                data = _VISIBLE_URL_RE.sub(
+                    lambda m: m.group(0) if _normalise_url(m.group(0)) in self.allowed_urls else "",
+                    data,
+                )
             self.output.append(escape(data, quote=False))
-
-
-def sanitize_answer_html(value: str) -> str:
-    parser = _Sanitizer()
+
+
+def sanitize_answer_html(value: str, allowed_urls: Collection[str] | None = None) -> str:
+    parser = _Sanitizer(allowed_urls)
     parser.feed(value or "")
     parser.close()
     return "".join(parser.output)