Some refactoring, refresh from test

This commit is contained in:
shamoon
2024-10-05 09:53:33 -07:00
parent 6a1253fea1
commit c798083eb5
4 changed files with 152 additions and 137 deletions

View File

@@ -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}")
success = self.refresh_token(account)
if success:
if refresh_oauth_token(account):
account.refresh_from_db()
else:
return total_processed_files

View File

@@ -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(
pk=request.data["id"],
).password
existing_account = MailAccount.objects.get(pk=request.data["id"])
serializer.validated_data["password"] = existing_account.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}",
)