From 3386eaee6d79229744dbf914c478b8fdd3743453 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Sat, 5 Oct 2024 01:05:29 -0700 Subject: [PATCH] Basic refresh --- src/documents/views.py | 11 +++-- src/paperless_mail/mail.py | 49 +++++++++++++++++++ ...nt_expiration_mailaccount_refresh_token.py | 34 +++++++++++++ src/paperless_mail/models.py | 19 +++++++ 4 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 src/paperless_mail/migrations/0027_mailaccount_expiration_mailaccount_refresh_token.py diff --git a/src/documents/views.py b/src/documents/views.py index 8986cb811..46785a940 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -8,6 +8,7 @@ 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 @@ -1562,7 +1563,7 @@ class UiSettingsView(GenericAPIView): redirect_uri = "http://localhost:8000/api/oauth/callback/" scope = "https://mail.google.com/" access_type = "offline" - url = f"{token_request_uri}?response_type={response_type}&client_id={client_id}&redirect_uri={redirect_uri}&scope={scope}&access_type={access_type}" + url = f"{token_request_uri}?response_type={response_type}&client_id={client_id}&redirect_uri={redirect_uri}&scope={scope}&access_type={access_type}&prompt=consent" return url def generate_outlook_oauth_url(self) -> str: @@ -2220,7 +2221,6 @@ class OauthCallbackView(GenericAPIView): } response = httpx.post(token_request_uri, data=data, headers=headers) data = response.json() - logger.debug(data) if "error" in data: logger.error(f"Error {response.status_code} getting access token: {data}") @@ -2229,13 +2229,14 @@ class OauthCallbackView(GenericAPIView): ) elif "access_token" in data: access_token = data["access_token"] - # if "refresh_token" in data: - # refresh_token = data["refresh_token"] - # expires_in = data["expires_in"] + 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( diff --git a/src/paperless_mail/mail.py b/src/paperless_mail/mail.py index 77d293ea0..b86414674 100644 --- a/src/paperless_mail/mail.py +++ b/src/paperless_mail/mail.py @@ -11,6 +11,7 @@ from fnmatch import fnmatch from pathlib import Path from typing import TYPE_CHECKING +import httpx import magic import pathvalidate from celery import chord @@ -18,6 +19,7 @@ from celery import shared_task from celery.canvas import Signature from django.conf import settings from django.db import DatabaseError +from django.utils import timezone from django.utils.timezone import is_naive from django.utils.timezone import make_aware from imap_tools import AND @@ -514,6 +516,46 @@ 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: + 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: + 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( + "https://login.microsoftonline.com/common/oauth2/v2.0/token", + data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + data = response.json() + if "access_token" in data: + account.token = data["access_token"] + account.expiration = datetime.datetime.now() + timedelta( + seconds=data["expires_in"], + ) + account.save() + 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. @@ -530,6 +572,13 @@ class MailAccountHandler(LoggingMixin): account.imap_port, account.imap_security, ) as M: + if account.is_token and account.expiration < timezone.now(): + self.log.debug(f"Attempting to refresh token for account {account}") + success = self.refresh_token(account) + if not success: + self.log.error(f"Failed to refresh token for account {account}") + return total_processed_files + supports_gmail_labels = "X-GM-EXT-1" in M.client.capabilities supports_auth_plain = "AUTH=PLAIN" in M.client.capabilities diff --git a/src/paperless_mail/migrations/0027_mailaccount_expiration_mailaccount_refresh_token.py b/src/paperless_mail/migrations/0027_mailaccount_expiration_mailaccount_refresh_token.py new file mode 100644 index 000000000..95e202fd1 --- /dev/null +++ b/src/paperless_mail/migrations/0027_mailaccount_expiration_mailaccount_refresh_token.py @@ -0,0 +1,34 @@ +# Generated by Django 5.1.1 on 2024-10-05 07:42 + +from django.db import migrations +from django.db import models + + +class Migration(migrations.Migration): + dependencies = [ + ("paperless_mail", "0026_mailrule_enabled"), + ] + + operations = [ + migrations.AddField( + model_name="mailaccount", + name="expiration", + field=models.DateTimeField( + blank=True, + help_text="The expiration date of the refresh token. ", + null=True, + verbose_name="expiration", + ), + ), + migrations.AddField( + model_name="mailaccount", + name="refresh_token", + field=models.CharField( + blank=True, + help_text="The refresh token to use for token authentication e.g. with oauth2.", + max_length=2048, + null=True, + verbose_name="refresh token", + ), + ), + ] diff --git a/src/paperless_mail/models.py b/src/paperless_mail/models.py index c23ea48c7..231b01659 100644 --- a/src/paperless_mail/models.py +++ b/src/paperless_mail/models.py @@ -51,6 +51,25 @@ class MailAccount(document_models.ModelWithOwner): ), ) + refresh_token = models.CharField( + _("refresh token"), + max_length=2048, + blank=True, + null=True, + help_text=_( + "The refresh token to use for token authentication e.g. with oauth2.", + ), + ) + + expiration = models.DateTimeField( + _("expiration"), + blank=True, + null=True, + help_text=_( + "The expiration date of the refresh token. ", + ), + ) + def __str__(self): return self.name