OAUTH: Add support to get an IMAP OAUTH token

Latest mutt can do this for MS and GMail providers, provide support for
getting the right scope and some examples how to set it up.

Signed-off-by: Jason Gunthorpe <jgg@mellanox.com>

1 files changed, 40 insertions, 10 deletions

diff --git a/cloud_mdir_sync/cms_oauth_main.py b/cloud_mdir_sync/cms_oauth_main.py
index c7c6699..514e411 100644
--- a/cloud_mdir_sync/cms_oauth_main.py
+++ b/cloud_mdir_sync/cms_oauth_main.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: GPL-2.0+
 import argparse
 import base64
+import contextlib
 import re
 import socket
 
@@ -9,6 +10,11 @@ def get_xoauth2_token(args):
     """Return the xoauth2 string. This is something like
         'user=foo^Aauth=Bearer bar^A^A'
     """
+    if args.test_smtp:
+        args.proto = "SMTP"
+    elif args.test_imap:
+        args.proto = "IMAP"
+
     with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
         sock.connect(args.cms_sock)
         sock.sendall(f"{args.proto} {args.user}".encode())
@@ -20,14 +26,27 @@ def get_xoauth2_token(args):
 
 
 def test_smtp(args, xoauth2_token):
-    """Initiate a testing SMTP connection to verify   """
+    """Initiate a testing SMTP connection to verify the token and server
+    work"""
     import smtplib
-    conn = smtplib.SMTP(args.test_smtp, 587)
-    conn.set_debuglevel(True)
-    conn.ehlo()
-    conn.starttls()
-    conn.ehlo()
-    conn.auth("xoauth2", lambda x: xoauth2_token, initial_response_ok=False)
+    with contextlib.closing(smtplib.SMTP(args.test_smtp, 587)) as conn:
+        conn.set_debuglevel(True)
+        conn.ehlo()
+        conn.starttls()
+        conn.ehlo()
+        conn.auth("xoauth2",
+                  lambda x: xoauth2_token,
+                  initial_response_ok=False)
+
+
+def test_imap(args, xoauth2_token):
+    """Initiate a testing IMAP connection to verify the token and server
+    work"""
+    import imaplib
+    with contextlib.closing(imaplib.IMAP4_SSL(args.test_imap)) as conn:
+        conn.debug = 4
+        conn.authenticate('XOAUTH2', lambda x: xoauth2_token.encode())
+        conn.select('INBOX')
 
 
 def main():
@@ -35,7 +54,7 @@ def main():
     parser.add_argument(
         "--proto",
         default="SMTP",
-        choices={"SMTP"},
+        choices={"SMTP", "IMAP"},
         help="""Select the protocol to get a token for. The protocol will
         automatically select the correct OAUTH scope.""")
     parser.add_argument(
@@ -57,19 +76,30 @@ def main():
         xoauth2 is used if the caller will provide the base64 conversion.
         token returns the bare access_token""")
 
-    parser.add_argument(
+    tests = parser.add_mutually_exclusive_group()
+    tests.add_argument(
         "--test-smtp",
         metavar="SMTP_SERVER",
         help=
         """If specified attempt to connect and authenticate to the given SMTP
-        sever. This can be used to test that the authentication method works
+        server. This can be used to test that the authentication method works
         properly on the server. Typical servers would be smtp.office365.com
         and smtp.gmail.com.""")
+    tests.add_argument(
+        "--test-imap",
+        metavar="IMAP_SERVER",
+        help=
+        """If specified attempt to connect and authenticate to the given IMAP
+        server. This can be used to test that the authentication method works
+        properly on the server. Typical servers would be outlook.office365.com
+        and imap.gmail.com.""")
     args = parser.parse_args()
 
     xoauth2_token = get_xoauth2_token(args)
     if args.test_smtp:
         return test_smtp(args, xoauth2_token)
+    if args.test_imap:
+        return test_imap(args, xoauth2_token)
 
     if args.output == "xoauth2-b64":
         print(base64.b64encode(xoauth2_token.encode()).decode())