test_middleware.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. import os
  2. import pytest
  3. from django.core.exceptions import PermissionDenied
  4. from django.core.files.base import ContentFile
  5. from django.core.files.uploadedfile import SimpleUploadedFile
  6. from s3file.middleware import S3FileMiddleware
  7. from s3file.storages import get_aws_location, storage
  8. class TestS3FileMiddleware:
  9. def test_get_files_from_storage(self, freeze_upload_folder):
  10. content = b"test_get_files_from_storage"
  11. name = storage.save(
  12. "tmp/s3file/test_get_files_from_storage", ContentFile(content)
  13. )
  14. files = S3FileMiddleware.get_files_from_storage(
  15. [os.path.join(get_aws_location(), name)],
  16. "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  17. )
  18. file = next(files)
  19. assert file.read() == content
  20. def test_process_request(self, freeze_upload_folder, rf):
  21. uploaded_file = SimpleUploadedFile("uploaded_file.txt", b"uploaded")
  22. request = rf.post("/", data={"file": uploaded_file})
  23. S3FileMiddleware(lambda x: None)(request)
  24. assert request.FILES.getlist("file")
  25. assert request.FILES.get("file").read() == b"uploaded"
  26. storage.save("tmp/s3file/s3_file.txt", ContentFile(b"s3file"))
  27. request = rf.post(
  28. "/",
  29. data={
  30. "file": "custom/location/tmp/s3file/s3_file.txt",
  31. "s3file": "file",
  32. "file-s3f-signature": "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  33. },
  34. )
  35. S3FileMiddleware(lambda x: None)(request)
  36. assert request.FILES.getlist("file")
  37. assert request.FILES.get("file").read() == b"s3file"
  38. def test_process_request__location_escape(self, freeze_upload_folder, rf):
  39. storage.save("secrets/passwords.txt", ContentFile(b"keep this secret"))
  40. request = rf.post(
  41. "/",
  42. data={
  43. "file": "custom/location/secrets/passwords.txt",
  44. "s3file": "file",
  45. "file-s3f-signature": "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  46. },
  47. )
  48. with pytest.raises(PermissionDenied) as e:
  49. S3FileMiddleware(lambda x: None)(request)
  50. assert "Illegal signature!" in str(e.value)
  51. def test_process_request__multiple_files(self, freeze_upload_folder, rf):
  52. storage.save("tmp/s3file/s3_file.txt", ContentFile(b"s3file"))
  53. storage.save("tmp/s3file/s3_other_file.txt", ContentFile(b"other s3file"))
  54. request = rf.post(
  55. "/",
  56. data={
  57. "file": [
  58. "custom/location/tmp/s3file/s3_file.txt",
  59. "custom/location/tmp/s3file/s3_other_file.txt",
  60. ],
  61. "file-s3f-signature": "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  62. "other_file-s3f-signature": "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  63. "s3file": ["file", "other_file"],
  64. },
  65. )
  66. S3FileMiddleware(lambda x: None)(request)
  67. files = request.FILES.getlist("file")
  68. assert files[0].read() == b"s3file"
  69. assert files[1].read() == b"other s3file"
  70. def test_process_request__no_location(self, freeze_upload_folder, rf, settings):
  71. settings.AWS_LOCATION = ""
  72. uploaded_file = SimpleUploadedFile("uploaded_file.txt", b"uploaded")
  73. request = rf.post("/", data={"file": uploaded_file})
  74. S3FileMiddleware(lambda x: None)(request)
  75. assert request.FILES.getlist("file")
  76. assert request.FILES.get("file").read() == b"uploaded"
  77. storage.save("tmp/s3file/s3_file.txt", ContentFile(b"s3file"))
  78. request = rf.post(
  79. "/",
  80. data={
  81. "file": "tmp/s3file/s3_file.txt",
  82. "s3file": "file",
  83. "file-s3f-signature": "pJYaM4x7RzLDLVXWuphK2dMqqc0oLr_jZFasfGU7BhU",
  84. },
  85. )
  86. S3FileMiddleware(lambda x: None)(request)
  87. assert request.FILES.getlist("file")
  88. assert request.FILES.get("file").read() == b"s3file"
  89. def test_process_request__no_file(self, freeze_upload_folder, rf, caplog):
  90. request = rf.post(
  91. "/",
  92. data={
  93. "file": "custom/location/tmp/s3file/does_not_exist.txt",
  94. "s3file": "file",
  95. "file-s3f-signature": "VRIPlI1LCjUh1EtplrgxQrG8gSAaIwT48mMRlwaCytI",
  96. },
  97. )
  98. S3FileMiddleware(lambda x: None)(request)
  99. assert not request.FILES.getlist("file")
  100. assert (
  101. "File not found: custom/location/tmp/s3file/does_not_exist.txt"
  102. in caplog.text
  103. )
  104. def test_process_request__no_signature(self, rf, caplog):
  105. request = rf.post(
  106. "/", data={"file": "tmp/s3file/does_not_exist.txt", "s3file": "file"}
  107. )
  108. with pytest.raises(PermissionDenied) as e:
  109. S3FileMiddleware(lambda x: None)(request)
  110. assert "No signature provided." in str(e.value)
  111. def test_process_request__wrong_signature(self, rf, caplog):
  112. request = rf.post(
  113. "/",
  114. data={
  115. "file": "tmp/s3file/does_not_exist.txt",
  116. "s3file": "file",
  117. "file-s3f-signature": "fake",
  118. },
  119. )
  120. with pytest.raises(PermissionDenied) as e:
  121. S3FileMiddleware(lambda x: None)(request)
  122. assert "Illegal signature!" in str(e.value)
  123. def test_sign_s3_key_prefix(self, rf):
  124. assert (
  125. S3FileMiddleware.sign_s3_key_prefix("test/test")
  126. == "a8KINhIf1IpSD5sgdXE4wEQodZorq_8CmwkqZ5V6nr4"
  127. )