From acd1726bf2ae317011865e3c3b1591d8ef28af1a Mon Sep 17 00:00:00 2001 From: Daniel Bankmann <6322723+dbankmann@users.noreply.github.com> Date: Sun, 25 Aug 2024 11:52:43 +0200 Subject: [PATCH] Add more tests for MailMessageDecryptor --- src/paperless_mail/tests/test_preprocessor.py | 99 ++++++++++++++----- 1 file changed, 73 insertions(+), 26 deletions(-) diff --git a/src/paperless_mail/tests/test_preprocessor.py b/src/paperless_mail/tests/test_preprocessor.py index 134a0fff1..17d5b05d2 100644 --- a/src/paperless_mail/tests/test_preprocessor.py +++ b/src/paperless_mail/tests/test_preprocessor.py @@ -1,23 +1,20 @@ +import email +import email.contentmanager import tempfile +from email.message import Message from email.mime.application import MIMEApplication from email.mime.multipart import MIMEMultipart -from typing import TYPE_CHECKING import gnupg -from django.test import TestCase from django.test import override_settings from imap_tools import MailMessage -from paperless_mail.mail import MailAccountHandler from paperless_mail.models import MailAccount from paperless_mail.models import MailRule from paperless_mail.preprocessor import MailMessageDecryptor -from paperless_mail.tests.test_mail import MailMocker +from paperless_mail.tests.test_mail import TestMail from paperless_mail.tests.test_mail import _AttachmentDef -if TYPE_CHECKING: - import email.contentmanager - class MessageEncryptor: def __init__(self): @@ -35,10 +32,24 @@ class MessageEncryptor: ) self.gpg.gen_key(input_data) + @staticmethod + def get_email_body_without_headers(email_message: Message) -> bytes: + """ + Filters some relevant headers from an EmailMessage and returns just the body. + """ + message_copy = email.message_from_bytes(email_message.as_bytes()) + + message_copy._headers = [ + header + for header in message_copy._headers + if header[0].lower() not in ("from", "to", "subject") + ] + return message_copy.as_bytes() + def encrypt(self, message): original_email: email.message.Message = message.obj encrypted_data = self.gpg.encrypt( - original_email.as_bytes(), + self.get_email_body_without_headers(original_email), self._testUser, armor=True, ) @@ -65,44 +76,66 @@ class MessageEncryptor: ) new_email.attach(encrypted_part) - encrypted_message = MailMessage( + encrypted_message: MailMessage = MailMessage( [(f"UID {message.uid}".encode(), new_email.as_bytes())], ) return encrypted_message -class TestPreprocessor(TestCase): +class TestMailMessageGpgDecryptor(TestMail): def setUp(self): - self.mailMocker = MailMocker() - self.mailMocker.setUp() - self.messageEncryptor = MessageEncryptor() with override_settings( EMAIL_GNUPG_HOME=self.messageEncryptor.gpg_home, EMAIL_ENABLE_GPG_DECRYPTOR=True, ): - self.mail_account_handler = MailAccountHandler() + super().setUp() - super().setUp() + def test_preprocessor_is_able_to_run(self): + with override_settings( + EMAIL_GNUPG_HOME=self.messageEncryptor.gpg_home, + EMAIL_ENABLE_GPG_DECRYPTOR=True, + ): + self.assertTrue(MailMessageDecryptor.able_to_run()) + + def test_preprocessor_is_able_to_run2(self): + with override_settings( + EMAIL_GNUPG_HOME=None, + EMAIL_ENABLE_GPG_DECRYPTOR=True, + ): + self.assertTrue(MailMessageDecryptor.able_to_run()) + + def test_is_not_able_to_run_disabled(self): + with override_settings( + EMAIL_ENABLE_GPG_DECRYPTOR=False, + ): + self.assertFalse(MailMessageDecryptor.able_to_run()) + + def test_is_not_able_to_run_bogus_path(self): + with override_settings( + EMAIL_ENABLE_GPG_DECRYPTOR=True, + EMAIL_GNUPG_HOME="_)@# notapath &%#$", + ): + self.assertFalse(MailMessageDecryptor.able_to_run()) + + def test_decrypt_fails(self): + encrypted_message, _ = self.create_encrypted_unencrypted_message_pair() + empty_gpg_home = tempfile.mkdtemp() + with override_settings( + EMAIL_ENABLE_GPG_DECRYPTOR=True, + EMAIL_GNUPG_HOME=empty_gpg_home, + ): + message_decryptor = MailMessageDecryptor() + self.assertRaises(Exception, message_decryptor.run, encrypted_message) def test_decrypt_encrypted_mail(self): """ Creates a mail with attachments. Then encrypts it with a new key. Verifies that this encrypted message can be decrypted with attachments intact. """ - message = self.mailMocker.messageBuilder.create_message( - body="Test message with 2 attachments", - attachments=[ - _AttachmentDef( - filename="f1.pdf", - disposition="inline", - ), - _AttachmentDef(filename="f2.pdf"), - ], - ) + encrypted_message, message = self.create_encrypted_unencrypted_message_pair() headers = message.headers text = message.text - encrypted_message = self.messageEncryptor.encrypt(message) self.assertEqual(len(encrypted_message.attachments), 1) self.assertEqual(encrypted_message.attachments[0].filename, "encrypted.asc") @@ -123,6 +156,20 @@ class TestPreprocessor(TestCase): self.assertEqual(decrypted_message.text, text) self.assertEqual(decrypted_message.uid, message.uid) + def create_encrypted_unencrypted_message_pair(self): + message = self.mailMocker.messageBuilder.create_message( + body="Test message with 2 attachments", + attachments=[ + _AttachmentDef( + filename="f1.pdf", + disposition="inline", + ), + _AttachmentDef(filename="f2.pdf"), + ], + ) + encrypted_message = self.messageEncryptor.encrypt(message) + return encrypted_message, message + def test_handle_encrypted_message(self): message = self.mailMocker.messageBuilder.create_message( subject="the message title",