1. import os

  2. 

  3. from django.core.exceptions import SuspiciousFileOperation

  4. from django.core.files.base import ContentFile

  5. from django.core.files.storage import FileSystemStorage, Storage

  6. from django.db.models import FileField

  7. from django.test import SimpleTestCase

  8. 

  9. 

  10. class AWSS3Storage(Storage):

  11.     """

  12.     Simulate an AWS S3 storage which uses Unix-like paths and allows any

  13.     characters in file names but where there aren't actual folders but just

  14.     keys.

  15.     """

  16. 

  17.     prefix = "mys3folder/"

  18. 

  19.     def _save(self, name, content):

  20.         """

  21.         This method is important to test that Storage.save() doesn't replace

  22.         '\' with '/' (rather FileSystemStorage.save() does).

  23.         """

  24.         return name

  25. 

  26.     def get_valid_name(self, name):

  27.         return name

  28. 

  29.     def get_available_name(self, name, max_length=None):

  30.         return name

  31. 

  32.     def generate_filename(self, filename):

  33.         """

  34.         This is the method that's important to override when using S3 so that

  35.         os.path() isn't called, which would break S3 keys.

  36.         """

  37.         return self.prefix + self.get_valid_name(filename)

  38. 

  39. 

  40. class GenerateFilenameStorageTests(SimpleTestCase):

  41.     def test_storage_dangerous_paths(self):

  42.         candidates = [

  43.             ("/tmp/..", ".."),

  44.             ("/tmp/.", "."),

  45.             ("", ""),

  46.         ]

  47.         s = FileSystemStorage()

  48.         msg = "Could not derive file name from '%s'"

  49.         for file_name, base_name in candidates:

  50.             with self.subTest(file_name=file_name):

  51.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg % base_name):

  52.                     s.get_available_name(file_name)

  53.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg % base_name):

  54.                     s.generate_filename(file_name)

  55. 

  56.     def test_storage_dangerous_paths_dir_name(self):

  57.         candidates = [

  58.             ("tmp/../path", "tmp/.."),

  59.             ("tmp\\..\\path", "tmp/.."),

  60.             ("/tmp/../path", "/tmp/.."),

  61.             ("\\tmp\\..\\path", "/tmp/.."),

  62.         ]

  63.         s = FileSystemStorage()

  64.         for file_name, path in candidates:

  65.             msg = "Detected path traversal attempt in '%s'" % path

  66.             with self.subTest(file_name=file_name):

  67.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  68.                     s.get_available_name(file_name)

  69.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  70.                     s.generate_filename(file_name)

  71. 

  72.     def test_filefield_dangerous_filename(self):

  73.         candidates = [

  74.             ("..", "some/folder/.."),

  75.             (".", "some/folder/."),

  76.             ("", "some/folder/"),

  77.             ("???", "???"),

  78.             ("$.$.$", "$.$.$"),

  79.         ]

  80.         f = FileField(upload_to="some/folder/")

  81.         for file_name, msg_file_name in candidates:

  82.             msg = f"Could not derive file name from '{msg_file_name}'"

  83.             with self.subTest(file_name=file_name):

  84.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  85.                     f.generate_filename(None, file_name)

  86. 

  87.     def test_filefield_dangerous_filename_dot_segments(self):

  88.         f = FileField(upload_to="some/folder/")

  89.         msg = "Detected path traversal attempt in 'some/folder/../path'"

  90.         with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  91.             f.generate_filename(None, "../path")

  92. 

  93.     def test_filefield_generate_filename_absolute_path(self):

  94.         f = FileField(upload_to="some/folder/")

  95.         candidates = [

  96.             "/tmp/path",

  97.             "/tmp/../path",

  98.         ]

  99.         for file_name in candidates:

  100.             msg = f"Detected path traversal attempt in '{file_name}'"

  101.             with self.subTest(file_name=file_name):

  102.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  103.                     f.generate_filename(None, file_name)

  104. 

  105.     def test_filefield_generate_filename(self):

  106.         f = FileField(upload_to="some/folder/")

  107.         self.assertEqual(

  108.             f.generate_filename(None, "test with space.txt"),

  109.             os.path.normpath("some/folder/test_with_space.txt"),

  110.         )

  111. 

  112.     def test_filefield_generate_filename_with_upload_to(self):

  113.         def upload_to(instance, filename):

  114.             return "some/folder/" + filename

  115. 

  116.         f = FileField(upload_to=upload_to)

  117.         self.assertEqual(

  118.             f.generate_filename(None, "test with space.txt"),

  119.             os.path.normpath("some/folder/test_with_space.txt"),

  120.         )

  121. 

  122.     def test_filefield_generate_filename_upload_to_overrides_dangerous_filename(self):

  123.         def upload_to(instance, filename):

  124.             return "test.txt"

  125. 

  126.         f = FileField(upload_to=upload_to)

  127.         candidates = [

  128.             "/tmp/.",

  129.             "/tmp/..",

  130.             "/tmp/../path",

  131.             "/tmp/path",

  132.             "some/folder/",

  133.             "some/folder/.",

  134.             "some/folder/..",

  135.             "some/folder/???",

  136.             "some/folder/$.$.$",

  137.             "some/../test.txt",

  138.             "",

  139.         ]

  140.         for file_name in candidates:

  141.             with self.subTest(file_name=file_name):

  142.                 self.assertEqual(f.generate_filename(None, file_name), "test.txt")

  143. 

  144.     def test_filefield_generate_filename_upload_to_absolute_path(self):

  145.         def upload_to(instance, filename):

  146.             return "/tmp/" + filename

  147. 

  148.         f = FileField(upload_to=upload_to)

  149.         candidates = [

  150.             "path",

  151.             "../path",

  152.             "???",

  153.             "$.$.$",

  154.         ]

  155.         for file_name in candidates:

  156.             msg = f"Detected path traversal attempt in '/tmp/{file_name}'"

  157.             with self.subTest(file_name=file_name):

  158.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  159.                     f.generate_filename(None, file_name)

  160. 

  161.     def test_filefield_generate_filename_upload_to_dangerous_filename(self):

  162.         def upload_to(instance, filename):

  163.             return "/tmp/" + filename

  164. 

  165.         f = FileField(upload_to=upload_to)

  166.         candidates = ["..", ".", ""]

  167.         for file_name in candidates:

  168.             msg = f"Could not derive file name from '/tmp/{file_name}'"

  169.             with self.subTest(file_name=file_name):

  170.                 with self.assertRaisesMessage(SuspiciousFileOperation, msg):

  171.                     f.generate_filename(None, file_name)

  172. 

  173.     def test_filefield_awss3_storage(self):

  174.         """

  175.         Simulate a FileField with an S3 storage which uses keys rather than

  176.         folders and names. FileField and Storage shouldn't have any os.path()

  177.         calls that break the key.

  178.         """

  179.         storage = AWSS3Storage()

  180.         folder = "not/a/folder/"

  181. 

  182.         f = FileField(upload_to=folder, storage=storage)

  183.         key = "my-file-key\\with odd characters"

  184.         data = ContentFile("test")

  185.         expected_key = AWSS3Storage.prefix + folder + key

  186. 

  187.         # Simulate call to f.save()

  188.         result_key = f.generate_filename(None, key)

  189.         self.assertEqual(result_key, expected_key)

  190. 

  191.         result_key = storage.save(result_key, data)

  192.         self.assertEqual(result_key, expected_key)

  193. 

  194.         # Repeat test with a callable.

  195.         def upload_to(instance, filename):

  196.             # Return a non-normalized path on purpose.

  197.             return folder + filename

  198. 

  199.         f = FileField(upload_to=upload_to, storage=storage)

  200. 

  201.         # Simulate call to f.save()

  202.         result_key = f.generate_filename(None, key)

  203.         self.assertEqual(result_key, expected_key)

  204. 

  205.         result_key = storage.save(result_key, data)

  206.         self.assertEqual(result_key, expected_key)
About Sign Up Leave Feedback Contact