diff --git a/docs/index.rst b/docs/index.rst index 49690a2..ef25d0c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -31,6 +31,30 @@ For example, to only allow ```` tags: >>> nh3.clean("Hello", tags={"b"}) 'Hello' +Adding to (or removing from) the defaults +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``tags``, ``clean_content_tags``, ``url_schemes`` and +``generic_attribute_prefixes`` options each *replace* the ammonia defaults. +To extend the defaults without copying them, the ``add_*`` and ``rm_*`` +companion options can be used instead: + +.. code-block:: pycon + + >>> # Allow a custom tag in addition to the defaults. + >>> nh3.clean("x", add_tags={"my-tag"}) + 'x' + + >>> # Forbid but keep the rest of the defaults. + >>> nh3.clean("x", rm_tags={"b"}) + 'x' + +The same pattern works for ``add_clean_content_tags`` / ``rm_clean_content_tags``, +``add_url_schemes`` / ``rm_url_schemes`` and +``add_generic_attribute_prefixes`` / ``rm_generic_attribute_prefixes``. +When combined with the replacement option, ``add_*`` and ``rm_*`` apply on +top of the supplied set. + API reference ------------- diff --git a/nh3.pyi b/nh3.pyi index 6a5a1f7..31e524a 100644 --- a/nh3.pyi +++ b/nh3.pyi @@ -9,15 +9,23 @@ class Cleaner: def __init__( self, tags: Optional[AbstractSet[str]] = None, + add_tags: Optional[AbstractSet[str]] = None, + rm_tags: Optional[AbstractSet[str]] = None, clean_content_tags: Optional[AbstractSet[str]] = None, + add_clean_content_tags: Optional[AbstractSet[str]] = None, + rm_clean_content_tags: Optional[AbstractSet[str]] = None, attributes: Optional[Mapping[str, AbstractSet[str]]] = None, attribute_filter: Optional[Callable[[str, str, str], Optional[str]]] = None, strip_comments: bool = True, link_rel: Optional[str] = "noopener noreferrer", generic_attribute_prefixes: Optional[AbstractSet[str]] = None, + add_generic_attribute_prefixes: Optional[AbstractSet[str]] = None, + rm_generic_attribute_prefixes: Optional[AbstractSet[str]] = None, tag_attribute_values: Optional[Mapping[str, Mapping[str, AbstractSet[str]]]] = None, set_tag_attribute_values: Optional[Mapping[str, Mapping[str, str]]] = None, url_schemes: Optional[AbstractSet[str]] = None, + add_url_schemes: Optional[AbstractSet[str]] = None, + rm_url_schemes: Optional[AbstractSet[str]] = None, allowed_classes: Optional[Mapping[str, AbstractSet[str]]] = None, filter_style_properties: Optional[AbstractSet[str]] = None, ) -> None: ... @@ -26,15 +34,23 @@ class Cleaner: def clean( html: str, tags: Optional[AbstractSet[str]] = None, + add_tags: Optional[AbstractSet[str]] = None, + rm_tags: Optional[AbstractSet[str]] = None, clean_content_tags: Optional[AbstractSet[str]] = None, + add_clean_content_tags: Optional[AbstractSet[str]] = None, + rm_clean_content_tags: Optional[AbstractSet[str]] = None, attributes: Optional[Mapping[str, AbstractSet[str]]] = None, attribute_filter: Optional[Callable[[str, str, str], Optional[str]]] = None, strip_comments: bool = True, link_rel: Optional[str] = "noopener noreferrer", generic_attribute_prefixes: Optional[AbstractSet[str]] = None, + add_generic_attribute_prefixes: Optional[AbstractSet[str]] = None, + rm_generic_attribute_prefixes: Optional[AbstractSet[str]] = None, tag_attribute_values: Optional[Mapping[str, Mapping[str, AbstractSet[str]]]] = None, set_tag_attribute_values: Optional[Mapping[str, Mapping[str, str]]] = None, url_schemes: Optional[AbstractSet[str]] = None, + add_url_schemes: Optional[AbstractSet[str]] = None, + rm_url_schemes: Optional[AbstractSet[str]] = None, allowed_classes: Optional[Mapping[str, AbstractSet[str]]] = None, filter_style_properties: Optional[AbstractSet[str]] = None, ) -> str: ... diff --git a/src/lib.rs b/src/lib.rs index 6540c80..e413a41 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,15 +8,23 @@ use pyo3::types::{PyString, PyTuple}; struct Config { tags: Option>, + add_tags: Option>, + rm_tags: Option>, clean_content_tags: Option>, + add_clean_content_tags: Option>, + rm_clean_content_tags: Option>, attributes: Option>>, attribute_filter: Option>, strip_comments: bool, link_rel: Option, generic_attribute_prefixes: Option>, + add_generic_attribute_prefixes: Option>, + rm_generic_attribute_prefixes: Option>, tag_attribute_values: Option>>>, set_tag_attribute_values: Option>>, url_schemes: Option>, + add_url_schemes: Option>, + rm_url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, } @@ -25,15 +33,23 @@ impl Default for Config { fn default() -> Self { Self { tags: None, + add_tags: None, + rm_tags: None, clean_content_tags: None, + add_clean_content_tags: None, + rm_clean_content_tags: None, attributes: None, attribute_filter: None, strip_comments: true, link_rel: Some("noopener noreferrer".to_string()), generic_attribute_prefixes: None, + add_generic_attribute_prefixes: None, + rm_generic_attribute_prefixes: None, tag_attribute_values: None, set_tag_attribute_values: None, url_schemes: None, + add_url_schemes: None, + rm_url_schemes: None, allowed_classes: None, filter_style_properties: None, } @@ -52,10 +68,26 @@ struct Inner { /// /// :param tags: Sets the tags that are allowed. /// :type tags: ``set[str]``, optional +/// :param add_tags: Adds tags to the allowed set, on top of the defaults or the +/// value supplied via ``tags``. Useful to extend the default whitelist +/// without copying it. +/// :type add_tags: ``set[str]``, optional +/// :param rm_tags: Removes tags from the allowed set, applied after ``tags`` +/// and ``add_tags``. Useful to forbid a small number of tags while keeping +/// the rest of the defaults. +/// :type rm_tags: ``set[str]``, optional /// :param clean_content_tags: Sets the tags whose contents will be completely removed from the output. -/// Must be disjoint from ``tags`` (or the default allowed set when ``tags`` -/// is omitted); a tag cannot be both kept and have its content stripped. +/// Must be disjoint from the effective ``tags`` set (defaults or the value +/// supplied via ``tags``, after ``add_tags``/``rm_tags``); a tag cannot be +/// both kept and have its content stripped. /// :type clean_content_tags: ``set[str]``, optional +/// :param add_clean_content_tags: Adds tags to the clean-content set, on top of +/// the defaults or the value supplied via ``clean_content_tags``. The same +/// disjointness rule as ``clean_content_tags`` applies. +/// :type add_clean_content_tags: ``set[str]``, optional +/// :param rm_clean_content_tags: Removes tags from the clean-content set, +/// applied after ``clean_content_tags`` and ``add_clean_content_tags``. +/// :type rm_clean_content_tags: ``set[str]``, optional /// :param attributes: Sets the HTML attributes that are allowed on specific tags, /// ``*`` key means the attributes are allowed on any tag. /// :type attributes: ``dict[str, set[str]]``, optional @@ -75,6 +107,14 @@ struct Inner { /// :type link_rel: ``str`` /// :param generic_attribute_prefixes: Sets the prefix of attributes that are allowed on any tag. /// :type generic_attribute_prefixes: ``set[str]``, optional +/// :param add_generic_attribute_prefixes: Adds prefixes to the +/// generic-attribute prefix set, on top of the defaults or the value +/// supplied via ``generic_attribute_prefixes``. +/// :type add_generic_attribute_prefixes: ``set[str]``, optional +/// :param rm_generic_attribute_prefixes: Removes prefixes from the +/// generic-attribute prefix set, applied after +/// ``generic_attribute_prefixes`` and ``add_generic_attribute_prefixes``. +/// :type rm_generic_attribute_prefixes: ``set[str]``, optional /// :param tag_attribute_values: Sets the values of HTML attributes that are allowed on specific tags. /// The value is structured as a map from tag names to a map from attribute names to a set of attribute values. /// If a tag is not itself whitelisted, adding entries to this map will do nothing. @@ -91,6 +131,12 @@ struct Inner { /// :type set_tag_attribute_values: ``dict[str, dict[str, str]]``, optional /// :param url_schemes: Sets the URL schemes permitted on ``href`` and ``src`` attributes. /// :type url_schemes: ``set[str]``, optional +/// :param add_url_schemes: Adds URL schemes to the permitted set, on top of the +/// defaults or the value supplied via ``url_schemes``. +/// :type add_url_schemes: ``set[str]``, optional +/// :param rm_url_schemes: Removes URL schemes from the permitted set, applied +/// after ``url_schemes`` and ``add_url_schemes``. +/// :type rm_url_schemes: ``set[str]``, optional /// :param allowed_classes: Sets the CSS classes that are allowed on specific tags. /// The values is structured as a map from tag names to a set of class names. /// The `class` attribute itself should not be whitelisted if this parameter is used. @@ -134,10 +180,22 @@ impl Cleaner { let tags: HashSet<&str> = tags.iter().map(|s| s.as_str()).collect(); builder.tags(tags); } + if let Some(add_tags) = config.add_tags.as_ref() { + builder.add_tags(add_tags.iter().map(|s| s.as_str())); + } + if let Some(rm_tags) = config.rm_tags.as_ref() { + builder.rm_tags(rm_tags.iter().map(|s| s.as_str())); + } if let Some(tags) = config.clean_content_tags.as_ref() { let tags: HashSet<&str> = tags.iter().map(|s| s.as_str()).collect(); builder.clean_content_tags(tags); } + if let Some(add_tags) = config.add_clean_content_tags.as_ref() { + builder.add_clean_content_tags(add_tags.iter().map(|s| s.as_str())); + } + if let Some(rm_tags) = config.rm_clean_content_tags.as_ref() { + builder.rm_clean_content_tags(rm_tags.iter().map(|s| s.as_str())); + } if let Some(attrs) = config.attributes.as_ref() { let attrs: HashMap<&str, HashSet<&str>> = attrs .iter() @@ -155,6 +213,12 @@ impl Cleaner { let prefixes: HashSet<&str> = prefixes.iter().map(|s| s.as_str()).collect(); builder.generic_attribute_prefixes(prefixes); } + if let Some(prefixes) = config.add_generic_attribute_prefixes.as_ref() { + builder.add_generic_attribute_prefixes(prefixes.iter().map(|s| s.as_str())); + } + if let Some(prefixes) = config.rm_generic_attribute_prefixes.as_ref() { + builder.rm_generic_attribute_prefixes(prefixes.iter().map(|s| s.as_str())); + } if let Some(values) = config.tag_attribute_values.as_ref() { let values: HashMap<&str, HashMap<&str, HashSet<&str>>> = values .iter() @@ -244,6 +308,12 @@ impl Cleaner { let url_schemes: HashSet<_> = url_schemes.iter().map(|s| s.as_str()).collect(); builder.url_schemes(url_schemes); } + if let Some(url_schemes) = config.add_url_schemes.as_ref() { + builder.add_url_schemes(url_schemes.iter().map(|s| s.as_str())); + } + if let Some(url_schemes) = config.rm_url_schemes.as_ref() { + builder.rm_url_schemes(url_schemes.iter().map(|s| s.as_str())); + } if let Some(allowed_classes) = config.allowed_classes.as_ref() { builder.allowed_classes( allowed_classes @@ -277,30 +347,47 @@ impl Cleaner { #[new] #[pyo3(signature = ( tags = None, + add_tags = None, + rm_tags = None, clean_content_tags = None, + add_clean_content_tags = None, + rm_clean_content_tags = None, attributes = None, attribute_filter = None, strip_comments = true, link_rel = "noopener noreferrer", generic_attribute_prefixes = None, + add_generic_attribute_prefixes = None, + rm_generic_attribute_prefixes = None, tag_attribute_values = None, set_tag_attribute_values = None, url_schemes = None, + add_url_schemes = None, + rm_url_schemes = None, allowed_classes = None, filter_style_properties = None ))] + #[allow(clippy::too_many_arguments)] fn py_new( py: Python, tags: Option>, + add_tags: Option>, + rm_tags: Option>, clean_content_tags: Option>, + add_clean_content_tags: Option>, + rm_clean_content_tags: Option>, attributes: Option>>, attribute_filter: Option>, strip_comments: bool, link_rel: Option<&str>, generic_attribute_prefixes: Option>, + add_generic_attribute_prefixes: Option>, + rm_generic_attribute_prefixes: Option>, tag_attribute_values: Option>>>, set_tag_attribute_values: Option>>, url_schemes: Option>, + add_url_schemes: Option>, + rm_url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, ) -> PyResult { @@ -322,22 +409,45 @@ impl Cleaner { } } } - if let Some(ref clean_tags) = clean_content_tags { - // A tag listed in both the allowed `tags` set and `clean_content_tags` - // makes ammonia panic. Raise an explicit ValueError instead. When the - // caller omits `tags`, ammonia falls back to its default allowed set, - // so check against that default in order to catch e.g. - // `clean_content_tags={"p"}`. - let conflict = match tags.as_ref() { - Some(allowed) => clean_tags.iter().find(|t| allowed.contains(t.as_str())), - None => { - let default_tags = ammonia::Builder::default().clone_tags(); - clean_tags - .iter() - .find(|t| default_tags.contains(t.as_str())) - } + // Compute the effective `tags` and `clean_content_tags` sets after + // applying the add_*/rm_* modifiers, then check they are disjoint. + // Otherwise ammonia would panic at clean-time. When the caller omits + // the replacement option, ammonia falls back to its own defaults. + if clean_content_tags.is_some() || add_clean_content_tags.is_some() { + let default_builder = ammonia::Builder::default(); + let effective_allowed: HashSet = { + let base: HashSet = match tags.as_ref() { + Some(t) => t.clone(), + None => default_builder + .clone_tags() + .into_iter() + .map(|s| s.to_string()) + .collect(), + }; + let added = add_tags.as_ref(); + let removed = rm_tags.as_ref(); + base.into_iter() + .chain(added.into_iter().flatten().cloned()) + .filter(|t| !removed.is_some_and(|r| r.contains(t))) + .collect() + }; + let effective_clean: HashSet = { + let base: HashSet = match clean_content_tags.as_ref() { + Some(t) => t.clone(), + None => default_builder + .clone_clean_content_tags() + .into_iter() + .map(|s| s.to_string()) + .collect(), + }; + let added = add_clean_content_tags.as_ref(); + let removed = rm_clean_content_tags.as_ref(); + base.into_iter() + .chain(added.into_iter().flatten().cloned()) + .filter(|t| !removed.is_some_and(|r| r.contains(t))) + .collect() }; - if let Some(tag) = conflict { + if let Some(tag) = effective_clean.intersection(&effective_allowed).next() { return Err(PyValueError::new_err(format!( "tag \"{}\" cannot appear in both `tags` and `clean_content_tags`; \ either remove it from `clean_content_tags` or pass an explicit \ @@ -348,15 +458,23 @@ impl Cleaner { } let config = Config { tags, + add_tags, + rm_tags, clean_content_tags, + add_clean_content_tags, + rm_clean_content_tags, attributes, attribute_filter, strip_comments, link_rel: link_rel.map(|s| s.to_string()), generic_attribute_prefixes, + add_generic_attribute_prefixes, + rm_generic_attribute_prefixes, tag_attribute_values, set_tag_attribute_values, url_schemes, + add_url_schemes, + rm_url_schemes, allowed_classes, filter_style_properties, }; @@ -496,15 +614,23 @@ impl Cleaner { #[pyfunction(signature = ( html, tags = None, + add_tags = None, + rm_tags = None, clean_content_tags = None, + add_clean_content_tags = None, + rm_clean_content_tags = None, attributes = None, attribute_filter = None, strip_comments = true, link_rel = "noopener noreferrer", generic_attribute_prefixes = None, + add_generic_attribute_prefixes = None, + rm_generic_attribute_prefixes = None, tag_attribute_values = None, set_tag_attribute_values = None, url_schemes = None, + add_url_schemes = None, + rm_url_schemes = None, allowed_classes = None, filter_style_properties = None ))] @@ -513,30 +639,46 @@ fn clean( py: Python, html: &str, tags: Option>, + add_tags: Option>, + rm_tags: Option>, clean_content_tags: Option>, + add_clean_content_tags: Option>, + rm_clean_content_tags: Option>, attributes: Option>>, attribute_filter: Option>, strip_comments: bool, link_rel: Option<&str>, generic_attribute_prefixes: Option>, + add_generic_attribute_prefixes: Option>, + rm_generic_attribute_prefixes: Option>, tag_attribute_values: Option>>>, set_tag_attribute_values: Option>>, url_schemes: Option>, + add_url_schemes: Option>, + rm_url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, ) -> PyResult { let cleaner = Cleaner::py_new( py, tags, + add_tags, + rm_tags, clean_content_tags, + add_clean_content_tags, + rm_clean_content_tags, attributes, attribute_filter, strip_comments, link_rel, generic_attribute_prefixes, + add_generic_attribute_prefixes, + rm_generic_attribute_prefixes, tag_attribute_values, set_tag_attribute_values, url_schemes, + add_url_schemes, + rm_url_schemes, allowed_classes, filter_style_properties, )?; diff --git a/tests/test_nh3.py b/tests/test_nh3.py index 0526d57..e9ff08b 100644 --- a/tests/test_nh3.py +++ b/tests/test_nh3.py @@ -67,6 +67,94 @@ def test_clean(): ) +def test_add_tags_extends_defaults(): + # "my-tag" is not in the defaults, so it would normally be stripped. + assert nh3.clean("x") == "x" + # add_tags extends the defaults without replacing them, so default tags + # like are still preserved. + assert ( + nh3.clean("x", add_tags={"my-tag"}) + == "x" + ) + + +def test_rm_tags_removes_from_defaults(): + # is allowed by default. + assert nh3.clean("x") == "x" + # rm_tags strips it while keeping the rest of the defaults. + assert nh3.clean("x", rm_tags={"b"}) == "x" + + +def test_add_tags_combined_with_explicit_tags(): + # Explicit `tags=` replaces the defaults; add_tags is applied on top of it. + assert ( + nh3.clean("xy", tags={"b"}, add_tags={"i"}) + == "xy" + ) + + +def test_add_and_rm_clean_content_tags(): + # add_clean_content_tags wipes the content of an extra tag. + assert ( + nh3.clean("secret", add_clean_content_tags={"my-tag"}) + == "" + ) + # rm_clean_content_tags lets a default clean-content tag through (its + # contents survive even though the tag itself is still stripped). + assert "alert" in nh3.clean( + "", rm_clean_content_tags={"script"} + ) + + +def test_add_and_rm_url_schemes(): + # add_url_schemes permits a custom scheme on a link. + assert ( + 'href="myapp:foo"' + in nh3.clean('x', add_url_schemes={"myapp"}) + ) + # rm_url_schemes strips a default-allowed scheme. + assert "href" not in nh3.clean( + 'x', rm_url_schemes={"https"} + ) + + +def test_add_and_rm_generic_attribute_prefixes(): + # add_generic_attribute_prefixes allows a custom prefix on any tag. + assert 'foo-bar="v"' in nh3.clean( + "

x

", add_generic_attribute_prefixes={"foo-"} + ) + # rm_generic_attribute_prefixes removes a prefix that was first added via + # generic_attribute_prefixes. + assert "data-x" not in nh3.clean( + "

x

", + generic_attribute_prefixes={"data-"}, + rm_generic_attribute_prefixes={"data-"}, + ) + + +def test_add_clean_content_tags_overlap_with_add_tags(): + # If a tag ends up in both effective sets via add_*, validation must fire. + with pytest.raises(ValueError, match="clean_content_tags"): + nh3.clean( + "x", + add_tags={"my-tag"}, + add_clean_content_tags={"my-tag"}, + ) + + +def test_rm_clean_content_tags_resolves_overlap(): + # `clean_content_tags={"b"}` would conflict with the default tag, but + # rm_tags removes from the allowed set first, so this is valid. + assert ( + nh3.clean( + "secretsafe", + rm_tags={"b"}, + add_clean_content_tags={"b"}, + ) + == "safe" + ) + + def test_clean_with_attribute_filter(): html = "Home"