Chore: use PinnedIMAP4

src/paperless_mail/mail.py

@@ -1,6 +1,8 @@
 import datetime
+import imaplib
 import itertools
 import logging
+import socket
 import ssl
 import tempfile
 import traceback
@@ -33,6 +35,7 @@ from imap_tools import MailMessageFlags
 from imap_tools import errors
 from imap_tools.mailbox import MailBoxStartTls
 from imap_tools.query import LogicOperator
+from imap_tools.utils import check_command_status
 
 from documents.data_models import ConsumableDocument
 from documents.data_models import DocumentMetadataOverrides
@@ -429,12 +432,95 @@ def make_criterias(rule: MailRule, *, supports_gmail_labels: bool):
         return AND(rule_query, **criterias)
 
 
+class PinnedIMAP4(imaplib.IMAP4):
+    """
+    IMAP4 client which connects to addresses that have already been resolved.
+    ``self.host`` keeps the original hostname for TLS SNI and cert verification
+
+    Without pinned addresses, and with the ssl_context of the matching imaplib
+    class, this behaves exactly like imaplib.IMAP4 / imaplib.IMAP4_SSL.
+    """
+
+    def __init__(self, host, port, pinned_ips, ssl_context=None, timeout=None) -> None:
+        self._pinned_ips = pinned_ips
+        self.ssl_context = ssl_context
+        super().__init__(host, port, timeout=timeout)
+
+    def _connect_pinned(self, timeout):
+        last_error: OSError | None = None
+        for ip_str in self._pinned_ips:
+            try:
+                address = (ip_str, self.port)
+                if timeout is not None:
+                    return socket.create_connection(address, timeout)
+                return socket.create_connection(address)
+            except OSError as e:
+                last_error = e
+        raise last_error or OSError(f"Could not connect to {self.host}")
+
+    def _create_socket(self, timeout):
+        if self._pinned_ips:
+            sock = self._connect_pinned(timeout)
+        else:
+            sock = super()._create_socket(timeout)
+        if self.ssl_context is None:
+            return sock
+        return self.ssl_context.wrap_socket(sock, server_hostname=self.host)
+
+
+class PinnedClientMixin:
+    """Builds the imaplib client against the pre-resolved addresses, if any."""
+
+    def __init__(self, *args, pinned_ips: list[str] | None, **kwargs) -> None:
+        self._pinned_ips = pinned_ips
+        super().__init__(*args, **kwargs)
+
+    def _pinned_client(self, ssl_context=None) -> imaplib.IMAP4:
+        return PinnedIMAP4(
+            self._host,
+            self._port,
+            self._pinned_ips,
+            ssl_context=ssl_context,
+            timeout=self._timeout,
+        )
+
+
+class PinnedMailBox(PinnedClientMixin, MailBox):
+    def _get_mailbox_client(self) -> imaplib.IMAP4:
+        return self._pinned_client(self._ssl_context)
+
+
+class PinnedMailBoxUnencrypted(PinnedClientMixin, MailBoxUnencrypted):
+    def _get_mailbox_client(self) -> imaplib.IMAP4:
+        return self._pinned_client()
+
+
+class PinnedMailBoxStartTls(PinnedClientMixin, MailBoxStartTls):
+    def _get_mailbox_client(self) -> imaplib.IMAP4:
+        if self._port == 993:
+            raise ValueError(
+                "Port 993 requires IMAP4_SSL. Use MailBox class for SSL/TLS connection.",
+            )
+        client = self._pinned_client()
+        check_command_status(
+            client.starttls(self._ssl_context),
+            errors.MailboxStarttlsError,
+        )
+        return client
+
+
 def get_mailbox(server, port, security) -> MailBox:
     """
     Returns the correct MailBox instance for the given configuration.
     """
+    pinned_ips: list[str] | None = None
     if not settings.EMAIL_ALLOW_INTERNAL_HOSTS:
-        for ip_str in resolve_hostname_ips(server):
+        try:
+            pinned_ips = resolve_hostname_ips(server)
+        except ValueError as e:
+            raise MailError(str(e)) from e
+
+        for ip_str in pinned_ips:
             if not is_public_ip(ip_str):
                 raise MailError(
                     f"Connection blocked: {server} resolves to a non-public address",
@@ -445,11 +531,21 @@ def get_mailbox(server, port, security) -> MailBox:
         ssl_context.load_verify_locations(cafile=settings.EMAIL_CERTIFICATE_FILE)
 
     if security == MailAccount.ImapSecurity.NONE:
-        mailbox = MailBoxUnencrypted(server, port)
+        mailbox = PinnedMailBoxUnencrypted(server, port, pinned_ips=pinned_ips)
     elif security == MailAccount.ImapSecurity.STARTTLS:
-        mailbox = MailBoxStartTls(server, port, ssl_context=ssl_context)
+        mailbox = PinnedMailBoxStartTls(
+            server,
+            port,
+            ssl_context=ssl_context,
+            pinned_ips=pinned_ips,
+        )
     elif security == MailAccount.ImapSecurity.SSL:
-        mailbox = MailBox(server, port, ssl_context=ssl_context)
+        mailbox = PinnedMailBox(
+            server,
+            port,
+            ssl_context=ssl_context,
+            pinned_ips=pinned_ips,
+        )
     else:
         raise NotImplementedError("Unknown IMAP security")  # pragma: no cover
     return mailbox

src/paperless_mail/tests/test_api.py

@@ -26,7 +26,7 @@ class TestAPIMailAccounts(DirectoriesMixin, APITestCase):
     def setUp(self) -> None:
         self.bogus_mailbox = BogusMailBox()
 
-        patcher = mock.patch("paperless_mail.mail.MailBox")
+        patcher = mock.patch("paperless_mail.mail.PinnedMailBox")
         m = patcher.start()
         m.return_value = self.bogus_mailbox
         self.addCleanup(patcher.stop)

src/paperless_mail/tests/test_mail.py

@@ -36,6 +36,7 @@ from paperless_mail.mail import MailAccountHandler
 from paperless_mail.mail import MailError
 from paperless_mail.mail import TagMailAction
 from paperless_mail.mail import apply_mail_action
+from paperless_mail.mail import get_mailbox
 from paperless_mail.models import MailAccount
 from paperless_mail.models import MailRule
 from paperless_mail.models import ProcessedMail
@@ -341,7 +342,7 @@ class MailMocker(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
 
         reset_bogus_mailbox(self.bogus_mailbox, self.messageBuilder)
 
-        patcher = mock.patch("paperless_mail.mail.MailBox")
+        patcher = mock.patch("paperless_mail.mail.PinnedMailBox")
         m = patcher.start()
         m.return_value = self.bogus_mailbox
         self.addCleanup(patcher.stop)
@@ -441,6 +442,7 @@ class TestMail(
     def setUp(self) -> None:
         self.mailMocker = MailMocker()
         self.mailMocker.setUp()
+        self.addCleanup(self.mailMocker.doCleanups)
         self.mail_account_handler = MailAccountHandler()
 
         super().setUp()
@@ -2190,6 +2192,7 @@ class TestMailAccountTestView(APITestCase):
     def setUp(self) -> None:
         self.mailMocker = MailMocker()
         self.mailMocker.setUp()
+        self.addCleanup(self.mailMocker.doCleanups)
         self.user = User.objects.create_user(
             username="testuser",
             password="testpassword",
@@ -2386,10 +2389,82 @@ class TestMailAccountTestView(APITestCase):
         self.assertEqual(response.content.decode(), "Insufficient permissions")
 
 
+class TestGetMailboxHostPinning(TestCase):
+    """
+    get_mailbox() must connect to the address it validated, so that a DNS answer
+    which changes between the check and the connection (DNS rebinding) cannot
+    redirect the connection to an internal host.
+    """
+
+    @override_settings(EMAIL_ALLOW_INTERNAL_HOSTS=False)
+    @mock.patch(
+        "paperless_mail.mail.resolve_hostname_ips",
+        return_value=["93.184.216.34"],
+    )
+    def test_connects_to_validated_ip(self, _mock_resolve) -> None:
+        with mock.patch(
+            "paperless_mail.mail.socket.create_connection",
+            side_effect=OSError("no connection in tests"),
+        ) as mock_connection:
+            with self.assertRaises(OSError):
+                get_mailbox(
+                    "mail.example.com",
+                    143,
+                    MailAccount.ImapSecurity.NONE,
+                )
+
+        # the hostname is never handed to the socket layer for a second lookup
+        mock_connection.assert_called_once_with(("93.184.216.34", 143))
+
+    @override_settings(EMAIL_ALLOW_INTERNAL_HOSTS=False)
+    @mock.patch(
+        "paperless_mail.mail.resolve_hostname_ips",
+        return_value=["93.184.216.34"],
+    )
+    def test_ssl_pins_ip_but_keeps_hostname_for_sni(self, _mock_resolve) -> None:
+        ssl_context = mock.MagicMock()
+        ssl_context.wrap_socket.return_value.makefile.side_effect = OSError(
+            "no connection in tests",
+        )
+
+        with (
+            mock.patch(
+                "paperless_mail.mail.ssl.create_default_context",
+                return_value=ssl_context,
+            ),
+            mock.patch(
+                "paperless_mail.mail.socket.create_connection",
+            ) as mock_connection,
+        ):
+            with self.assertRaises(OSError):
+                get_mailbox(
+                    "mail.example.com",
+                    993,
+                    MailAccount.ImapSecurity.SSL,
+                )
+
+        mock_connection.assert_called_once_with(("93.184.216.34", 993))
+        # certificate verification still happens against the hostname
+        ssl_context.wrap_socket.assert_called_once_with(
+            mock_connection.return_value,
+            server_hostname="mail.example.com",
+        )
+
+    @override_settings(EMAIL_ALLOW_INTERNAL_HOSTS=False)
+    @mock.patch(
+        "paperless_mail.mail.resolve_hostname_ips",
+        return_value=["93.184.216.34", "127.0.0.1"],
+    )
+    def test_blocks_when_any_resolved_address_is_internal(self, _mock_resolve) -> None:
+        with self.assertRaises(MailError):
+            get_mailbox("mail.example.com", 993, MailAccount.ImapSecurity.SSL)
+
+
 class TestMailAccountProcess(APITestCase):
     def setUp(self) -> None:
         self.mailMocker = MailMocker()
         self.mailMocker.setUp()
+        self.addCleanup(self.mailMocker.doCleanups)
         self.user = User.objects.create_superuser(
             username="testuser",
             password="testpassword",