Fix xss.

src/snek/service/channel_message.py

@@ -69,10 +69,10 @@ class ChannelMessageService(BaseService):
                 "color": user["color"],
             }
         )
+        context['message'] = whitelist_attributes(context['message'])
         try:
             template = self.app.jinja2_env.get_template("message.html")
             model["html"] = template.render(**context)
-            model["html"] = whitelist_attributes(model["html"])
         except Exception as ex:
             print(ex, flush=True)
 
@@ -118,6 +118,7 @@ class ChannelMessageService(BaseService):
     async def save(self, model):
         context = {}
         context.update(model.record)
+        context['message'] = whitelist_attributes(context['message'])
         user = await self.app.services.user.get(model["user_uid"])
         context.update(
             {
@@ -129,7 +130,6 @@ class ChannelMessageService(BaseService):
         )
         template = self.app.jinja2_env.get_template("message.html")
         model["html"] = template.render(**context)
-        model["html"] = whitelist_attributes(model["html"])
         return await super().save(model)
 
     async def offset(self, channel_uid, page=0, timestamp=None, page_size=30):

src/snek/system/template.py

@@ -79,44 +79,12 @@ emoji.EMOJI_DATA[
 ] = {"en": ":a1:", "status": 2, "E": 0.6, "alias": [":a1:"]}
 
 
-ALLOWED_TAGS = list(bleach.sanitizer.ALLOWED_TAGS) + [
-    "img",
-    "video",
-    "audio",
-    "source",
-    "iframe",
-    "picture",
-    "span",
-]
-ALLOWED_ATTRIBUTES = {
-    **bleach.sanitizer.ALLOWED_ATTRIBUTES,
-    "img": ["src", "alt", "title", "width", "height"],
-    "a": ["href", "title", "target", "rel", "referrerpolicy", "class"],
-    "iframe": [
-        "src",
-        "width",
-        "height",
-        "frameborder",
-        "allow",
-        "allowfullscreen",
-        "title",
-        "referrerpolicy",
-        "style",
-    ],
-    "video": ["src", "controls", "width", "height"],
-    "audio": ["src", "controls"],
-    "source": ["src", "type"],
-    "span": ["class"],
-    "picture": [],
-}
-
+ALLOWED_TAGS = list(bleach.sanitizer.ALLOWED_TAGS) + ["picture"]
 
 def sanitize_html(value):
     return bleach.clean(
         value,
-        tags=ALLOWED_TAGS,
-        attributes=ALLOWED_ATTRIBUTES,
-        protocols=bleach.sanitizer.ALLOWED_PROTOCOLS + ["data"],
+        protocols=list(bleach.sanitizer.ALLOWED_PROTOCOLS) + ["data"],
         strip=True,
     )
 
@@ -132,50 +100,8 @@ def set_link_target_blank(text):
 
     return str(soup)
 
-
-SAFE_ATTRIBUTES = {
-    "href",
-    "src",
-    "alt",
-    "title",
-    "width",
-    "height",
-    "style",
-    "id",
-    "class",
-    "rel",
-    "type",
-    "name",
-    "value",
-    "placeholder",
-    "aria-hidden",
-    "aria-label",
-    "srcset",
-    "target",
-    "rel",
-    "referrerpolicy",
-    "controls",
-    "frameborder",
-    "allow",
-    "allowfullscreen",
-    "referrerpolicy",
-}
-
-
 def whitelist_attributes(html):
-    soup = BeautifulSoup(html, "html.parser")
-
-    for tag in soup.find_all():
-        if hasattr(tag, "attrs"):
-            if tag.name in ["script", "form", "input"]:
-                tag.replace_with("")
-                continue
-            attrs = dict(tag.attrs)
-            for attr in list(attrs):
-                # Check if attribute is in the safe list or is a data-* attribute
-                if not (attr in SAFE_ATTRIBUTES or attr.startswith("data-")):
-                    del tag.attrs[attr]
-    return str(soup)
+    return sanitize_html(html)
 
 
 def embed_youtube(text):