Some refactoring, refresh from test

src/documents/views.py

@@ -8,14 +8,12 @@ import tempfile
 import urllib
 import zipfile
 from datetime import datetime
-from datetime import timedelta
 from pathlib import Path
 from time import mktime
 from unicodedata import normalize
 from urllib.parse import quote
 from urllib.parse import urlparse
 
-import httpx
 import pathvalidate
 from django.apps import apps
 from django.conf import settings
@@ -2158,87 +2156,3 @@ class TrashView(ListModelMixin, PassUserMixin):
                 doc_ids = [doc.id for doc in docs]
             empty_trash(doc_ids=doc_ids)
         return Response({"result": "OK", "doc_ids": doc_ids})
-
-
-# Outlook https://stackoverflow.com/questions/73902642/office-365-imap-authentication-via-oauth2-and-python-msal-library
-class OauthCallbackView(GenericAPIView):
-    # permission_classes = (AllowAny,)
-
-    def get(self, request, format=None):
-        code = request.query_params.get("code")
-        # Gmail passes scope as a query param, Outlook does not
-        scope = request.query_params.get("scope")
-
-        if code is None:
-            logger.error(
-                f"Invalid oauth callback request, code: {code}, scope: {scope}",
-            )
-            return HttpResponseBadRequest("Invalid request, see logs for more detail")
-
-        if scope is not None and "google" in scope:
-            # Google
-            # Gmail setup guide: https://postmansmtp.com/how-to-configure-post-smtp-with-gmailgsuite-using-oauth/
-            imap_server = "imap.gmail.com"
-            defaults = {
-                "name": f"Gmail OAuth {datetime.now()}",
-                "username": "",
-                "imap_security": MailAccount.ImapSecurity.SSL,
-                "imap_port": 993,
-            }
-
-            token_request_uri = "https://accounts.google.com/o/oauth2/token"
-            client_id = settings.GMAIL_OAUTH_CLIENT_ID
-            client_secret = settings.GMAIL_OAUTH_CLIENT_SECRET
-            scope = "https://mail.google.com/"
-        elif scope is None:
-            # Outlook
-            # Outlok setup guide: https://medium.com/@manojkumardhakad/python-read-and-send-outlook-mail-using-oauth2-token-and-graph-api-53de606ecfa1
-            imap_server = "outlook.office365.com"
-            defaults = {
-                "name": f"Outlook OAuth {datetime.now()}",
-                "username": "",
-                "imap_security": MailAccount.ImapSecurity.SSL,
-                "imap_port": 993,
-            }
-
-            token_request_uri = (
-                "https://login.microsoftonline.com/common/oauth2/v2.0/token"
-            )
-            client_id = settings.OUTLOOK_OAUTH_CLIENT_ID
-            client_secret = settings.OUTLOOK_OAUTH_CLIENT_SECRET
-            scope = "offline_access https://outlook.office.com/IMAP.AccessAsUser.All"
-
-        data = {
-            "code": code,
-            "client_id": client_id,
-            "client_secret": client_secret,
-            "scope": scope,
-            "redirect_uri": "http://localhost:8000/api/oauth/callback/",
-            "grant_type": "authorization_code",
-        }
-        headers = {
-            "Content-Type": "application/x-www-form-urlencoded",
-        }
-        response = httpx.post(token_request_uri, data=data, headers=headers)
-        data = response.json()
-
-        if "error" in data:
-            logger.error(f"Error {response.status_code} getting access token: {data}")
-            return HttpResponseRedirect(
-                "http://localhost:4200/mail?oauth_success=0",
-            )
-        elif "access_token" in data:
-            access_token = data["access_token"]
-            refresh_token = data["refresh_token"]
-            expires_in = data["expires_in"]
-            account, _ = MailAccount.objects.update_or_create(
-                password=access_token,
-                is_token=True,
-                imap_server=imap_server,
-                refresh_token=refresh_token,
-                expiration=timezone.now() + timedelta(seconds=expires_in),
-                defaults=defaults,
-            )
-            return HttpResponseRedirect(
-                f"http://localhost:4200/mail?oauth_success=1&account_id={account.pk}",
-            )

src/paperless/urls.py

@@ -24,7 +24,6 @@ from documents.views import DocumentTypeViewSet
 from documents.views import GlobalSearchView
 from documents.views import IndexView
 from documents.views import LogViewSet
-from documents.views import OauthCallbackView
 from documents.views import PostDocumentView
 from documents.views import RemoteVersionView
 from documents.views import SavedViewViewSet
@@ -55,6 +54,7 @@ from paperless.views import UserViewSet
 from paperless_mail.views import MailAccountTestView
 from paperless_mail.views import MailAccountViewSet
 from paperless_mail.views import MailRuleViewSet
+from paperless_mail.views import OauthCallbackView
 
 api_router = DefaultRouter()
 api_router.register(r"correspondents", CorrespondentViewSet)

src/paperless_mail/mail.py

@@ -422,6 +422,53 @@ def get_mailbox(server, port, security) -> MailBox:
     return mailbox
 
 
+def refresh_oauth_token(self, account: MailAccount) -> bool:
+    """
+    Refreshes the token for the given mail account.
+    """
+    self.log.debug(f"Attempting to refresh oauth token for account {account}")
+    if not account.refresh_token:
+        self.log.error(f"Account {account}: No refresh token available.")
+        return False
+
+    if "gmail" in account.imap_server:
+        url = "https://accounts.google.com/o/oauth2/token"
+        data = {
+            "client_id": settings.GMAIL_OAUTH_CLIENT_ID,
+            "client_secret": settings.GMAIL_OAUTH_CLIENT_SECRET,
+            "refresh_token": account.refresh_token,
+            "grant_type": "refresh_token",
+        }
+    elif "outlook" in account.imap_server:
+        url = "https://login.microsoftonline.com/common/oauth2/v2.0/token"
+        data = {
+            "client_id": settings.OUTLOOK_OAUTH_CLIENT_ID,
+            "client_secret": settings.OUTLOOK_OAUTH_CLIENT_SECRET,
+            "refresh_token": account.refresh_token,
+            "grant_type": "refresh_token",
+        }
+
+    response = httpx.post(
+        url=url,
+        data=data,
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+    )
+    data = response.json()
+    if response.status_code < 400 and "access_token" in data:
+        account.password = data["access_token"]
+        account.expiration = timezone.now() + timedelta(
+            seconds=data["expires_in"],
+        )
+        account.save()
+        self.log.debug(f"Successfully refreshed oauth token for account {account}")
+        return True
+    else:
+        self.log.error(
+            f"Failed to refresh oauth token for account {account}: {response}",
+        )
+        return False
+
+
 class MailAccountHandler(LoggingMixin):
     """
     The main class that handles mail accounts.
@@ -516,49 +563,6 @@ class MailAccountHandler(LoggingMixin):
                 "Unknown correspondent selector",
             )  # pragma: no cover
 
-    def refresh_token(self, account: MailAccount) -> bool:
-        """
-        Refreshes the token for the given mail account.
-        """
-        if not account.refresh_token:
-            self.log.error(f"Account {account}: No refresh token available.")
-            return False
-
-        if "gmail" in account.imap_server:
-            url = "https://accounts.google.com/o/oauth2/token"
-            data = {
-                "client_id": settings.GMAIL_OAUTH_CLIENT_ID,
-                "client_secret": settings.GMAIL_OAUTH_CLIENT_SECRET,
-                "refresh_token": account.refresh_token,
-                "grant_type": "refresh_token",
-            }
-        elif "outlook" in account.imap_server:
-            url = "https://login.microsoftonline.com/common/oauth2/v2.0/token"
-            data = {
-                "client_id": settings.OUTLOOK_OAUTH_CLIENT_ID,
-                "client_secret": settings.OUTLOOK_OAUTH_CLIENT_SECRET,
-                "refresh_token": account.refresh_token,
-                "grant_type": "refresh_token",
-            }
-
-        response = httpx.post(
-            url=url,
-            data=data,
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-        )
-        data = response.json()
-        if "access_token" in data:
-            account.password = data["access_token"]
-            account.expiration = timezone.now() + timedelta(
-                seconds=data["expires_in"],
-            )
-            account.save()
-            self.log.debug(f"Successfully refreshed token for account {account}")
-            return True
-        else:
-            self.log.error(f"Failed to refresh token for account {account}: {data}")
-            return False
-
     def handle_mail_account(self, account: MailAccount):
         """
         Main entry method to handle a specific mail account.
@@ -580,9 +584,7 @@ class MailAccountHandler(LoggingMixin):
                     and account.expiration is not None
                     and account.expiration < timezone.now()
                 ):
-                    self.log.debug(f"Attempting to refresh token for account {account}")
+                    if refresh_oauth_token(account):
-                    success = self.refresh_token(account)
-                    if success:
                         account.refresh_from_db()
                     else:
                         return total_processed_files

src/paperless_mail/views.py

@@ -1,7 +1,12 @@
 import datetime
 import logging
+from datetime import timedelta
 
+import httpx
+from django.conf import settings
 from django.http import HttpResponseBadRequest
+from django.http import HttpResponseRedirect
+from django.utils import timezone
 from rest_framework.generics import GenericAPIView
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
@@ -14,6 +19,7 @@ from paperless.views import StandardPagination
 from paperless_mail.mail import MailError
 from paperless_mail.mail import get_mailbox
 from paperless_mail.mail import mailbox_login
+from paperless_mail.mail import refresh_oauth_token
 from paperless_mail.models import MailAccount
 from paperless_mail.models import MailRule
 from paperless_mail.serialisers import MailAccountSerializer
@@ -50,14 +56,15 @@ class MailAccountTestView(GenericAPIView):
         serializer = self.get_serializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
-        # account exists, use the password from there instead of ***
+        # account exists, use the password from there instead of *** and refresh_token / expiration
         if (
             len(serializer.validated_data.get("password").replace("*", "")) == 0
             and request.data["id"] is not None
         ):
-            serializer.validated_data["password"] = MailAccount.objects.get(
+            existing_account = MailAccount.objects.get(pk=request.data["id"])
-                pk=request.data["id"],
+            serializer.validated_data["password"] = existing_account.password
-            ).password
+            serializer.validated_data["refresh_token"] = existing_account.refresh_token
+            serializer.validated_data["expiration"] = existing_account.expiration
 
         account = MailAccount(**serializer.validated_data)
 
@@ -67,6 +74,16 @@ class MailAccountTestView(GenericAPIView):
             account.imap_security,
         ) as M:
             try:
+                if (
+                    account.is_token
+                    and account.expiration is not None
+                    and account.expiration < timezone.now()
+                ):
+                    if refresh_oauth_token(account):
+                        account.refresh_from_db()
+                    else:
+                        raise MailError("Unable to refresh oauth token")
+
                 mailbox_login(M, account)
                 return Response({"success": True})
             except MailError:
@@ -74,3 +91,85 @@ class MailAccountTestView(GenericAPIView):
                     f"Mail account {account} test failed",
                 )
                 return HttpResponseBadRequest("Unable to connect to server")
+
+
+class OauthCallbackView(GenericAPIView):
+    def get(self, request, format=None):
+        logger = logging.getLogger("paperless_mail")
+        code = request.query_params.get("code")
+        # Gmail passes scope as a query param, Outlook does not
+        scope = request.query_params.get("scope")
+
+        if code is None:
+            logger.error(
+                f"Invalid oauth callback request, code: {code}, scope: {scope}",
+            )
+            return HttpResponseBadRequest("Invalid request, see logs for more detail")
+
+        if scope is not None and "google" in scope:
+            # Google
+            # Gmail setup guide: https://postmansmtp.com/how-to-configure-post-smtp-with-gmailgsuite-using-oauth/
+            imap_server = "imap.gmail.com"
+            defaults = {
+                "name": f"Gmail OAuth {datetime.now()}",
+                "username": "",
+                "imap_security": MailAccount.ImapSecurity.SSL,
+                "imap_port": 993,
+            }
+
+            token_request_uri = "https://accounts.google.com/o/oauth2/token"
+            client_id = settings.GMAIL_OAUTH_CLIENT_ID
+            client_secret = settings.GMAIL_OAUTH_CLIENT_SECRET
+            scope = "https://mail.google.com/"
+        elif scope is None:
+            # Outlook
+            # Outlok setup guide: https://medium.com/@manojkumardhakad/python-read-and-send-outlook-mail-using-oauth2-token-and-graph-api-53de606ecfa1
+            imap_server = "outlook.office365.com"
+            defaults = {
+                "name": f"Outlook OAuth {datetime.now()}",
+                "username": "",
+                "imap_security": MailAccount.ImapSecurity.SSL,
+                "imap_port": 993,
+            }
+
+            token_request_uri = (
+                "https://login.microsoftonline.com/common/oauth2/v2.0/token"
+            )
+            client_id = settings.OUTLOOK_OAUTH_CLIENT_ID
+            client_secret = settings.OUTLOOK_OAUTH_CLIENT_SECRET
+            scope = "offline_access https://outlook.office.com/IMAP.AccessAsUser.All"
+
+        data = {
+            "code": code,
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "scope": scope,
+            "redirect_uri": "http://localhost:8000/api/oauth/callback/",
+            "grant_type": "authorization_code",
+        }
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+        response = httpx.post(token_request_uri, data=data, headers=headers)
+        data = response.json()
+
+        if "error" in data:
+            logger.error(f"Error {response.status_code} getting access token: {data}")
+            return HttpResponseRedirect(
+                "http://localhost:4200/mail?oauth_success=0",
+            )
+        elif "access_token" in data:
+            access_token = data["access_token"]
+            refresh_token = data["refresh_token"]
+            expires_in = data["expires_in"]
+            account, _ = MailAccount.objects.update_or_create(
+                password=access_token,
+                is_token=True,
+                imap_server=imap_server,
+                refresh_token=refresh_token,
+                expiration=timezone.now() + timedelta(seconds=expires_in),
+                defaults=defaults,
+            )
+            return HttpResponseRedirect(
+                f"http://localhost:4200/mail?oauth_success=1&account_id={account.pk}",
+            )