diff --git a/hobbit_core/flask_hobbit/utils.py b/hobbit_core/flask_hobbit/utils.py index e92c977..de4c352 100644 --- a/hobbit_core/flask_hobbit/utils.py +++ b/hobbit_core/flask_hobbit/utils.py @@ -1,5 +1,8 @@ # -*- encoding: utf-8 -*- +import os +import re import six +from unicodedata import normalize class dict2object(dict): @@ -7,7 +10,7 @@ class dict2object(dict): 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 diff --git a/tests/__init__.py b/tests/__init__.py index c77e923..f81a274 100644 --- a/tests/__init__.py +++ b/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') diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..7048cac --- /dev/null +++ b/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]