Whitelist the allowed CSS classes

readme_renderer/clean.py

@@ -19,6 +19,8 @@
 import bleach.callbacks
 import bleach.linkifier
 import bleach.sanitizer
+import html5lib.filters.base
+import pygments.token
 
 
 ALLOWED_TAGS = [
@@ -54,21 +56,75 @@
     "width", "height",
 ]
 
+ALLOWED_CLASSES = {
+    "img": ["align-left", "align-center", "align-right"],
+    "span": [c for c in pygments.token.STANDARD_TYPES.values() if c],
+}
+
+
+class _CSSClassFilter(html5lib.filters.base.Filter):
+    def __init__(self, *args, allowed_classes=None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if allowed_classes is None:
+            allowed_classes = {}
+        self.allowed_classes = allowed_classes
+
+    def __iter__(self):
+        for token in super().__iter__():
+            token = self.sanitize_token(token)
+            if token:
+                yield token
+
+    def sanitize_token(self, token):
+        if token["type"] in {"StartTag", "EndTag", "EmptyTag"}:
+            name = token["name"]
+
+            if "data" in token:
+                attrs = token["data"]
+
+                if (None, "class") in attrs:
+                    new_classes = self.sanitize_css_classes(
+                        name,
+                        attrs[(None, "class")]
+                    )
+
+                    if new_classes:
+                        attrs[(None, "class")] = new_classes
+                    else:
+                        del attrs[(None, "class")]
+
+                token["data"] = attrs
+
+        return token
+
+    def sanitize_css_classes(self, name, classes):
+        classes = classes.split()
+        allowed = set(self.allowed_classes.get(name, []))
+        classes = sorted(set(classes) & allowed)
+        return " ".join(classes)
+
 
-def clean(html, tags=None, attributes=None, styles=None):
+def clean(html, tags=None, attributes=None, styles=None, classes=None):
     if tags is None:
         tags = ALLOWED_TAGS
     if attributes is None:
         attributes = ALLOWED_ATTRIBUTES
     if styles is None:
         styles = ALLOWED_STYLES
+    if classes is None:
+        classes = ALLOWED_CLASSES
 
     # Clean the output using Bleach
     cleaner = bleach.sanitizer.Cleaner(
         tags=tags,
         attributes=attributes,
         styles=styles,
         filters=[
+            # Bleach by default doesn't allow whitelisting what CSS classes
+            # are available to be used, so we'll override that behavior with
+            # our own filter which does.
+            functools.partial(_CSSClassFilter, allowed_classes=classes),
             # Bleach Linkify makes it easy to modify links, however, we will
             # not be using it to create additional links.
             functools.partial(

tests/test_clean.py

@@ -3,3 +3,8 @@
 
 def test_invalid_link():
     assert clean('<a href="http://exam](ple.com">foo</a>') == "<a>foo</a>"
+
+
+def test_css_sanitizer():
+    r = clean("<span class='foo'><img class='align-right bar'></span>")
+    assert r == '<span><img class="align-right"></span>'