From 09b1413748195bfbcd1f3f71fb7e8058ccc659e9 Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Wed, 22 Mar 2023 11:54:20 -0700
Subject: [PATCH] Allows users to use OAuth tokens instead of passwords

---
 .../mail-account-edit-dialog.component.html   |  1 +
 .../mail-account-edit-dialog.component.ts     |  1 +
 src-ui/src/app/data/paperless-mail-account.ts |  2 +
 src/paperless_mail/mail.py                    | 25 ++++-----
 .../migrations/0020_mailaccount_is_token.py   | 19 +++++++
 src/paperless_mail/models.py                  |  2 +
 src/paperless_mail/serialisers.py             |  1 +
 src/paperless_mail/tests/test_mail.py         | 53 +++++++++++++++++++
 8 files changed, 92 insertions(+), 12 deletions(-)
 create mode 100644 src/paperless_mail/migrations/0020_mailaccount_is_token.py

diff --git a/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.html b/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.html
index 8164fca9a..356c0a300 100644
--- a/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.html
+++ b/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.html
@@ -15,6 +15,7 @@
       <div class="col">
         <app-input-text i18n-title title="Username" formControlName="username" [error]="error?.username"></app-input-text>
         <app-input-password i18n-title title="Password" formControlName="password" [error]="error?.password"></app-input-password>
+        <app-input-check i18n-title title="Is Token?" formControlName="is_token" [error]="error?.is_token"></app-input-check>
         <app-input-text i18n-title title="Character Set" formControlName="character_set" [error]="error?.character_set"></app-input-text>
       </div>
     </div>
diff --git a/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.ts b/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.ts
index 1c42ffeb0..558ec0229 100644
--- a/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.ts
+++ b/src-ui/src/app/components/common/edit-dialog/mail-account-edit-dialog/mail-account-edit-dialog.component.ts
@@ -45,6 +45,7 @@ export class MailAccountEditDialogComponent extends EditDialogComponent<Paperles
       imap_security: new FormControl(IMAPSecurity.SSL),
       username: new FormControl(null),
       password: new FormControl(null),
+      is_token: new FormControl(false),
       character_set: new FormControl('UTF-8'),
     })
   }
diff --git a/src-ui/src/app/data/paperless-mail-account.ts b/src-ui/src/app/data/paperless-mail-account.ts
index 9f875e783..484997213 100644
--- a/src-ui/src/app/data/paperless-mail-account.ts
+++ b/src-ui/src/app/data/paperless-mail-account.ts
@@ -20,4 +20,6 @@ export interface PaperlessMailAccount extends ObjectWithId {
   password: string
 
   character_set?: string
+
+  is_token: boolean
 }
diff --git a/src/paperless_mail/mail.py b/src/paperless_mail/mail.py
index 5d06db733..50a578563 100644
--- a/src/paperless_mail/mail.py
+++ b/src/paperless_mail/mail.py
@@ -202,20 +202,21 @@ def mailbox_login(mailbox: MailBox, account: MailAccount):
 
     try:
 
-        mailbox.login(account.username, account.password)
+        if account.is_token:
+            mailbox.xoauth2(account.username, account.password)
+        else:
+            try:
+                _ = account.password.encode("ascii")
+                use_ascii_login = True
+            except UnicodeEncodeError:
+                use_ascii_login = False
 
-    except UnicodeEncodeError:
-        logger.debug("Falling back to AUTH=PLAIN")
+            if use_ascii_login:
+                mailbox.login(account.username, account.password)
+            else:
+                logger.debug("Falling back to AUTH=PLAIN")
+                mailbox.login_utf8(account.username, account.password)
 
-        try:
-            mailbox.login_utf8(account.username, account.password)
-        except Exception as e:
-            logger.error(
-                "Unable to authenticate with mail server using AUTH=PLAIN",
-            )
-            raise MailError(
-                f"Error while authenticating account {account}",
-            ) from e
     except Exception as e:
         logger.error(
             f"Error while authenticating account {account}: {e}",
diff --git a/src/paperless_mail/migrations/0020_mailaccount_is_token.py b/src/paperless_mail/migrations/0020_mailaccount_is_token.py
new file mode 100644
index 000000000..5a2f8c572
--- /dev/null
+++ b/src/paperless_mail/migrations/0020_mailaccount_is_token.py
@@ -0,0 +1,19 @@
+# Generated by Django 4.1.7 on 2023-03-22 17:51
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("paperless_mail", "0019_mailrule_filter_to"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="mailaccount",
+            name="is_token",
+            field=models.BooleanField(
+                default=False, verbose_name="Is token authentication"
+            ),
+        ),
+    ]
diff --git a/src/paperless_mail/models.py b/src/paperless_mail/models.py
index d0289ef31..2b2631664 100644
--- a/src/paperless_mail/models.py
+++ b/src/paperless_mail/models.py
@@ -38,6 +38,8 @@ class MailAccount(document_models.ModelWithOwner):
 
     password = models.CharField(_("password"), max_length=256)
 
+    is_token = models.BooleanField(_("Is token authentication"), default=False)
+
     character_set = models.CharField(
         _("character set"),
         max_length=256,
diff --git a/src/paperless_mail/serialisers.py b/src/paperless_mail/serialisers.py
index 28511fcb8..133227468 100644
--- a/src/paperless_mail/serialisers.py
+++ b/src/paperless_mail/serialisers.py
@@ -34,6 +34,7 @@ class MailAccountSerializer(OwnedObjectSerializer):
             "username",
             "password",
             "character_set",
+            "is_token",
         ]
 
     def update(self, instance, validated_data):
diff --git a/src/paperless_mail/tests/test_mail.py b/src/paperless_mail/tests/test_mail.py
index b1d821824..c0bfccba5 100644
--- a/src/paperless_mail/tests/test_mail.py
+++ b/src/paperless_mail/tests/test_mail.py
@@ -83,6 +83,8 @@ class BogusMailBox(ContextManager):
     ASCII_PASSWORD: str = "secret"
     # Note the non-ascii characters here
     UTF_PASSWORD: str = "w57äöüw4b6huwb6nhu"
+    # A dummy access token
+    ACCESS_TOKEN = "ea7e075cd3acf2c54c48e600398d5d5a"
 
     def __init__(self):
         self.messages: List[MailMessage] = []
@@ -112,6 +114,10 @@ class BogusMailBox(ContextManager):
         if username != self.USERNAME or password != self.UTF_PASSWORD:
             raise MailboxLoginError("BAD", "OK")
 
+    def xoauth2(self, username: str, access_token: str):
+        if username != self.USERNAME or access_token != self.ACCESS_TOKEN:
+            raise MailboxLoginError("BAD", "OK")
+
     def fetch(self, criteria, mark_seen, charset=""):
         msg = self.messages
 
@@ -737,6 +743,14 @@ class TestMail(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
         self.assertEqual(len(self.bogus_mailbox.messages), 3)
 
     def test_error_login(self):
+        """
+        GIVEN:
+            - Account configured with incorrect password
+        WHEN:
+            - Account tried to login
+        THEN:
+            - MailError with correct message raised
+        """
         account = MailAccount.objects.create(
             name="test",
             imap_server="",
@@ -1007,6 +1021,8 @@ class TestMail(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
         """
         GIVEN:
             - Mail account with password containing non-ASCII characters
+        WHEN:
+            - Mail account is handled
         THEN:
             - Should still authenticate to the mail account
         """
@@ -1040,6 +1056,8 @@ class TestMail(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
         GIVEN:
             - Mail account with password containing non-ASCII characters
             - Incorrect password value
+        WHEN:
+            - Mail account is handled
         THEN:
             - Should raise a MailError for the account
         """
@@ -1064,6 +1082,41 @@ class TestMail(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
             account,
         )
 
+    def test_auth_with_valid_token(self):
+        """
+        GIVEN:
+            - Mail account configured with access token
+        WHEN:
+            - Mail account is handled
+        THEN:
+            - Should still authenticate to the mail account
+        """
+        account = MailAccount.objects.create(
+            name="test",
+            imap_server="",
+            username=BogusMailBox.USERNAME,
+            # Note the non-ascii characters here
+            password=BogusMailBox.ACCESS_TOKEN,
+            is_token=True,
+        )
+
+        _ = MailRule.objects.create(
+            name="testrule",
+            account=account,
+            action=MailRule.MailAction.MARK_READ,
+        )
+
+        self.assertEqual(len(self.bogus_mailbox.messages), 3)
+        self.assertEqual(self._queue_consumption_tasks_mock.call_count, 0)
+        self.assertEqual(len(self.bogus_mailbox.fetch("UNSEEN", False)), 2)
+
+        self.mail_account_handler.handle_mail_account(account)
+        self.apply_mail_actions()
+
+        self.assertEqual(self._queue_consumption_tasks_mock.call_count, 2)
+        self.assertEqual(len(self.bogus_mailbox.fetch("UNSEEN", False)), 0)
+        self.assertEqual(len(self.bogus_mailbox.messages), 3)
+
     def assert_queue_consumption_tasks_call_args(self, expected_call_args: List):
         """
         Verifies that queue_consumption_tasks has been called with the expected arguments.