Repository: mozilla/bleach

Issue:
improved callable filter support
this is perhaps a bad title.  i'll explain the problem:

as mentioned in an earlier ticket, i needed to support the  "data-*" namespace on attributes.  the suggested route is callable filters.

the two main problems are as follows
- `*` doesn't accept callable filters.  based on the current setup, it requires an iterable.
- specific tags can have a callable filter OR a list of attributes.  not both.

both of these are invalid:

```
macro_attributes_ok = {
    '*': allowable_macro_attributes,
    'a' : ['href', 'title', 'target', 'rel', allowable_macro_attributes],
}
```

the callable_filters, as-is, must be implemented with a dedicated function per TAG.  This gets messy as a lot of code must be written.

```
macro_attributes_ok = {
    'a' : allowable_attributes_a,
    'div' : allowable_attributes_div,
    'span' : allowable_attributes_span,
}
```

approaches that could be nicer / easier to deal with:

• allow for callables to be within a list of text attributes

```
macro_attributes_ok = {
    'a' : [ 'href', 'link', allowable_attributes_setA ],
    'div' : [ 'class', 'style', allowable_attributes_setA ],
}
```

• pass the tag's name in as a kwarg to the callable.  this would allow callables to reflect the current tag -- and be backwards compatible.  

```
def filter_src(name, value, tag=None):
    if tag != 'img':
         return False
    if name in ('alt', 'height', 'width'):
        return True
```
-  the `attributes` kwarg to clean could be a callable itself; one that expects a triplet of data
  
     def validate_attribute( tag, attribute, value ):
           return True
  
    clean( text, attributes=validate_attribute)


Discussion:

Patch:
diff --git a/CHANGES b/CHANGES
index 3ff2b78..79f56a9 100644
--- a/CHANGES
+++ b/CHANGES
@@ -25,6 +25,12 @@ Version 2.0 (in development)
   Amongst other things, this version will add end tags even if the tag in
   question is to be escaped.
 
+* ``bleach.clean`` and friends attribute callables now take three arguments:
+  tag, attribute name and attribute value. Previously they only took attribute
+  name and attribute value.
+
+  All attribute callables will need to be updated.
+
 * ``bleach.linkify`` was rewritten
 
   ``linkify`` was reimplemented as an html5lib Filter. As such, it no longer
@@ -52,6 +58,8 @@ Version 2.0 (in development)
   don't then html5lib will raise an assertion error that the value is not
   unicode.
 
+  All linkify filters will need to be updated.
+
 **Changes**
 
 * Supports Python 3.6.
@@ -68,6 +76,8 @@ Version 2.0 (in development)
   you can pass as a filter to ``bleach.sanitizer.Cleaner`` allowing you to clean
   and linkify in one pass.
 
+* ``bleach.clean`` and friends can now take a callable as an attributes arg value.
+
 * Tons of bug fixes.
 
 * Cleaned up tests.
diff --git a/bleach/sanitizer.py b/bleach/sanitizer.py
index 06c9066..1223e79 100644
--- a/bleach/sanitizer.py
+++ b/bleach/sanitizer.py
@@ -128,7 +128,7 @@ class Cleaner(object):
             source=self.walker(dom),
 
             # Bleach-sanitizer-specific things
-            allowed_attributes_map=self.attributes,
+            attributes=self.attributes,
             strip_disallowed_elements=self.strip,
             strip_html_comments=self.strip_comments,
 
@@ -146,22 +146,58 @@ class Cleaner(object):
         return self.serializer.render(filtered)
 
 
+def attribute_filter_factory(attributes):
+    """Generates attribute filter function for the given attributes value
+
+    The attributes value can take one of several shapes. This returns a filter
+    function appropriate to the attributes value. One nice thing about this is
+    that there's less if/then shenanigans in the ``allow_token`` method.
+
+    """
+    if callable(attributes):
+        return attributes
+
+    if isinstance(attributes, dict):
+        def _attr_filter(tag, attr, value):
+            if tag in attributes:
+                attr_val = attributes[tag]
+                if callable(attr_val):
+                    return attr_val(tag, attr, value)
+
+                if attr in attr_val:
+                    return True
+
+            if '*' in attributes:
+                attr_val = attributes['*']
+                if callable(attr_val):
+                    return attr_val(tag, attr, value)
+
+                return attr in attr_val
+
+            return False
+
+        return _attr_filter
+
+    if isinstance(attributes, list):
+        def _attr_filter(tag, attr, value):
+            return attr in attributes
+
+        return _attr_filter
+
+    raise ValueError('attributes needs to be a callable, a list or a dict')
+
+
 class BleachSanitizerFilter(sanitizer.Filter):
     """html5lib Filter that sanitizes text
 
     This filter can be used anywhere html5lib filters can be used.
 
     """
-    def __init__(self, source, allowed_attributes_map,
+    def __init__(self, source, attributes=ALLOWED_ATTRIBUTES,
                  strip_disallowed_elements=False, strip_html_comments=True,
                  **kwargs):
 
-        if isinstance(allowed_attributes_map, dict):
-            self.wildcard_attributes = allowed_attributes_map.get('*', [])
-            self.allowed_attributes_map = allowed_attributes_map
-        else:
-            self.wildcard_attributes = allowed_attributes_map
-            self.allowed_attributes_map = {}
+        self.attr_filter = attribute_filter_factory(attributes)
 
         self.strip_disallowed_elements = strip_disallowed_elements
         self.strip_html_comments = strip_html_comments
@@ -205,10 +241,6 @@ class BleachSanitizerFilter(sanitizer.Filter):
     def allow_token(self, token):
         """Handles the case where we're allowing the tag"""
         if 'data' in token:
-            allowed_attributes = self.allowed_attributes_map.get(token['name'], [])
-            if not callable(allowed_attributes):
-                allowed_attributes += self.wildcard_attributes
-
             # Loop through all the attributes and drop the ones that are not
             # allowed, are unsafe or break other rules. Additionally, fix
             # attribute values that need fixing.
@@ -220,11 +252,10 @@ class BleachSanitizerFilter(sanitizer.Filter):
                 namespace, name = namespaced_name
 
                 # Drop attributes that are not explicitly allowed
-                if callable(allowed_attributes):
-                    if not allowed_attributes(name, val):
-                        continue
-
-                elif name not in allowed_attributes:
+                #
+                # NOTE(willkg): We pass in the attribute name--not a namespaced
+                # name.
+                if not self.attr_filter(token['name'], name, val):
                     continue
 
                 # Look at attributes that have uri values
diff --git a/docs/clean.rst b/docs/clean.rst
index 161e435..b02c452 100644
--- a/docs/clean.rst
+++ b/docs/clean.rst
@@ -55,8 +55,8 @@ The default value is also a conservative dict found in
 As a list
 ---------
 
-The ``attributes`` value can be a list, in which case the attributes are allowed
-for any tag.
+The ``attributes`` value can be a list which specifies the list of attributes
+allowed for any tag.
 
 For example:
 
@@ -76,10 +76,12 @@ For example:
 As a dict
 ---------
 
-The ``attributes`` value can be a dict, in which case the keys are tag names (or
-a wildcard: ``*`` for all tags) and the values are lists of allowed attributes.
+The ``attributes`` value can be a dict which maps tags to what attributes they can have.
 
-For example:
+You can also specify ``*``, which will match any tag.
+
+For example, this allows "href" and "rel" for "a" tags, "alt" for the "img" tag
+and "class" for any tag (including "a" and "img"):
 
 .. doctest::
 
@@ -99,48 +101,66 @@ For example:
    u'<img alt="an example">'
 
 
-In this case, ``class`` is allowed on any allowed element (from the ``tags``
-argument), ``<a>`` tags are allowed to have ``href`` and ``rel`` attributes,
-and so on.
-
-
 Using functions
 ---------------
 
-You can also use callables. If the callable returns ``True``, the attribute is
-allowed. Otherwise, it is stripped. For example:
+You can also use callables that take the tag, attribute name and attribute value
+and returns ``True`` to keep the attribute or ``False`` to drop it.
+
+You can pass a callable as the attributes argument value and it'll run for
+every tag/attr.
+
+For example:
+
+.. doctest::
+
+   >>> import bleach
+
+   >>> def allow_h(tag, name, value):
+   ...     return name[0] == 'h'
+
+   >>> bleach.clean(
+   ...    u'<a href="http://example.com" title="link">link</a>',
+   ...    tags=['a'],
+   ...    attributes=allow_h,
+   ... )
+   u'<a href="http://example.com">link</a>'
+
+
+You can also pass a callable as a value in an attributes dict and it'll run for
+attributes for specified tags:
 
 .. doctest::
 
-    >>> from urlparse import urlparse
-    >>> import bleach
+   >>> from urlparse import urlparse
+   >>> import bleach
 
-    >>> def allow_src(name, value):
-    ...     if name in ('alt', 'height', 'width'):
-    ...         return True
-    ...     if name == 'src':
-    ...         p = urlparse(value)
-    ...         return (not p.netloc) or p.netloc == 'mydomain.com'
-    ...     return False
+   >>> def allow_src(tag, name, value):
+   ...     if name in ('alt', 'height', 'width'):
+   ...         return True
+   ...     if name == 'src':
+   ...         p = urlparse(value)
+   ...         return (not p.netloc) or p.netloc == 'mydomain.com'
+   ...     return False
 
-    >>> bleach.clean(
-    ...    u'<img src="http://example.com" alt="an example">',
-    ...    tags=['img'],
-    ...    attributes={
-    ...        'img': allow_src
-    ...    }
-    ... )
-    u'<img alt="an example">'
+   >>> bleach.clean(
+   ...    u'<img src="http://example.com" alt="an example">',
+   ...    tags=['img'],
+   ...    attributes={
+   ...        'img': allow_src
+   ...    }
+   ... )
+   u'<img alt="an example">'
 
 
 Allowed styles (``styles``)
 ===========================
 
-If you allow the ``style`` attribute, you will also need to whitelist styles
-users are allowed to set, for example ``color`` and ``background-color``.
+If you allow the ``style`` attribute, you will also need to specify the allowed
+styles users are allowed to set, for example ``color`` and ``background-color``.
 
-The default value is an empty list, i.e., the ``style`` attribute will be
-allowed but no values will be.
+The default value is an empty list. In other words, the ``style`` attribute will
+be allowed but no style declaration names will be allowed.
 
 For example, to allow users to set the color and font-weight of text:
 
@@ -205,8 +225,8 @@ Default protocols are in ``bleach.ALLOWED_PROTOCOLS``.
 Stripping markup (``strip``)
 ============================
 
-By default, Bleach *escapes* tags that aren't specified in the tags
-whitelist and invalid markup. For example:
+By default, Bleach *escapes* tags that aren't specified in the allowed tags list
+and invalid markup. For example:
 
 .. doctest::
 


Test patch:
diff --git a/tests/test_basics.py b/tests/test_basics.py
index 031ab66..5b59ebf 100644
--- a/tests/test_basics.py
+++ b/tests/test_basics.py
@@ -164,23 +164,46 @@ class TestClean:
         clean = '<em class="FOO">BAR</em>'
         assert bleach.clean(dirty, attributes=['class']) == clean
 
-    def test_wildcard_attributes(self):
+    def test_attributes_callable(self):
+        """Verify attributes can take a callable"""
+        ATTRS = lambda tag, name, val: name == 'title'
+        TAGS = ['a']
+
+        assert (
+            bleach.clean(u'<a href="/foo" title="blah">example</a>', tags=TAGS, attributes=ATTRS) ==
+            u'<a title="blah">example</a>'
+        )
+
+    def test_attributes_wildcard(self):
+        """Verify attributes[*] works"""
         ATTRS = {
             '*': ['id'],
             'img': ['src'],
         }
-        TAG = ['img', 'em']
+        TAGS = ['img', 'em']
         dirty = ('both <em id="foo" style="color: black">can</em> have '
                  '<img id="bar" src="foo"/>')
         assert (
-            bleach.clean(dirty, tags=TAG, attributes=ATTRS) ==
+            bleach.clean(dirty, tags=TAGS, attributes=ATTRS) ==
             'both <em id="foo">can</em> have <img id="bar" src="foo">'
         )
 
-    def test_callable_attributes(self):
-        """Verify callable attributes work and get correct arg values"""
-        def img_test(attr, val):
-            return attr == 'src' and val.startswith('https')
+    def test_attributes_wildcard_callable(self):
+        """Verify attributes[*] callable works"""
+        ATTRS = {
+            '*': lambda tag, name, val: name == 'title'
+        }
+        TAGS = ['a']
+
+        assert (
+            bleach.clean(u'<a href="/foo" title="blah">example</a>', tags=TAGS, attributes=ATTRS) ==
+            u'<a title="blah">example</a>'
+        )
+
+    def test_attributes_tag_callable(self):
+        """Verify attributes[tag] callable works"""
+        def img_test(tag, name, val):
+            return name == 'src' and val.startswith('https')
 
         ATTRS = {
             'img': img_test,
@@ -198,6 +221,28 @@ class TestClean:
             u'foo <img src="https://example.com"> baz'
         )
 
+    def test_attributes_tag_list(self):
+        """Verify attributes[tag] list works"""
+        ATTRS = {
+            'a': ['title']
+        }
+        TAGS = ['a']
+
+        assert (
+            bleach.clean(u'<a href="/foo" title="blah">example</a>', tags=TAGS, attributes=ATTRS) ==
+            u'<a title="blah">example</a>'
+        )
+
+    def test_attributes_list(self):
+        """Verify attributes list works"""
+        ATTRS = ['title']
+        TAGS = ['a']
+
+        assert (
+            bleach.clean(u'<a href="/foo" title="blah">example</a>', tags=TAGS, attributes=ATTRS) ==
+            u'<a title="blah">example</a>'
+        )
+
     def test_svg_attr_val_allows_ref(self):
         """Unescape values in svg attrs that allow url references"""
         # Local IRI, so keep it
diff --git a/tests/test_security.py b/tests/test_security.py
index 2aac020..da0fe92 100644
--- a/tests/test_security.py
+++ b/tests/test_security.py
@@ -75,7 +75,7 @@ def test_invalid_href_attr():
 def test_invalid_filter_attr():
     IMG = ['img', ]
     IMG_ATTR = {
-        'img': lambda attr, val: attr == 'src' and val == "http://example.com/"
+        'img': lambda tag, name, val: name == 'src' and val == "http://example.com/"
     }
 
     assert (