Add shared helpers for OAuth2 setup and IMAP config building (#24)

* Add shared helpers for OAuth2 setup and IMAP config building

Introduces acquire_token(), auth_description(), and build_imap_conf() so
each script delegates provider detection, token acquisition, error handling,
and conf dict assembly to a single implementation instead of duplicating it
inline.

* Add comprehensive tests for OAuth2 token acquisition and IMAP config
This commit is contained in:
Nathan Moinvaziri 2026-02-08 06:26:42 -08:00 committed by GitHub
parent dd5251f55d
commit 9dca673e0f
9 changed files with 292 additions and 192 deletions

View File

@ -804,44 +804,14 @@ def main():
args = parser.parse_args()
use_oauth2 = bool(args.src_client_id)
global MAX_WORKERS, BATCH_SIZE
MAX_WORKERS = args.workers
BATCH_SIZE = args.batch
# Acquire OAuth2 token if configured
oauth2_token = None
oauth2_provider = None
if use_oauth2:
oauth2_provider = imap_oauth2.detect_oauth2_provider(args.src_host)
if not oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{args.src_host}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token ({oauth2_provider})...")
oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
oauth2_provider, args.src_client_id, args.src_user, args.src_client_secret
)
if not oauth2_token:
print("Error: Failed to acquire OAuth2 token.")
sys.exit(1)
print("OAuth2 token acquired successfully.\n")
# Use a dict so token updates propagate to worker threads
src_conf = {
"host": args.src_host,
"user": args.src_user,
"password": args.src_pass,
"oauth2_token": oauth2_token,
"oauth2": {
"provider": oauth2_provider,
"client_id": args.src_client_id,
"email": args.src_user,
"client_secret": args.src_client_secret,
}
if use_oauth2
else None,
}
# Build connection config (acquires OAuth2 token if configured)
src_conf = imap_session.build_imap_conf(
args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
)
# Expand path (~/...)
local_path = os.path.expanduser(args.dest_path)
@ -856,7 +826,7 @@ def main():
print("\n--- Configuration Summary ---")
print(f"Source Host : {args.src_host}")
print(f"Source User : {args.src_user}")
print(f"Auth Method : {'OAuth2/' + oauth2_provider + ' (XOAUTH2)' if use_oauth2 else 'Basic (password)'}")
print(f"Auth Method : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
print(f"Destination Path: {local_path}")
if args.gmail_mode:
print("Mode : Gmail Backup (All Mail + Labels + Flags)")

View File

@ -59,7 +59,6 @@ Examples:
import argparse
import os
import sys
import imap_common
import imap_oauth2
@ -205,41 +204,20 @@ def main():
DEST_HOST = args.dest_host
DEST_USER = args.dest_user
src_use_oauth2 = bool(args.src_client_id) and not src_is_local
dest_use_oauth2 = bool(args.dest_client_id) and not dest_is_local
# Acquire OAuth2 tokens if configured
src_oauth2_token = None
src_oauth2_provider = None
if src_use_oauth2:
src_oauth2_provider = imap_oauth2.detect_oauth2_provider(SRC_HOST)
if not src_oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{SRC_HOST}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token for source ({src_oauth2_provider})...")
src_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
src_oauth2_provider, args.src_client_id, SRC_USER, args.src_client_secret
if not src_is_local and args.src_client_id:
src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
)
if not src_oauth2_token:
print("Error: Failed to acquire OAuth2 token for source.")
sys.exit(1)
print("Source OAuth2 token acquired successfully.\n")
dest_oauth2_token = None
dest_oauth2_provider = None
if dest_use_oauth2:
dest_oauth2_provider = imap_oauth2.detect_oauth2_provider(DEST_HOST)
if not dest_oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{DEST_HOST}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token for destination ({dest_oauth2_provider})...")
dest_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
dest_oauth2_provider, args.dest_client_id, DEST_USER, args.dest_client_secret
if not dest_is_local and args.dest_client_id:
dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
)
if not dest_oauth2_token:
print("Error: Failed to acquire OAuth2 token for destination.")
sys.exit(1)
print("Destination OAuth2 token acquired successfully.\n")
print("\n--- Configuration Summary ---")
if src_is_local:
@ -247,18 +225,14 @@ def main():
else:
print(f"Source Host : {args.src_host}")
print(f"Source User : {args.src_user}")
print(
f"Source Auth : {'OAuth2/' + src_oauth2_provider + ' (XOAUTH2)' if src_use_oauth2 else 'Basic (password)'}"
)
print(f"Source Auth : {imap_oauth2.auth_description(src_oauth2_provider)}")
if dest_is_local:
print(f"Destination (Local): {args.dest_path}")
else:
print(f"Destination Host: {args.dest_host}")
print(f"Destination User: {args.dest_user}")
print(
f"Destination Auth: {'OAuth2/' + dest_oauth2_provider + ' (XOAUTH2)' if dest_use_oauth2 else 'Basic (password)'}"
)
print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
print("-----------------------------\n")
src = None

View File

@ -218,29 +218,18 @@ def main(argv: Optional[list[str]] = None) -> None:
USERNAME = args.user
PASSWORD = args.password
use_oauth2 = bool(args.client_id)
# Acquire OAuth2 token if configured
oauth2_token = None
oauth2_provider = None
if use_oauth2:
oauth2_provider = imap_oauth2.detect_oauth2_provider(IMAP_SERVER)
if not oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{IMAP_SERVER}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token ({oauth2_provider})...")
oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
oauth2_provider, args.client_id, USERNAME, args.client_secret
if args.client_id:
oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
IMAP_SERVER, args.client_id, USERNAME, args.client_secret
)
if not oauth2_token:
print("Error: Failed to acquire OAuth2 token.")
sys.exit(1)
print("OAuth2 token acquired successfully.\n")
print("\n--- Configuration Summary ---")
print(f"Host : {IMAP_SERVER}")
print(f"User : {USERNAME}")
print(f"Auth Method : {'OAuth2/' + oauth2_provider + ' (XOAUTH2)' if use_oauth2 else 'Basic (password)'}")
print(f"Auth Method : {imap_oauth2.auth_description(oauth2_provider)}")
print("-----------------------------\n")
count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)

View File

@ -7,6 +7,7 @@ Dispatches to provider-specific modules (oauth2_microsoft, oauth2_google).
Provider is auto-detected from the IMAP host string.
"""
import sys
import threading
import oauth2_google
@ -113,6 +114,59 @@ def acquire_oauth2_token_for_provider(provider, client_id, email, client_secret=
return None
def acquire_token(host, client_id, email, client_secret=None, label=None):
"""
Detect the OAuth2 provider from the host and acquire a token.
Prints status messages and calls sys.exit(1) on failure.
Args:
host: IMAP host string (used to detect provider)
client_id: OAuth2 client ID
email: User's email address
client_secret: OAuth2 client secret (required for Google)
label: Optional context label for messages (e.g. "source", "destination")
Returns:
(token, provider) tuple on success.
"""
provider = detect_oauth2_provider(host)
if not provider:
print(f"Error: Could not detect OAuth2 provider from host '{host}'.")
sys.exit(1)
if label:
print(f"Acquiring OAuth2 token for {label} ({provider})...")
else:
print(f"Acquiring OAuth2 token ({provider})...")
token = acquire_oauth2_token_for_provider(provider, client_id, email, client_secret)
if not token:
if label:
print(f"Error: Failed to acquire OAuth2 token for {label}.")
else:
print("Error: Failed to acquire OAuth2 token.")
sys.exit(1)
if label:
print(f"{label.capitalize()} OAuth2 token acquired successfully.\n")
else:
print("OAuth2 token acquired successfully.\n")
return token, provider
def auth_description(provider):
"""
Return a human-readable auth description for config summaries.
Args:
provider: OAuth2 provider string (e.g. "microsoft", "google") or None for password auth.
Returns:
"OAuth2/{provider} (XOAUTH2)" or "Basic (password)".
"""
if provider:
return f"OAuth2/{provider} (XOAUTH2)"
return "Basic (password)"
def refresh_oauth2_token(conf, old_token):
"""
Thread-safe OAuth2 token refresh using double-checked locking.

View File

@ -9,6 +9,46 @@ import imap_common
import imap_oauth2
def build_imap_conf(host, user, password, client_id=None, client_secret=None, label=None):
"""
Build a standard IMAP connection config dict.
If client_id is provided, acquires an OAuth2 token (with error handling and
sys.exit(1) on failure). Otherwise, builds a password-auth config.
Args:
host: IMAP host
user: IMAP username / email
password: IMAP password (used for password auth or as fallback)
client_id: OAuth2 client ID (triggers OAuth2 flow if provided)
client_secret: OAuth2 client secret (required for Google)
label: Optional context label for status messages (e.g. "source", "destination")
Returns:
Dict with keys: host, user, password, oauth2_token, oauth2
"""
oauth2_token = None
oauth2_provider = None
oauth2_info = None
if client_id:
oauth2_token, oauth2_provider = imap_oauth2.acquire_token(host, client_id, user, client_secret, label)
oauth2_info = {
"provider": oauth2_provider,
"client_id": client_id,
"email": user,
"client_secret": client_secret,
}
return {
"host": host,
"user": user,
"password": password,
"oauth2_token": oauth2_token,
"oauth2": oauth2_info,
}
def ensure_connection(conn, conf):
"""
Refresh OAuth2 token if needed and ensure connection is healthy.

View File

@ -942,53 +942,21 @@ def main():
MAX_WORKERS = args.workers
BATCH_SIZE = args.batch
src_use_oauth2 = bool(args.src_client_id)
dest_use_oauth2 = bool(args.dest_client_id)
# Acquire OAuth2 tokens if configured
src_oauth2_token = None
src_oauth2_provider = None
if src_use_oauth2:
src_oauth2_provider = imap_oauth2.detect_oauth2_provider(SRC_HOST)
if not src_oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{SRC_HOST}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token for source ({src_oauth2_provider})...")
src_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
src_oauth2_provider, args.src_client_id, SRC_USER, args.src_client_secret
)
if not src_oauth2_token:
print("Error: Failed to acquire OAuth2 token for source.")
sys.exit(1)
print("Source OAuth2 token acquired successfully.\n")
dest_oauth2_token = None
dest_oauth2_provider = None
if dest_use_oauth2:
dest_oauth2_provider = imap_oauth2.detect_oauth2_provider(DEST_HOST)
if not dest_oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{DEST_HOST}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token for destination ({dest_oauth2_provider})...")
dest_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
dest_oauth2_provider, args.dest_client_id, DEST_USER, args.dest_client_secret
)
if not dest_oauth2_token:
print("Error: Failed to acquire OAuth2 token for destination.")
sys.exit(1)
print("Destination OAuth2 token acquired successfully.\n")
# Build connection configs (acquires OAuth2 tokens if configured)
src_conf = imap_session.build_imap_conf(
SRC_HOST, SRC_USER, SRC_PASS, args.src_client_id, args.src_client_secret, "source"
)
dest_conf = imap_session.build_imap_conf(
DEST_HOST, DEST_USER, DEST_PASS, args.dest_client_id, args.dest_client_secret, "destination"
)
print("\n--- Configuration Summary ---")
print(f"Source Host : {SRC_HOST}")
print(f"Source User : {SRC_USER}")
print(
f"Source Auth : {'OAuth2/' + src_oauth2_provider + ' (XOAUTH2)' if src_use_oauth2 else 'Basic (password)'}"
)
print(f"Source Auth : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
print(f"Destination Host: {DEST_HOST}")
print(f"Destination User: {DEST_USER}")
print(
f"Destination Auth: {'OAuth2/' + dest_oauth2_provider + ' (XOAUTH2)' if dest_use_oauth2 else 'Basic (password)'}"
)
print(f"Destination Auth: {imap_oauth2.auth_description(dest_conf['oauth2'] and dest_conf['oauth2']['provider'])}")
print(f"Delete fm Source: {DELETE_SOURCE}")
print(f"Dest Delete : {DEST_DELETE}")
print(f"Preserve Flags : {preserve_flags}")
@ -997,36 +965,6 @@ def main():
print(f"Target Folder : {TARGET_FOLDER}")
print("-----------------------------\n")
# Use dicts so token updates propagate to worker threads
src_conf = {
"host": SRC_HOST,
"user": SRC_USER,
"password": SRC_PASS,
"oauth2_token": src_oauth2_token,
"oauth2": {
"provider": src_oauth2_provider,
"client_id": args.src_client_id,
"email": SRC_USER,
"client_secret": args.src_client_secret,
}
if src_use_oauth2
else None,
}
dest_conf = {
"host": DEST_HOST,
"user": DEST_USER,
"password": DEST_PASS,
"oauth2_token": dest_oauth2_token,
"oauth2": {
"provider": dest_oauth2_provider,
"client_id": args.dest_client_id,
"email": DEST_USER,
"client_secret": args.dest_client_secret,
}
if dest_use_oauth2
else None,
}
progress_cache_file = None
progress_cache_data = None
progress_cache_lock = None

View File

@ -1046,44 +1046,14 @@ def main():
args = parser.parse_args()
dest_use_oauth2 = bool(args.dest_client_id)
global MAX_WORKERS, BATCH_SIZE
MAX_WORKERS = args.workers
BATCH_SIZE = args.batch
# Acquire OAuth2 token if configured
dest_oauth2_token = None
dest_oauth2_provider = None
if dest_use_oauth2:
dest_oauth2_provider = imap_oauth2.detect_oauth2_provider(args.dest_host)
if not dest_oauth2_provider:
print(f"Error: Could not detect OAuth2 provider from host '{args.dest_host}'.")
sys.exit(1)
print(f"Acquiring OAuth2 token for destination ({dest_oauth2_provider})...")
dest_oauth2_token = imap_oauth2.acquire_oauth2_token_for_provider(
dest_oauth2_provider, args.dest_client_id, args.dest_user, args.dest_client_secret
)
if not dest_oauth2_token:
print("Error: Failed to acquire OAuth2 token for destination.")
sys.exit(1)
print("Destination OAuth2 token acquired successfully.\n")
# Use a dict so token updates propagate to worker threads
dest_conf = {
"host": args.dest_host,
"user": args.dest_user,
"password": args.dest_pass,
"oauth2_token": dest_oauth2_token,
"oauth2": {
"provider": dest_oauth2_provider,
"client_id": args.dest_client_id,
"email": args.dest_user,
"client_secret": args.dest_client_secret,
}
if dest_use_oauth2
else None,
}
# Build connection config (acquires OAuth2 token if configured)
dest_conf = imap_session.build_imap_conf(
args.dest_host, args.dest_user, args.dest_pass, args.dest_client_id, args.dest_client_secret, "destination"
)
# Expand path
local_path = os.path.expanduser(args.src_path)
@ -1124,9 +1094,7 @@ def main():
print(f"Source Path : {local_path}")
print(f"Destination Host: {args.dest_host}")
print(f"Destination User: {args.dest_user}")
print(
f"Destination Auth: {'OAuth2/' + dest_oauth2_provider + ' (XOAUTH2)' if dest_use_oauth2 else 'Basic (password)'}"
)
print(f"Destination Auth: {imap_oauth2.auth_description(dest_conf['oauth2'] and dest_conf['oauth2']['provider'])}")
print(f"Workers : {args.workers}")
if args.gmail_mode:
print("Mode : Gmail Restore with Labels + Flags")

View File

@ -322,3 +322,97 @@ class TestIsAuthError:
"""Test returns False for folder errors."""
error = Exception("Folder not found: INBOX")
assert imap_oauth2.is_auth_error(error) is False
class TestAcquireToken:
"""Tests for the acquire_token convenience function."""
def test_returns_token_and_provider_for_microsoft(self, capsys):
"""Test successful token acquisition for Microsoft host."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="ms_token"):
token, provider = imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com")
assert token == "ms_token"
assert provider == "microsoft"
out = capsys.readouterr().out
assert "Acquiring OAuth2 token (microsoft)" in out
assert "OAuth2 token acquired successfully" in out
def test_returns_token_and_provider_for_google(self, capsys):
"""Test successful token acquisition for Google host."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="g_token"):
token, provider = imap_oauth2.acquire_token(
"imap.gmail.com", "cid", "user@gmail.com", client_secret="secret"
)
assert token == "g_token"
assert provider == "google"
def test_label_appears_in_messages(self, capsys):
"""Test that label is included in status messages."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="tok"):
imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com", label="source")
out = capsys.readouterr().out
assert "Acquiring OAuth2 token for source (microsoft)" in out
assert "Source OAuth2 token acquired successfully" in out
def test_exits_on_unrecognized_host(self):
"""Test sys.exit(1) when provider cannot be detected from host."""
with pytest.raises(SystemExit) as exc_info:
imap_oauth2.acquire_token("imap.example.com", "cid", "user@example.com")
assert exc_info.value.code == 1
def test_exits_on_token_acquisition_failure(self):
"""Test sys.exit(1) when token acquisition returns None."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
with pytest.raises(SystemExit) as exc_info:
imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com")
assert exc_info.value.code == 1
def test_exit_message_includes_label(self, capsys):
"""Test that failure message includes the label when provided."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
with pytest.raises(SystemExit):
imap_oauth2.acquire_token("outlook.office365.com", "cid", "user@example.com", label="destination")
out = capsys.readouterr().out
assert "Failed to acquire OAuth2 token for destination" in out
def test_exit_message_without_label(self, capsys):
"""Test that failure message works without a label."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
with pytest.raises(SystemExit):
imap_oauth2.acquire_token("imap.gmail.com", "cid", "user@gmail.com", client_secret="s")
out = capsys.readouterr().out
assert "Error: Failed to acquire OAuth2 token." in out
def test_passes_client_secret_to_provider(self):
"""Test that client_secret is forwarded to acquire_oauth2_token_for_provider."""
with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="tok") as mock_acq:
imap_oauth2.acquire_token("imap.gmail.com", "cid", "user@gmail.com", client_secret="my_secret")
mock_acq.assert_called_once_with("google", "cid", "user@gmail.com", "my_secret")
class TestAuthDescription:
"""Tests for the auth_description function."""
def test_microsoft_provider(self):
"""Test returns OAuth2 description for Microsoft provider."""
assert imap_oauth2.auth_description("microsoft") == "OAuth2/microsoft (XOAUTH2)"
def test_google_provider(self):
"""Test returns OAuth2 description for Google provider."""
assert imap_oauth2.auth_description("google") == "OAuth2/google (XOAUTH2)"
def test_none_provider(self):
"""Test returns Basic description when provider is None."""
assert imap_oauth2.auth_description(None) == "Basic (password)"
def test_falsy_provider(self):
"""Test returns Basic description for any falsy value."""
assert imap_oauth2.auth_description("") == "Basic (password)"
assert imap_oauth2.auth_description(0) == "Basic (password)"
assert imap_oauth2.auth_description(False) == "Basic (password)"

View File

@ -17,6 +17,79 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../s
import imap_session
class TestBuildImapConf:
"""Tests for build_imap_conf function."""
def test_password_auth_returns_correct_dict(self):
"""Test builds correct config for password authentication."""
conf = imap_session.build_imap_conf("imap.example.com", "user@example.com", "pass123")
assert conf["host"] == "imap.example.com"
assert conf["user"] == "user@example.com"
assert conf["password"] == "pass123"
assert conf["oauth2_token"] is None
assert conf["oauth2"] is None
def test_oauth2_auth_acquires_token(self):
"""Test acquires token and builds oauth2 config when client_id is provided."""
with patch.object(
imap_session.imap_oauth2, "acquire_token", return_value=("my_token", "microsoft")
) as mock_acquire:
conf = imap_session.build_imap_conf(
"outlook.office365.com",
"user@example.com",
"pass",
client_id="cid",
client_secret="csecret",
label="source",
)
mock_acquire.assert_called_once_with("outlook.office365.com", "cid", "user@example.com", "csecret", "source")
assert conf["host"] == "outlook.office365.com"
assert conf["user"] == "user@example.com"
assert conf["password"] == "pass"
assert conf["oauth2_token"] == "my_token"
assert conf["oauth2"] == {
"provider": "microsoft",
"client_id": "cid",
"email": "user@example.com",
"client_secret": "csecret",
}
def test_no_client_id_skips_oauth2(self):
"""Test that oauth2 is None when client_id is not provided."""
with patch.object(imap_session.imap_oauth2, "acquire_token") as mock_acquire:
conf = imap_session.build_imap_conf("imap.example.com", "user", "pass")
mock_acquire.assert_not_called()
assert conf["oauth2"] is None
assert conf["oauth2_token"] is None
def test_empty_client_id_skips_oauth2(self):
"""Test that empty string client_id is treated as no OAuth2."""
with patch.object(imap_session.imap_oauth2, "acquire_token") as mock_acquire:
conf = imap_session.build_imap_conf("imap.example.com", "user", "pass", client_id="")
mock_acquire.assert_not_called()
assert conf["oauth2"] is None
def test_none_client_secret_passed_through(self):
"""Test that client_secret=None is correctly passed to acquire_token and stored."""
with patch.object(imap_session.imap_oauth2, "acquire_token", return_value=("tok", "microsoft")):
conf = imap_session.build_imap_conf("outlook.office365.com", "user@example.com", "pass", client_id="cid")
assert conf["oauth2"]["client_secret"] is None
def test_label_forwarded_to_acquire_token(self):
"""Test that label is passed through to acquire_token."""
with patch.object(imap_session.imap_oauth2, "acquire_token", return_value=("tok", "google")) as mock_acquire:
imap_session.build_imap_conf(
"imap.gmail.com", "user@gmail.com", "pass", client_id="cid", client_secret="sec", label="destination"
)
mock_acquire.assert_called_once_with("imap.gmail.com", "cid", "user@gmail.com", "sec", "destination")
class TestEnsureConnection:
"""Tests for ensure_connection function."""