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())