Add more tests for MailMessageDecryptor

This commit is contained in:
Daniel Bankmann 2024-08-25 11:52:43 +02:00 committed by shamoon
parent 9113efd918
commit acd1726bf2

View File

@ -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",