"""What may be uploaded, and what the bytes have to say about it.
Run: DATABASE_URL=sqlite:// PYTHONPATH=backend python -m unittest discover -s backend/tests
"""
import os
os.environ.setdefault("DATABASE_URL", "sqlite://")
import io
import unittest
import zipfile
from app.services import file_intake
def docx(text: str = "Febrile seizures are common between six months and five years.") -> bytes:
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, "w") as bundle:
bundle.writestr("[Content_Types].xml", "")
bundle.writestr(
"word/document.xml",
f"{text}")
return buffer.getvalue()
class Upload:
"""The two attributes the reader touches on a Starlette UploadFile."""
def __init__(self, data: bytes):
self.file = io.BytesIO(data)
PNG = b"\x89PNG\r\n\x1a\n" + b"0" * 40
class SniffingTests(unittest.TestCase):
def test_the_bytes_decide_not_the_extension(self):
self.assertEqual(file_intake.kind_of(b"%PDF-1.7\nstuff"), file_intake.PDF)
self.assertEqual(file_intake.kind_of(docx()), file_intake.DOCX)
self.assertEqual(file_intake.kind_of(PNG), file_intake.IMAGE)
self.assertEqual(file_intake.kind_of(b"RIFF\x00\x00\x00\x00WEBPmore"), file_intake.IMAGE)
def test_everything_else_is_refused(self):
# A shell script called report.pdf is not a PDF; a zip that is not a
# Word document is not one either.
self.assertIsNone(file_intake.kind_of(b"#!/bin/sh\nrm -rf /"))
self.assertIsNone(file_intake.kind_of(b""))
self.assertIsNone(file_intake.kind_of(b"id,name\n1,two\n"))
plain_zip = io.BytesIO()
with zipfile.ZipFile(plain_zip, "w") as bundle:
bundle.writestr("notes.txt", "hello")
self.assertIsNone(file_intake.kind_of(plain_zip.getvalue()))
self.assertIsNone(file_intake.kind_of(b""))
class SizeTests(unittest.TestCase):
def test_over_the_cap_is_refused_by_the_byte_past_it(self):
with self.assertRaises(file_intake.Rejected) as refusal:
file_intake.read(Upload(b"x" * (file_intake.MAX_BYTES + 1)))
self.assertIn("2 MB", str(refusal.exception))
def test_exactly_the_cap_is_allowed_and_empty_is_not(self):
self.assertEqual(len(file_intake.read(Upload(b"x" * file_intake.MAX_BYTES))),
file_intake.MAX_BYTES)
with self.assertRaises(file_intake.Rejected):
file_intake.read(Upload(b""))
class TextTests(unittest.TestCase):
def test_a_word_document_gives_up_its_words(self):
text = file_intake.text_from(docx(), file_intake.DOCX, 4000)
self.assertIn("Febrile seizures are common", text)
# Tags never reach the text, so no markup can be smuggled through one.
self.assertNotIn("", text)
def test_a_zip_bomb_is_refused_before_it_is_read(self):
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as bundle:
bundle.writestr("[Content_Types].xml", "")
bundle.writestr("word/document.xml", "" + ("a" * (file_intake.MAX_UNZIPPED + 10)))
with self.assertRaises(file_intake.Rejected):
file_intake.text_from(buffer.getvalue(), file_intake.DOCX, 4000)
def test_control_characters_go_and_the_words_stay(self):
self.assertEqual(file_intake.clean("a\x00b\x07 c"), "a b c")
def test_an_image_without_a_reader_says_so_rather_than_failing(self):
with self.assertRaises(file_intake.Rejected) as refusal:
file_intake.text_from(PNG, file_intake.IMAGE, 4000, describe=None)
self.assertIn("tool model", str(refusal.exception))
def test_an_image_with_a_reader_is_read(self):
text = file_intake.text_from(
PNG, file_intake.IMAGE, 4000,
describe=lambda data, mime: f"A slide about croup ({mime})")
self.assertEqual(text, "A slide about croup (image/png)")
if __name__ == "__main__":
unittest.main()