From 55089aab32f0d8df41f7505ba351f6236e1fa335 Mon Sep 17 00:00:00 2001
From: Trenton H <holmes.trenton@gmail.com>
Date: Fri, 7 Oct 2022 07:57:26 -0700
Subject: [PATCH] Fixes handling of gmail label extension to IMAP

---
 src/paperless_mail/mail.py            | 41 +++++++++++++++++++++------
 src/paperless_mail/tests/test_mail.py | 20 +++++++------
 2 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/src/paperless_mail/mail.py b/src/paperless_mail/mail.py
index ebab59a88..f28586a2a 100644
--- a/src/paperless_mail/mail.py
+++ b/src/paperless_mail/mail.py
@@ -4,6 +4,7 @@ import tempfile
 from datetime import date
 from datetime import timedelta
 from fnmatch import fnmatch
+from typing import Dict
 
 import magic
 import pathvalidate
@@ -30,7 +31,7 @@ class MailError(Exception):
 
 
 class BaseMailAction:
-    def get_criteria(self):
+    def get_criteria(self) -> Dict:
         return {}
 
     def post_consume(self, M, message_uids, parameter):
@@ -78,7 +79,7 @@ class TagMailAction(BaseMailAction):
             M.flag(message_uids, [self.keyword], True)
 
 
-def get_rule_action(rule):
+def get_rule_action(rule) -> BaseMailAction:
     if rule.action == MailRule.MailAction.FLAG:
         return FlagMailAction()
     elif rule.action == MailRule.MailAction.DELETE:
@@ -108,7 +109,7 @@ def make_criterias(rule):
     return {**criterias, **get_rule_action(rule).get_criteria()}
 
 
-def get_mailbox(server, port, security):
+def get_mailbox(server, port, security) -> MailBox:
     if security == MailAccount.ImapSecurity.NONE:
         mailbox = MailBoxUnencrypted(server, port)
     elif security == MailAccount.ImapSecurity.STARTTLS:
@@ -167,7 +168,7 @@ class MailAccountHandler(LoggingMixin):
                 "Unknown correspondent selector",
             )  # pragma: nocover
 
-    def handle_mail_account(self, account):
+    def handle_mail_account(self, account: MailAccount):
 
         self.renew_logging_group()
 
@@ -181,7 +182,14 @@ class MailAccountHandler(LoggingMixin):
                 account.imap_security,
             ) as M:
 
+                supports_gmail_labels = "X-GM-EXT-1" in M.client.capabilities
+                supports_auth_plain = "AUTH=PLAIN" in M.client.capabilities
+
+                self.log("debug", f"GMAIL Label Support: {supports_gmail_labels}")
+                self.log("debug", f"AUTH=PLAIN Support: {supports_auth_plain}")
+
                 try:
+
                     M.login(account.username, account.password)
 
                 except UnicodeEncodeError:
@@ -215,7 +223,11 @@ class MailAccountHandler(LoggingMixin):
 
                 for rule in account.rules.order_by("order"):
                     try:
-                        total_processed_files += self.handle_mail_rule(M, rule)
+                        total_processed_files += self.handle_mail_rule(
+                            M,
+                            rule,
+                            supports_gmail_labels,
+                        )
                     except Exception as e:
                         self.log(
                             "error",
@@ -233,7 +245,12 @@ class MailAccountHandler(LoggingMixin):
 
         return total_processed_files
 
-    def handle_mail_rule(self, M: MailBox, rule):
+    def handle_mail_rule(
+        self,
+        M: MailBox,
+        rule: MailRule,
+        supports_gmail_labels: bool = False,
+    ):
 
         self.log("debug", f"Rule {rule}: Selecting folder {rule.folder}")
 
@@ -261,11 +278,19 @@ class MailAccountHandler(LoggingMixin):
             ) from err
 
         criterias = make_criterias(rule)
-        criterias_imap = AND(**criterias)
+
+        # Deal with the Gmail label extension
         if "gmail_label" in criterias:
+
             gmail_label = criterias["gmail_label"]
             del criterias["gmail_label"]
-            criterias_imap = AND(NOT(gmail_label=gmail_label), **criterias)
+
+            if not supports_gmail_labels:
+                criterias_imap = AND(**criterias)
+            else:
+                criterias_imap = AND(NOT(gmail_label=gmail_label), **criterias)
+        else:
+            criterias_imap = AND(**criterias)
 
         self.log(
             "debug",
diff --git a/src/paperless_mail/tests/test_mail.py b/src/paperless_mail/tests/test_mail.py
index be016a79a..997184fd2 100644
--- a/src/paperless_mail/tests/test_mail.py
+++ b/src/paperless_mail/tests/test_mail.py
@@ -47,15 +47,16 @@ class BogusFolderManager:
 
 
 class BogusClient:
+    def __init__(self, messages):
+        self.messages: List[MailMessage] = messages
+        self.capabilities: List[str] = []
+
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
 
-    def __init__(self, messages):
-        self.messages: List[MailMessage] = messages
-
     def authenticate(self, mechanism, authobject):
         # authobject must be a callable object
         auth_bytes = authobject(None)
@@ -80,12 +81,6 @@ class BogusMailBox(ContextManager):
     # Note the non-ascii characters here
     UTF_PASSWORD: str = "w57äöüw4b6huwb6nhu"
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
     def __init__(self):
         self.messages: List[MailMessage] = []
         self.messages_spam: List[MailMessage] = []
@@ -93,6 +88,12 @@ class BogusMailBox(ContextManager):
         self.client = BogusClient(self.messages)
         self._host = ""
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
     def updateClient(self):
         self.client = BogusClient(self.messages)
 
@@ -648,6 +649,7 @@ class TestMail(DirectoriesMixin, TestCase):
 
     def test_handle_mail_account_tag_gmail(self):
         self.bogus_mailbox._host = "imap.gmail.com"
+        self.bogus_mailbox.client.capabilities = ["X-GM-EXT-1"]
 
         account = MailAccount.objects.create(
             name="test",