Merge pull request #16 from TTWShell/feature/secure_filename

Feature/secure filename

Author: TTWShell

hobbit_core/flask_hobbit/utils.py

@@ -1,13 +1,16 @@
 # -*- encoding: utf-8 -*-
+import os
+import re
 import six
+from unicodedata import normalize
 
 
 class dict2object(dict):
     """Dict to fake object that can use getattr.
 
     Examples::
 
-        In [2]: obj = dict2object({'a':2, 'c':3})
+        In [2]: obj = dict2object({'a': 2, 'c': 3})
 
         In [3]: obj.a
         Out[3]: 2
@@ -26,3 +29,48 @@ def __setattr__(self, name, value):
         if not isinstance(name, six.string_types):
             raise TypeError('key must be string type.')
         self[name] = value
+
+
+def secure_filename(filename):
+    """Borrowed from werkzeug.utils.secure_filename. **Python3 only**.
+
+    Pass it a filename and it will return a secure version of it. This
+    filename can then safely be stored on a regular file system and passed
+    to :func:`os.path.join`.
+
+    On windows systems the function also makes sure that the file is not
+    named after one of the special device files.
+
+        >>> secure_filename("My cool movie.mov")
+        'My_cool_movie.mov'
+        >>> secure_filename("../../../etc/passwd")
+        'etc_passwd'
+        >>> secure_filename(u'i contain cool \xfcml\xe4uts.txt')
+        'i_contain_cool_umlauts.txt'
+    """
+    for sep in os.path.sep, os.path.altsep:
+        if sep:
+            filename = filename.replace(sep, ' ')
+
+    filename = '_'.join(filename.split())
+
+    if isinstance(filename, six.text_type):
+        filename = normalize('NFKD', filename).encode('utf-8')
+        if not six.PY2:
+            filename = filename.decode('utf-8')
+
+    filename_strip_re = re.compile(r'[^A-Za-z0-9\u4e00-\u9fa5_.-]')
+    filename = filename_strip_re.sub('', filename).strip('._')
+
+    # on nt a couple of special files are present in each folder.  We
+    # have to ensure that the target file is not such a filename.  In
+    # this case we prepend an underline
+    windows_device_files = (
+        'CON', 'AUX', 'COM1', 'COM2', 'COM3', 'COM4', 'LPT1',
+        'LPT2', 'LPT3', 'PRN', 'NUL',
+    )
+    if os.name == 'nt' and filename and \
+            filename.split('.')[0].upper() in windows_device_files:
+        filename = '_' + filename
+
+    return filename

tests/__init__.py

@@ -1,7 +1,10 @@
 import os
 import shutil
+import six
 import functools
 
+import pytest
+
 
 class BaseTest(object):
     root_path = os.path.split(os.path.abspath(__name__))[0]
@@ -30,3 +33,6 @@ def inner(*args, **kwargs):
             os.chdir(cwd)
         return inner
     return wrapper
+
+
+python3_only = pytest.mark.skipif(six.PY2, reason='only support Python3')

tests/test_utils.py

@@ -0,0 +1,37 @@
+# -*- encoding: utf-8 -*-
+import pytest
+
+from hobbit_core.flask_hobbit import utils
+
+from . import BaseTest, python3_only
+
+
+class TestUtils(BaseTest):
+
+    def test_dict2object(self):
+        obj = utils.dict2object({'a': 2, 'c': 3})
+        assert obj.a == 2
+        assert obj.c == 3
+
+        # test setattr
+        obj.a = 4
+        assert obj.a == 4
+
+        # test getattr
+        with pytest.raises(AttributeError):
+            print(obj.b)
+
+    @python3_only
+    def test_secure_filename(self):
+        filenames = (
+            '哈哈.zip', '../../../etc/passwd', 'My cool movie.mov',
+            '__filename__', 'foo$&^*)bar',
+            'i contain cool \xfcml\xe4uts.txt',
+        )
+        excepted = (
+            '哈哈.zip', 'etc_passwd', 'My_cool_movie.mov',
+            'filename', 'foobar',
+            'i_contain_cool_umlauts.txt',
+        )
+        for i, filename in enumerate(filenames):
+            assert utils.secure_filename(filename) == excepted[i]