From ec4364c65f9ffe161c0f2dca6bc07ce7b9860c66 Mon Sep 17 00:00:00 2001
From: r1cardohj <34265457+r1cardohj@users.noreply.github.com>
Date: Tue, 12 Dec 2023 22:14:31 +0800
Subject: [PATCH] feat: add `utils.cleanify` function for HTML sanitization
 (#75)

---
 CHANGES.rst              |  7 +++++
 docs/api.rst             |  1 +
 docs/basic.rst           | 29 ++++++++++++++++--
 flask_ckeditor/utils.py  | 20 ++++++++++++-
 requirements/example.txt |  1 +
 requirements/tests.in    |  1 +
 requirements/tests.txt   | 18 ++++++++++--
 setup.py                 |  3 ++
 test_flask_ckeditor.py   | 63 ++++++++++++++++++++++++++++++++++++++++
 tox.ini                  |  1 +
 10 files changed, 138 insertions(+), 6 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 9e7a871..f027993 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,13 @@ Changelog
 
 Release date: -
 
+0.5.2
+-----
+
+Release date: N/A
+
+- Add ``cleanify`` function to ``flask_ckeditor.utils`` for HTML sanitization.
+
 
 0.5.1
 -----
diff --git a/docs/api.rst b/docs/api.rst
index ad4d2c6..669a794 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -30,3 +30,4 @@ Utils
 
 .. autofunction:: get_url
 .. autofunction:: random_filename
+.. autofunction:: cleanify
diff --git a/docs/basic.rst b/docs/basic.rst
index ba4dc2a..4f0fc26 100644
--- a/docs/basic.rst
+++ b/docs/basic.rst
@@ -62,7 +62,7 @@ to True to use built-in resources. You can use ``custom_url`` to load your custo
 CKEditor provides five types of preset (see `comparison table <https://ckeditor.com/cke4/presets-all>`_ for the differences):
 
 - ``basic``
-- ``standard`` （default value）
+- ``standard`` (default value)
 - ``full``
 - ``standard-all`` (only available from CDN)
 - ``full-all`` (only available from CDN)
@@ -100,7 +100,7 @@ It's quite simple, just call ``ckeditor.create()`` in the template:
        <input type="submit">
    </form>
 
-You can use ``value`` parameter to pass preset value (i.e. ``ckeditor.create(value='blah...blah...')``.
+You can use ``value`` parameter to pass preset value (i.e. ``ckeditor.create(value='blah...blah...')``).
 
 Get the Data
 ------------
@@ -119,6 +119,31 @@ from ``request.form`` by passing ``ckeditor`` as key:
 
         return render_template('index.html')
 
+Clean the Data
+--------------
+
+It's recommended to sanitize the HTML input from user before saving it to the database.
+
+The Flask-CKEditor provides a helper function `cleanify`. To use it, install the extra dependencies:
+
+.. code-block:: bash
+
+   $ pip install flask-ckeditor[all]
+
+Then call it for your form data (you could use ``allowed_tags`` to pass a list of custom allowed HTML tags):
+
+.. code-block:: python
+
+    from flask import request, render_template
+    from flask_ckeditor.utils import cleanify
+
+    @app.route('/write')
+    def new_post():
+        if request.method == 'POST':
+            data = cleanify(request.form.get('ckeditor'))  # <--
+
+        return render_template('index.html')
+
 Working with Flask-WTF/WTForms
 -------------------------------
 
diff --git a/flask_ckeditor/utils.py b/flask_ckeditor/utils.py
index 711a00e..aa85c96 100644
--- a/flask_ckeditor/utils.py
+++ b/flask_ckeditor/utils.py
@@ -1,8 +1,13 @@
 import os
 import uuid
-
+import warnings
 from flask import url_for
 
+try:
+    import bleach
+except ImportError:
+    warnings.warn('The "bleach" library is not installed, `cleanify` function will not be available.')
+
 
 def get_url(endpoint_or_url):
     if endpoint_or_url.startswith(('https://', 'http://', '/')):
@@ -15,3 +20,16 @@ def random_filename(old_filename):
     ext = os.path.splitext(old_filename)[1]
     new_filename = uuid.uuid4().hex + ext
     return new_filename
+
+
+def cleanify(text, *, allow_tags=None):
+    """Clean the input from client, this function rely on bleach.
+
+    :parm text: input str
+    :parm allow_tags: if you don't want to use default `allow_tags`,
+        you can provide a Iterable which include html tag string like ['a', 'li',...].
+    """
+    default_allowed_tags = {'a', 'abbr', 'b', 'blockquote', 'code',
+                            'em', 'i', 'li', 'ol', 'pre', 'strong', 'ul',
+                            'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p'}
+    return bleach.clean(text, tags=allow_tags or default_allowed_tags)
diff --git a/requirements/example.txt b/requirements/example.txt
index eab17f8..e66159b 100644
--- a/requirements/example.txt
+++ b/requirements/example.txt
@@ -48,3 +48,4 @@ wtforms==3.1.1
     # via
     #   flask-admin
     #   flask-wtf
+    
\ No newline at end of file
diff --git a/requirements/tests.in b/requirements/tests.in
index da0230e..d747657 100644
--- a/requirements/tests.in
+++ b/requirements/tests.in
@@ -5,3 +5,4 @@ flask-wtf
 flask-admin
 flask-sqlalchemy
 tablib
+bleach
diff --git a/requirements/tests.txt b/requirements/tests.txt
index 721c0dc..9223811 100644
--- a/requirements/tests.txt
+++ b/requirements/tests.txt
@@ -6,14 +6,16 @@
 #
 --index-url https://pypi.tuna.tsinghua.edu.cn/simple
 
+bleach==6.1.0
+    # via -r requirements/tests.in
 blinker==1.7.0
     # via flask
 click==8.1.7
     # via flask
 coverage[toml]==7.3.2
-    # via
-    #   coverage
-    #   pytest-cov
+    # via pytest-cov
+exceptiongroup==1.2.0
+    # via pytest
 flask==3.0.0
     # via
     #   -r requirements/tests.in
@@ -26,6 +28,8 @@ flask-sqlalchemy==3.1.1
     # via -r requirements/tests.in
 flask-wtf==1.2.1
     # via -r requirements/tests.in
+greenlet==3.0.2
+    # via sqlalchemy
 iniconfig==2.0.0
     # via pytest
 itsdangerous==2.1.2
@@ -49,12 +53,20 @@ pytest==7.4.3
     #   pytest-cov
 pytest-cov==4.1.0
     # via -r requirements/tests.in
+six==1.16.0
+    # via bleach
 sqlalchemy==2.0.23
     # via flask-sqlalchemy
 tablib==3.5.0
     # via -r requirements/tests.in
+tomli==2.0.1
+    # via
+    #   coverage
+    #   pytest
 typing-extensions==4.8.0
     # via sqlalchemy
+webencodings==0.5.1
+    # via bleach
 werkzeug==3.0.1
     # via flask
 wtforms==3.1.1
diff --git a/setup.py b/setup.py
index dd8e646..a1607a7 100644
--- a/setup.py
+++ b/setup.py
@@ -33,6 +33,9 @@
     install_requires=[
         'Flask'
     ],
+    extras_require={
+        'all': ['flask-wtf', 'bleach']
+    },
     classifiers=[
         'Environment :: Web Environment',
         'Intended Audience :: Developers',
diff --git a/test_flask_ckeditor.py b/test_flask_ckeditor.py
index 015dd94..4a4424b 100644
--- a/test_flask_ckeditor.py
+++ b/test_flask_ckeditor.py
@@ -9,11 +9,14 @@
 """
 import json
 import unittest
+import sys
+import builtins
 
 from flask import Flask, render_template_string, current_app
 from flask_wtf import FlaskForm, CSRFProtect
 
 from flask_ckeditor import CKEditorField, _CKEditor, CKEditor, upload_success, upload_fail
+from flask_ckeditor.utils import cleanify
 
 
 class CKEditorTestCase(unittest.TestCase):
@@ -287,6 +290,66 @@ def test_upload_fail(self):
             {'uploaded': 0, 'error': {'message': 'new error message'}}
         )
 
+    def test_cleanify_input_js(self):
+        input = 'an <script>evil()</script> example'
+        clean_ouput = cleanify(input)
+        self.assertEqual(clean_ouput,
+                         u'an &lt;script&gt;evil()&lt;/script&gt; example')
+
+    def test_cleanify_by_allow_tags(self):
+        input = '<b> hello <a> this is a url </a> !</b> <h1> this is h1 </h1>'
+        clean_out = cleanify(input, allow_tags=['b'])
+        self.assertEqual(clean_out,
+                         '<b> hello &lt;a&gt; this is a url &lt;/a&gt; !</b> &lt;h1&gt; this is h1 &lt;/h1&gt;')
+
+    def test_cleanify_by_default_allow_tags(self):
+        self.maxDiff = None
+        input = """<a>xxxxx</a>
+                <abbr>xxxxx</abbr>
+                <b>xxxxxxx</b>
+                <blockquote>xxxxxxx</blockquote>
+                <code>print(hello)</code>
+                <em>xxxxx</em>
+                <i>xxxxxx</i>
+                <li>xxxxxx</li>
+                <ol>xxxxxx</ol>
+                <pre>xxxxxx</pre>
+                <strong>xxxxxx</strong>
+                <ul>xxxxxx</ul>
+                <h1>xxxxxxx</h1>
+                <h2>xxxxxxx</h2>
+                <h3>xxxxxxx</h3>
+                <h4>xxxxxxx</h4>
+                <h5>xxxxxxx</h5>
+                <p>xxxxxxxx</p>
+        """
+        clean_out = cleanify(input)
+        self.assertEqual(clean_out, input)
+
+    def test_import_cleanify_without_install_bleach(self):
+        origin_import = builtins.__import__
+        origin_modules = sys.modules.copy()
+
+        def import_hook(name, *args, **kwargs):
+            if name == 'bleach':
+                raise ImportError('test case module')
+            else:
+                return origin_import(name, *args, **kwargs)
+
+        if 'flask_ckeditor.utils' in sys.modules:
+            del sys.modules['flask_ckeditor.utils']
+        builtins.__import__ = import_hook
+
+        with self.assertWarns(UserWarning) as w:
+            from flask_ckeditor.utils import cleanify  # noqa: F401
+
+        self.assertEqual(str(w.warning),
+                         'The "bleach" library is not installed, `cleanify` function will not be available.')
+
+        # recover default
+        builtins.__import__ = origin_import
+        sys.modules = origin_modules
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tox.ini b/tox.ini
index 642b196..95abc78 100644
--- a/tox.ini
+++ b/tox.ini
@@ -10,6 +10,7 @@ deps =
     pytest
     coverage
     flask_wtf
+    bleach
 
 [testenv:coverage]
 commands =