Initial plan

* Initial plan

* Add comprehensive unit tests for ensure_folder_exists and append_email

Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com>

README.md

@@ -19,8 +19,9 @@ This repository contains a set of Python scripts designed to migrate emails betw
 1. **`migrate_imap_emails.py`** (The Solution)
    - Migrates emails folder-by-folder.
    - **Multi-threaded**: Uses a thread pool to copy messages in parallel for high speed.
-   - **Smart De-duplication**: Checks if a message already exists in the destination (matching Message-ID and Size) and skips it if found.
-   - **Robust**: Preserves read/unread status (flags) and original dates.
+  - **Smart De-duplication**: Checks if a message already exists in the destination (matching Message-ID) and skips it if found.
+  - **Robust**: Preserves original dates and can preserve IMAP flags with `--preserve-flags`.
+  - **Gmail Mode**: For Gmail -> Gmail migrations, use `--gmail-mode` to migrate `[Gmail]/All Mail` (no duplicates) and apply Gmail labels by copying messages into label folders.
    - **Cleanup**: Optionally deletes messages from the source after successful transfer (effectively a "Move" operation).
      - *Improved for Gmail*: Automatically detects "Trash" folders to ensure emails are properly binned rather than just archived.
    - **Sync Mode**: Optionally deletes emails from destination that no longer exist in source (`--dest-delete`).
@@ -29,10 +30,12 @@ This repository contains a set of Python scripts designed to migrate emails betw
 2. **`compare_imap_folders.py`** (The Validator)
    - Connects to both Source and Destination accounts.
    - Prints a side-by-side comparison table of message counts for every folder.
-   - Essential for verifying that the migration was successful and that counts match.
+  - Supports comparing IMAP to a local backup folder (`.eml` files) as either the source or destination.
+  - Essential for verifying that the migration was successful and that counts match.
 
 3. **`count_imap_emails.py`** (The Investigator)
-   - A simple utility to connect to a single account and count emails in all folders. Useful for initial assessment.
+  - Counts emails in all folders for a single IMAP account (initial assessment / sizing).
+  - Also supports counting a local backup folder (`.eml` files) via `--path` (or `BACKUP_LOCAL_PATH`).
 
 4. **`backup_imap_emails.py`** (The Backup)
    - Downloads emails from an IMAP account to a local disk.
@@ -46,17 +49,20 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - Uploads emails from a local backup to an IMAP server.
    - **Format**: Reads `.eml` files and uploads them preserving original dates.
    - **Structure**: Recreates the folder hierarchy on the destination server.
-   - **Incremental**: Skips emails that already exist (based on Message-ID and size), but still syncs labels and flags.
+  - **Incremental**: Skips emails that already exist (based on Message-ID), but still syncs labels and flags.
    - **Sync Mode**: Optionally deletes emails from destination that no longer exist in local backup (`--dest-delete`).
    - **Gmail Labels Restoration**: Applies labels from `labels_manifest.json` to recreate the original Gmail label structure.
 
 ## Getting Started
 
 ### 1. Prerequisites
-- **Python 3.6+**
-- **No external installations required.**
-  The scripts use only the Python Standard Library, which is installed automatically with Python. You do **not** need to install anything else (no `pip install`).
-  *Used libraries: `imaplib`, `email`, `concurrent.futures`, `re`, `os`, `sys`, `threading`.*
+- **Python 3.9+**
+- **No external installations required for basic (password) authentication.**
+  The scripts use only the Python Standard Library for standard IMAP login.
+
+- **Optional: OAuth2 authentication** requires one additional package depending on your provider:
+  - **Microsoft (Outlook/Office 365):** `pip install msal`
+  - **Google (Gmail):** `pip install google-auth-oauthlib`
 
 ### 2. Installation
 
@@ -105,9 +111,17 @@ You can configure the scripts using **Environment Variables** (recommended for s
    export DEST_IMAP_USERNAME="dest@domain.com"
    export DEST_IMAP_PASSWORD="dest-app-password"
 
+   # OAuth2 (Optional - instead of password)
+   export SRC_OAUTH2_CLIENT_ID="your-client-id"
+   export SRC_OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
+   export DEST_OAUTH2_CLIENT_ID="your-dest-client-id"
+   export DEST_OAUTH2_CLIENT_SECRET="your-dest-client-secret"  # Required for Google
+
    # Options (Optional)
    export DELETE_FROM_SOURCE="false"  # Set to "true" to delete from source after copy
    export DEST_DELETE="false"         # Set to "true" to delete orphans from destination (sync mode)
+   export PRESERVE_FLAGS="false"      # Set to "true" to preserve IMAP flags (read/starred/etc)
+   export GMAIL_MODE="false"          # Set to "true" for Gmail mode (All Mail + label application)
    export MAX_WORKERS=4               # Number of parallel threads
    export BATCH_SIZE=10               # Emails per batch
    ```
@@ -139,22 +153,81 @@ All scripts support command-line arguments which take precedence over environmen
 
 **Migration:**
 ```bash
-python3 migrate_imap_emails.py --src-user "me@gmail.com" --dest-user "you@domain.com" --workers 4 --delete
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "me@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "you@domain.com" \
+  --dest-pass "your-app-password" \
+  --workers 4 \
+  --src-delete
 ```
 
 **Comparison:**
 ```bash
-python3 compare_imap_folders.py --src-host "imap.gmail.com" --dest-host "imap.other.com"
+python3 compare_imap_folders.py \
+  --src-host "imap.gmail.com" \
+  --src-user "me@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "you@domain.com" \
+  --dest-pass "your-app-password"
+
+# Compare IMAP source to a local backup folder
+python3 compare_imap_folders.py \
+  --src-host "imap.gmail.com" \
+  --src-user "me@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "./my_backup"
+
+# Compare a local backup folder to an IMAP destination
+python3 compare_imap_folders.py \
+  --src-path "./my_backup" \
+  --dest-host "imap.other.com" \
+  --dest-user "you@domain.com" \
+  --dest-pass "your-app-password"
 ```
 
 **Counting:**
 ```bash
-python3 count_imap_emails.py --host "imap.gmail.com" --user "me@gmail.com" --pass "secret"
+python3 count_imap_emails.py \
+   --host "imap.gmail.com" \
+   --user "me@gmail.com" \
+   --pass "secret"
+
+# Count a local backup folder
+python3 count_imap_emails.py \
+   --path "./my_backup"
+
+# Or via environment variable
+export BACKUP_LOCAL_PATH="./my_backup"
+python3 count_imap_emails.py
+```
+
+**Counting (OAuth2):**
+```bash
+python3 count_imap_emails.py \
+   --host "imap.gmail.com" \
+   --user "me@gmail.com" \
+   --oauth2-client-id "id" \
+   --oauth2-client-secret "secret"
+
+# Or via environment variables (single-account script)
+export IMAP_HOST="imap.gmail.com"
+export IMAP_USERNAME="me@gmail.com"
+export OAUTH2_CLIENT_ID="id"
+export OAUTH2_CLIENT_SECRET="secret"  # Required for Google
+python3 count_imap_emails.py
 ```
 
 **Backup:**
 ```bash
-python3 backup_imap_emails.py --src-user "me@gmail.com" --dest-path "./my_backup"
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "me@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "./my_backup"
 ```
 
 ## Usage Examples
@@ -162,37 +235,118 @@ python3 backup_imap_emails.py --src-user "me@gmail.com" --dest-path "./my_backup
 ### 1. Full Migration
 Migrate all folders from Source to Destination.
 ```bash
-python3 migrate_imap_emails.py
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password"
+```
+
+### 1a. Gmail Mode Migration (Gmail -> Gmail)
+For Gmail -> Gmail migrations, `--gmail-mode` migrates only `[Gmail]/All Mail` (no duplicates) and applies labels by copying messages into label folders.
+
+```bash
+python3 migrate_imap_emails.py \
+  --gmail-mode \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.gmail.com" \
+  --dest-user "dest@gmail.com" \
+  --dest-pass "dest-app-password"
+```
+
+Notes:
+- `--preserve-flags` is enabled automatically in `--gmail-mode`.
+- `--dest-delete` is not supported in `--gmail-mode`.
+
+### 1b. Preserve Flags (Any IMAP Server)
+Preserve IMAP flags (`\Seen`, `\Flagged`, `\Answered`, `\Draft`) during migration.
+
+If an email already exists on the destination (duplicate), the script can still sync missing flags on the existing message.
+
+```bash
+python3 migrate_imap_emails.py \
+  --preserve-flags \
+  --src-host "imap.example.com" \
+  --src-user "source@example.com" \
+  --src-pass "source-password" \
+  --dest-host "imap.example.com" \
+  --dest-user "dest@example.com" \
+  --dest-pass "dest-password"
 ```
 
 ### 2. Single Folder Migration
 Migrate ONLY a specific folder (e.g., trying to fix just "Important" or "Sent").
 ```bash
 # Syntax: python3 migrate_imap_emails.py "[Folder Name]"
-python3 migrate_imap_emails.py "[Gmail]/Important"
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password" \
+  "[Gmail]/Important"
 ```
 
 ### 3. Move Instead of Copy
 Migrate and **delete** from source immediately after verifying the copy.
 ```bash
 # Using flag
-python3 migrate_imap_emails.py --delete
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password" \
+  --src-delete
 
 # Or specific folder with delete
-python3 migrate_imap_emails.py "Inbox" --delete
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password" \
+  "INBOX" \
+  --src-delete
 ```
 
 ### 4. Sync Mode (Delete from Destination)
 Keep destination in sync by deleting emails that no longer exist in the source.
+
+Note: `--dest-delete` is not supported in `--gmail-mode`.
 ```bash
 # Migration: Delete destination emails not found in source
-python3 migrate_imap_emails.py --dest-delete
+python3 migrate_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password" \
+  --dest-delete
 
 # Backup: Delete local .eml files not found on server
-python3 backup_imap_emails.py --dest-path "./backup" --dest-delete
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-path "./backup" \
+  --dest-delete
 
 # Restore: Delete server emails not found in local backup
-python3 restore_imap_emails.py --src-path "./backup" --dest-delete
+python3 restore_imap_emails.py \
+  --src-path "./backup" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password" \
+  --dest-delete
 ```
 
 **Warning:** The `--dest-delete` flag permanently removes emails/files from the destination. Use with caution and always verify your backup is complete before enabling this option.
@@ -200,7 +354,27 @@ python3 restore_imap_emails.py --src-path "./backup" --dest-delete
 ### 5. Verify Migration
 Compare counts between source and destination.
 ```bash
-python3 compare_imap_folders.py
+python3 compare_imap_folders.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password"
+
+# IMAP source -> local backup destination
+python3 compare_imap_folders.py \
+  --src-host "imap.gmail.com" \
+  --src-user "source@gmail.com" \
+  --src-pass "source-app-password" \
+  --dest-path "./my_backup"
+
+# local backup source -> IMAP destination
+python3 compare_imap_folders.py \
+  --src-path "./my_backup" \
+  --dest-host "imap.other.com" \
+  --dest-user "dest@domain.com" \
+  --dest-pass "dest-app-password"
 ```
 *Output Example:*
 ```
@@ -213,15 +387,65 @@ INBOX                   | 1250         | 1250       | MATCH
 ### 6. Local Backup
 Download all your emails to your computer as `.eml` files.
 ```bash
-# Set destination path
-export BACKUP_LOCAL_PATH="./backup_folder"
-python3 backup_imap_emails.py
+# Backup all folders from an IMAP account to a local folder
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "you@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "./backup_folder"
 
 # Or via command line
-python3 backup_imap_emails.py --dest-path "/Users/jdoe/Documents/Emails" 
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "you@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "/Users/jdoe/Documents/Emails"
 
 # Backup single folder
-python3 backup_imap_emails.py --dest-path "./my_backup" "[Gmail]/Sent Mail"
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "you@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "./my_backup" \
+  "[Gmail]/Sent Mail"
+```
+
+### 6a. Compare IMAP vs Local Backup
+Use `compare_imap_folders.py` to validate an IMAP account against a local backup created by `backup_imap_emails.py`.
+
+```bash
+# Option 1: IMAP source -> local destination
+python3 compare_imap_folders.py \
+  --src-host "imap.gmail.com" \
+  --src-user "you@gmail.com" \
+  --src-pass "your-app-password" \
+  --dest-path "./my_backup"
+
+# Option 2: local source -> IMAP destination
+python3 compare_imap_folders.py \
+  --src-path "./my_backup" \
+  --dest-host "imap.other.com" \
+  --dest-user "you@domain.com" \
+  --dest-pass "your-app-password"
+```
+
+You can also set local paths via environment variables:
+
+```bash
+export SRC_LOCAL_PATH="./my_backup"
+export DEST_LOCAL_PATH="./my_backup"
+```
+
+### 6b. Count a Local Backup
+Use `count_imap_emails.py` to get per-folder counts from a local backup created by `backup_imap_emails.py`.
+
+```bash
+# Option 1: explicit path
+python3 count_imap_emails.py --path "./my_backup"
+
+# Option 2: environment variable
+export BACKUP_LOCAL_PATH="./my_backup"
+python3 count_imap_emails.py
 ```
 
 ### 7. Gmail Backup with Labels Preservation
@@ -319,6 +543,9 @@ This creates a `flags_manifest.json` with the status of each email.
 ### 9. Restore Backup to IMAP Server
 Restore emails from a local backup to any IMAP server.
 
+By default, restore runs in **incremental mode**: it uploads only emails that are not already present on the destination (based on `Message-ID`).
+Use `--full-restore` to force the legacy behavior (process all emails and re-sync labels/flags for already-present messages).
+
 ```bash
 # Restore all folders from backup
 python3 restore_imap_emails.py \
@@ -327,6 +554,14 @@ python3 restore_imap_emails.py \
   --dest-user "you@gmail.com" \
   --dest-pass "your-app-password"
 
+# Force full restore (legacy behavior)
+python3 restore_imap_emails.py \
+  --src-path "./my_backup" \
+  --dest-host "imap.gmail.com" \
+  --dest-user "you@gmail.com" \
+  --dest-pass "your-app-password" \
+  --full-restore
+
 # Restore with flags (read/starred status)
 python3 restore_imap_emails.py \
   --src-path "./my_backup" \
@@ -358,7 +593,8 @@ python3 restore_imap_emails.py \
 
 **How Gmail restore works:**
 1. Reads emails from the backup (typically `[Gmail]/All Mail`)
-2. Uploads each email to INBOX with original flags (read/starred/etc)
+2. Uploads each email to the first usable label folder (preserving original flags)
+  - If an email has no usable labels, it is uploaded to `Restored/Unlabeled`
 3. Looks up the Message-ID in `labels_manifest.json`
 4. Copies the email to each label folder (e.g., "Work", "Personal", "Projects/2024")
 
@@ -373,20 +609,102 @@ python3 restore_imap_emails.py \
   --apply-flags
 ```
 
+## OAuth2 Authentication
+
+All scripts support OAuth2 as an alternative to password-based authentication. The OAuth2 provider is **auto-detected** from the IMAP host:
+
+| IMAP Host contains | Detected Provider |
+|---|---|
+| `outlook`, `office365`, `microsoft` | Microsoft |
+| `gmail`, `google` | Google |
+
+To use OAuth2, pass `--oauth2-client-id` (or `--src-oauth2-client-id`/`--dest-oauth2-client-id` for dual-account scripts) instead of the password argument.
+
+Environment variable equivalents:
+- Dual-account scripts: `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET` and `DEST_OAUTH2_CLIENT_ID`, `DEST_OAUTH2_CLIENT_SECRET`.
+- Single-account scripts (like `count_imap_emails.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
+
+### Microsoft (Outlook / Office 365)
+
+Requires the `msal` package (`pip install msal`). Uses the **device code flow** — no browser redirect needed. The tenant ID is auto-discovered from the user's email domain.
+
+#### Creating an App Registration in Microsoft Entra
+
+1. Go to [Microsoft Entra admin center](https://entra.microsoft.com/) → **Identity** → **Applications** → **App registrations**
+2. Click **New registration**
+   - **Name**: e.g., "IMAP Migration Tool"
+   - **Supported account types**: Select based on your needs (single tenant or multi-tenant)
+   - **Redirect URI**: Leave blank (not needed for device code flow)
+3. Click **Register**
+4. Copy the **Application (client) ID** — this is your `--oauth2-client-id`
+
+#### Adding API Permissions
+
+1. In your app registration, go to **API permissions**
+2. Click **Add a permission** → **APIs my organization uses**
+3. Search for and select **Office 365 Exchange Online**
+4. Select **Delegated permissions**
+5. Check **IMAP.AccessAsUser.All**
+6. Click **Add permissions**
+7. (Optional) Click **Grant admin consent** if you have admin rights and want to pre-approve for all users
+
+#### Enabling Public Client Flow
+
+1. Go to **Authentication**
+2. Under **Advanced settings**, set **Allow public client flows** to **Yes**
+3. Click **Save**
+
+#### Usage
+
+```bash
+# Install dependency
+pip install msal
+
+# Migration with Microsoft OAuth2 on source
+python3 migrate_imap_emails.py \
+  --src-host "outlook.office365.com" \
+  --src-user "user@contoso.com" \
+  --src-oauth2-client-id "your-azure-app-client-id" \
+  --dest-host "imap.other.com" \
+  --dest-user "user@other.com" \
+  --dest-pass "password"
+```
+
+The script will print a device code and URL. Open the URL in a browser, enter the code, and sign in to authorize access.
+
+### Google (Gmail)
+
+Requires the `google-auth-oauthlib` package (`pip install google-auth-oauthlib`). Uses the **installed app flow** — opens a browser window for consent. Both `--oauth2-client-id` and `--oauth2-client-secret` are required.
+
+```bash
+# Install dependency
+pip install google-auth-oauthlib
+
+# Backup with Google OAuth2
+python3 backup_imap_emails.py \
+  --src-host "imap.gmail.com" \
+  --src-user "you@gmail.com" \
+  --src-oauth2-client-id "your-google-client-id" \
+  --src-oauth2-client-secret "your-google-client-secret" \
+  --dest-path "./gmail_backup"
+```
+
+The script will open your default browser for Google sign-in. After authorizing, the token is returned automatically via a local HTTP redirect.
+
 ## Troubleshooting
 
-- **"Too many simultaneous connections"**: 
+- **"Too many simultaneous connections"**:
   IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `migrate_imap_emails.py` uses multiple threads, you may hit this limit.
   **Solution**: Reduce `MAX_WORKERS` to `4` or `2` using the environment variable.
 
-- **Authentication Errors**: 
+- **Authentication Errors**:
   If you are using Gmail or Google Workspace, you generally **cannot** use your regular login password. You must enable 2-Step Verification and generate an **App Password**. Use that App Password in the `_PASSWORD` variable.
 
 - **Timeouts / Socket Errors**:
   Migrating 100k+ emails is network intensive. If the script crashes, simply run it again. The built-in de-duplication will skip already migrated messages and resume where it left off.
 
 ### Gmail "All Mail" & Deletion
-If you are migrating **from** a Gmail account and using the `--delete` option:
+If you are migrating **from** a Gmail account and using the `--src-delete` option:
 - The script attempts to detect your Trash folder (e.g., `[Gmail]/Trash` or `[Gmail]/Bin`).
 - Instead of simply marking emails as deleted (which Gmail often treats as "Archive"), the script **copies the email to the Trash folder** and then marks the original as deleted.
 - This ensures that the storage count in `[Gmail]/All Mail` actually decreases, as the emails are moved to the Trash (which is auto-emptied by Google after 30 days) rather than remaining in your "All Mail" archive.

src/backup_imap_emails.py

@@ -13,7 +13,13 @@ Features:
 - Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
 
 Configuration (Environment Variables):
-  SRC_IMAP_HOST, SRC_IMAP_USERNAME, SRC_IMAP_PASSWORD: Source credentials.
+    SRC_IMAP_HOST, SRC_IMAP_USERNAME: Source credentials.
+    SRC_IMAP_PASSWORD: Source password (or App Password).
+
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID: OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
+
   BACKUP_LOCAL_PATH: Destination local directory.
   MAX_WORKERS: Number of concurrent threads (default: 10).
   BATCH_SIZE: Number of emails to process per batch (default: 10).
@@ -25,10 +31,20 @@ Configuration (Environment Variables):
               Default is "false".
 
 Usage:
-  python3 backup_imap_emails.py --dest-path "./my_backup"
+    python3 backup_imap_emails.py \
+        --src-host "imap.example.com" \
+        --src-user "you@example.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup"
 
 Gmail Labels:
-  python3 backup_imap_emails.py --dest-path "./my_backup" --preserve-labels "[Gmail]/All Mail"
+    python3 backup_imap_emails.py \
+        --src-host "imap.gmail.com" \
+        --src-user "you@gmail.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup" \
+        --preserve-labels \
+        "[Gmail]/All Mail"
   This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
   file that maps each email's Message-ID to its Gmail labels for later restoration.
 """
@@ -41,25 +57,17 @@ import sys
 import threading
 
 import imap_common
+import imap_oauth2
 
 # Defaults
 MAX_WORKERS = 10
 BATCH_SIZE = 10
+MANIFEST_FILENAME = "labels_manifest.json"
 
 # Thread-local storage
 thread_local = threading.local()
 print_lock = threading.Lock()
 
-# Gmail-specific folders to exclude from label mapping
-GMAIL_SYSTEM_FOLDERS = {
-    "[Gmail]/All Mail",
-    "[Gmail]/Spam",
-    "[Gmail]/Trash",
-    "[Gmail]/Drafts",
-    "[Gmail]/Bin",
-    "[Gmail]/Important",  # This is actually a label, but often system-managed
-}
-
 
 def safe_print(message):
     t_name = threading.current_thread().name
@@ -69,17 +77,31 @@ def safe_print(message):
 
 
 def get_thread_connection(src_conf):
+    # Backwards-compatible: tests and some internal call sites may still pass
+    # (host, user, password) tuples. Normalize to the dict form used by
+    # imap_common.get_imap_connection_from_conf().
+    if isinstance(src_conf, tuple) and len(src_conf) == 3:
+        src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
+
     if not hasattr(thread_local, "src") or thread_local.src is None:
-        thread_local.src = imap_common.get_imap_connection(*src_conf)
+        thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
     try:
         if thread_local.src:
             thread_local.src.noop()
-    except:
-        thread_local.src = imap_common.get_imap_connection(*src_conf)
+    except Exception:
+        thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
+        # If reconnection failed (possibly expired token), try refreshing
+        if thread_local.src is None and src_conf.get("oauth2"):
+            old_token = src_conf["oauth2_token"]
+            imap_oauth2.refresh_oauth2_token(src_conf, old_token)
+            thread_local.src = imap_common.get_imap_connection_from_conf(src_conf)
     return thread_local.src
 
 
 def process_batch(uids, folder_name, src_conf, local_folder_path):
+    if isinstance(src_conf, tuple) and len(src_conf) == 3:
+        src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
+
     src = get_thread_connection(src_conf)
     if not src:
         safe_print("Error: Could not establish connection for batch.")
@@ -93,32 +115,10 @@ def process_batch(uids, folder_name, src_conf, local_folder_path):
 
     for uid in uids:
         try:
-            # 1. Fetch Subject for Filename
-            # We fetch headers first to generate the nice filename
-            msg_id, size, subject = imap_common.get_msg_details(src, uid)
-
             # Helper to handle byte UIDs
             uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
 
-            if not subject:
-                clean_subject = "No Subject"
-            else:
-                clean_subject = imap_common.sanitize_filename(subject)
-                # Limit subject len to avoid FS path length limits (max 255 usually)
-                clean_subject = clean_subject[:100]
-
-            filename = f"{uid_str}_{clean_subject}.eml"
-            full_path = os.path.join(local_folder_path, filename)
-
-            # Double check existence (in case optimization missed it or naming collision)
-            # Actually, the main purpose here is just to save.
-            # However, if we migrated to a new naming convention, we might have duplicates with different names?
-            # The incremental check in `backup_folder` relies on UID prefix, so we are safe.
-
-            if os.path.exists(full_path):
-                continue
-
-            # 2. Fetch Full Content
+            # Fetch Full Content
             # RFC822 gets the whole message including headers and attachments
             resp, data = src.uid("fetch", uid, "(RFC822)")
             if resp != "OK" or not data or data[0] is None:
@@ -131,6 +131,22 @@ def process_batch(uids, folder_name, src_conf, local_folder_path):
                     raw_email = item[1]
                     break
 
+            # Derive Subject for filename from the already-fetched message bytes.
+            _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
+            if not subject:
+                clean_subject = "No Subject"
+            else:
+                clean_subject = imap_common.sanitize_filename(subject)
+                # Limit subject len to avoid FS path length limits (max 255 usually)
+                clean_subject = clean_subject[:100]
+
+            filename = f"{uid_str}_{clean_subject}.eml"
+            full_path = os.path.join(local_folder_path, filename)
+
+            # Double check existence (in case incremental UID checks missed it or naming collision).
+            if os.path.exists(full_path):
+                continue
+
             if raw_email:
                 try:
                     with open(full_path, "wb") as f:
@@ -174,15 +190,15 @@ def is_gmail_label_folder(folder_name):
     Excludes system folders like All Mail, Spam, Trash, Drafts.
     """
     # Exclude system folders that aren't really "labels"
-    if folder_name in GMAIL_SYSTEM_FOLDERS:
+    if folder_name in imap_common.GMAIL_SYSTEM_FOLDERS:
         return False
 
-    # INBOX is a special case - it's a label in Gmail
-    if folder_name == "INBOX":
+    # imap_common.FOLDER_INBOX is a special case - it's a label in Gmail
+    if folder_name == imap_common.FOLDER_INBOX:
         return True
 
     # [Gmail]/Sent Mail and [Gmail]/Starred are labels worth preserving
-    if folder_name in ("[Gmail]/Sent Mail", "[Gmail]/Starred"):
+    if folder_name in (imap_common.GMAIL_SENT, imap_common.GMAIL_STARRED):
         return True
 
     # Any folder NOT under [Gmail]/ is a user label
@@ -195,7 +211,7 @@ def is_gmail_label_folder(folder_name):
 # Standard IMAP flags that can be preserved during migration
 # \Recent is session-specific and cannot be set by clients
 # \Deleted should not be preserved as it marks messages for removal
-PRESERVABLE_FLAGS = {"\\Seen", "\\Answered", "\\Flagged", "\\Draft"}
+PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
 
 
 def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
@@ -249,21 +265,14 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
 
                         # Extract all preservable flags from the metadata
                         flags = []
-                        for flag in PRESERVABLE_FLAGS:
+                        for flag in imap_common.PRESERVABLE_FLAGS:
                             if flag in meta_str:
                                 flags.append(flag)
 
                         # Second element contains the header
-                        header_data = item[1]
-                        if isinstance(header_data, bytes):
-                            header_str = header_data.decode("utf-8", errors="ignore")
-                            # Extract Message-ID from header
-                            for line in header_str.split("\n"):
-                                if line.lower().startswith("message-id:"):
-                                    msg_id = line.split(":", 1)[1].strip()
-                                    if msg_id:
-                                        message_info[msg_id] = {"flags": flags}
-                                    break
+                        msg_id = imap_common.extract_message_id(item[1])
+                        if msg_id:
+                            message_info[msg_id] = {"flags": flags}
             except Exception as e:
                 safe_print(f"Error fetching batch in {folder_name}: {e}")
                 # Try to keep connection alive
@@ -306,14 +315,10 @@ def build_labels_manifest(imap_conn, local_path):
 
     safe_print("--- Building Gmail Labels Manifest ---")
 
-    # Get all folders
-    try:
-        typ, folders = imap_conn.list()
-        if typ != "OK":
-            safe_print("Error: Could not list folders for label mapping.")
-            return manifest
-    except Exception as e:
-        safe_print(f"Error listing folders: {e}")
+    # Get all selectable folders and filter to label folders
+    all_folders = imap_common.list_selectable_folders(imap_conn)
+    if not all_folders:
+        safe_print("Error: Could not list folders for label mapping.")
         return manifest
 
     # First, scan [Gmail]/All Mail to get the authoritative flags for all emails
@@ -328,7 +333,7 @@ def build_labels_manifest(imap_conn, local_path):
             flush=True,
         )
 
-    all_mail_info = get_message_info_in_folder(imap_conn, "[Gmail]/All Mail", all_mail_progress_cb)
+    all_mail_info = get_message_info_in_folder(imap_conn, imap_common.GMAIL_ALL_MAIL, all_mail_progress_cb)
     print()  # New line after progress
 
     # Initialize manifest with flags from All Mail
@@ -337,8 +342,8 @@ def build_labels_manifest(imap_conn, local_path):
 
     all_mail_elapsed = time.time() - all_mail_start
     # Count flag statistics
-    read_count = sum(1 for m in manifest.values() if "\\Seen" in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if "\\Flagged" in m.get("flags", []))
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
     safe_print(f"  -> {len(all_mail_info)} emails scanned ({all_mail_elapsed:.1f}s)")
     safe_print(f"  -> Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}\n")
 
@@ -349,11 +354,7 @@ def build_labels_manifest(imap_conn, local_path):
         pass
 
     # Parse folder names and filter to label folders
-    label_folders = []
-    for f_info in folders:
-        name = imap_common.normalize_folder_name(f_info)
-        if is_gmail_label_folder(name):
-            label_folders.append(name)
+    label_folders = [f for f in all_folders if is_gmail_label_folder(f)]
 
     total_folders = len(label_folders)
     safe_print(f"Found {total_folders} label folders to scan.\n")
@@ -404,8 +405,8 @@ def build_labels_manifest(imap_conn, local_path):
 
     # Summary
     total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if "\\Seen" in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if "\\Flagged" in m.get("flags", []))
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
     safe_print("\nManifest building complete:")
     safe_print(f"  - Folders scanned: {total_folders + 1}")  # +1 for All Mail
     safe_print(f"  - Total email-label mappings: {total_emails_scanned}")
@@ -414,7 +415,7 @@ def build_labels_manifest(imap_conn, local_path):
     safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
 
     # Save manifest
-    manifest_path = os.path.join(local_path, "labels_manifest.json")
+    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
     try:
         with open(manifest_path, "w", encoding="utf-8") as f:
             json.dump(manifest, f, indent=2, ensure_ascii=False)
@@ -489,8 +490,8 @@ def build_flags_manifest(imap_conn, local_path, folders_to_scan=None):
 
     # Summary
     total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if "\\Seen" in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if "\\Flagged" in m.get("flags", []))
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
     safe_print("\nFlags manifest building complete:")
     safe_print(f"  - Folders scanned: {total_folders}")
     safe_print(f"  - Unique emails: {len(manifest)}")
@@ -514,7 +515,7 @@ def load_labels_manifest(local_path):
     Loads an existing labels manifest from the backup directory.
     Returns the manifest dict or empty dict if not found.
     """
-    manifest_path = os.path.join(local_path, "labels_manifest.json")
+    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
     if os.path.exists(manifest_path):
         try:
             with open(manifest_path, encoding="utf-8") as f:
@@ -662,13 +663,50 @@ def main():
     parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
 
     # Source
-    parser.add_argument("--src-host", default=os.getenv("SRC_IMAP_HOST"), help="Source IMAP Server")
-    parser.add_argument("--src-user", default=os.getenv("SRC_IMAP_USERNAME"), help="Source Username")
-    parser.add_argument("--src-pass", default=os.getenv("SRC_IMAP_PASSWORD"), help="Source Password")
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+
+    # Authentication: require either password OR OAuth2 client-id (unless provided via env vars)
+    auth_required = not bool(default_src_pass or default_src_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    # OAuth2
+    auth_group.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
 
     # Destination (Local Path)
     env_path = os.getenv("BACKUP_LOCAL_PATH")
-    parser.add_argument("--dest-path", default=env_path, help="Local destination path (Mandatory)")
+    parser.add_argument(
+        "--dest-path",
+        default=env_path,
+        required=not bool(env_path),
+        help="Local destination path (or BACKUP_LOCAL_PATH)",
+    )
 
     # Config
     parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
@@ -717,29 +755,44 @@ def main():
 
     args = parser.parse_args()
 
-    # Validate
-    missing = []
-    if not args.src_host:
-        missing.append("SRC_IMAP_HOST")
-    if not args.src_user:
-        missing.append("SRC_IMAP_USERNAME")
-    if not args.src_pass:
-        missing.append("SRC_IMAP_PASSWORD")
-
-    if missing:
-        print(f"Error: Missing credentials: {', '.join(missing)}")
-        sys.exit(1)
-
-    if not args.dest_path:
-        print("Error: Destination path is required.")
-        print("Please provide --dest-path or set environment variable BACKUP_LOCAL_PATH.")
-        sys.exit(1)
+    use_oauth2 = bool(args.src_client_id)
 
     global MAX_WORKERS, BATCH_SIZE
     MAX_WORKERS = args.workers
     BATCH_SIZE = args.batch
 
-    src_conf = (args.src_host, args.src_user, args.src_pass)
+    # 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,
+    }
 
     # Expand path (~/...)
     local_path = os.path.expanduser(args.dest_path)
@@ -754,6 +807,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"Destination Path: {local_path}")
     if args.gmail_mode:
         print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
@@ -770,7 +824,7 @@ def main():
     print("-----------------------------\n")
 
     try:
-        src = imap_common.get_imap_connection(*src_conf)
+        src = imap_common.get_imap_connection_from_conf(src_conf)
         if not src:
             sys.exit(1)
 
@@ -793,7 +847,7 @@ def main():
         # If manifest-only mode, we're done
         if args.manifest_only:
             src.logout()
-            manifest_path = os.path.join(local_path, "labels_manifest.json")
+            manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
             print("\nManifest-only mode complete.")
             print(f"Labels manifest saved to: {manifest_path}")
             print("\nTo download emails, run again without --manifest-only:")
@@ -802,15 +856,20 @@ def main():
 
         # Gmail mode: backup only [Gmail]/All Mail
         if args.gmail_mode:
-            backup_folder(src, "[Gmail]/All Mail", local_path, src_conf, args.dest_delete)
+            backup_folder(src, imap_common.GMAIL_ALL_MAIL, local_path, src_conf, args.dest_delete)
         elif args.folder:
             backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
         else:
-            typ, folders = src.list()
-            if typ == "OK":
-                for f_info in folders:
-                    name = imap_common.normalize_folder_name(f_info)
-                    backup_folder(src, name, local_path, src_conf, args.dest_delete)
+            folders = imap_common.list_selectable_folders(src)
+            for name in folders:
+                try:
+                    src.noop()
+                except Exception:
+                    if src_conf.get("oauth2"):
+                        safe_print("Refreshing OAuth2 token...")
+                        imap_oauth2.refresh_oauth2_token(src_conf, src_conf["oauth2_token"])
+
+                backup_folder(src, name, local_path, src_conf, args.dest_delete)
 
         src.logout()
         print("\nBackup completed successfully.")

src/compare_imap_folders.py

@@ -1,7 +1,8 @@
 """
 IMAP Folder Comparison Script
 
-This script compares email counts between a source IMAP account and a destination IMAP account.
+This script compares email counts between a source and a destination.
+Each side can be either an IMAP account or a local backup folder.
 It iterates through all folders found in the source account and checks the corresponding
 folder in the destination account.
 
@@ -11,23 +12,62 @@ Configuration (Environment Variables):
     SRC_IMAP_USERNAME   : Source Username/Email
     SRC_IMAP_PASSWORD   : Source Password
 
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
   Destination Account:
     DEST_IMAP_HOST      : Destination IMAP Host
     DEST_IMAP_USERNAME  : Destination Username/Email
     DEST_IMAP_PASSWORD  : Destination Password
 
+    OAuth2 (Optional - instead of password):
+    DEST_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    DEST_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
+Also supports local folders as source and/or destination:
+    SRC_LOCAL_PATH      : Source local folder (backup root)
+    DEST_LOCAL_PATH     : Destination local folder (backup root)
+
 Usage:
   python3 compare_imap_folders.py
+
+Examples:
+        # IMAP -> IMAP
+        python3 compare_imap_folders.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # Local -> IMAP
+        python3 compare_imap_folders.py \
+            --src-path "./my_backup" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # IMAP -> Local
+        python3 compare_imap_folders.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-path "./my_backup"
 """
 
 import argparse
 import os
 import sys
+from typing import Optional
 
 import imap_common
+import imap_oauth2
 
 
 def get_email_count(conn, folder_name):
+    """Return the IMAP message count for a folder, or None on error."""
     try:
         # Select folder in read-only mode
         # Quote folder name handles spaces
@@ -46,76 +86,258 @@ def get_email_count(conn, folder_name):
         return None
 
 
+def _is_ignored_local_dir(dirname: str) -> bool:
+    return dirname.startswith(".") or dirname == "__pycache__"
+
+
+def list_local_folders(local_root: str) -> list[str]:
+    """List all folders under a local backup root in IMAP-style names.
+
+    The local backup format is expected to mirror IMAP folder hierarchy using
+    subdirectories (e.g. "[Gmail]/All Mail" becomes "[Gmail]/All Mail/").
+    """
+    folders: set[str] = set()
+
+    for dirpath, dirnames, _filenames in os.walk(local_root):
+        dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)]
+
+        if os.path.abspath(dirpath) == os.path.abspath(local_root):
+            continue
+
+        rel = os.path.relpath(dirpath, local_root)
+        if rel == ".":
+            continue
+
+        parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)]
+        if not parts:
+            continue
+
+        folders.add("/".join(parts))
+
+    return sorted(folders)
+
+
+def get_local_email_count(local_root: str, folder_name: str) -> Optional[int]:
+    """Return the count of .eml files in a local folder, or None if missing/unreadable."""
+    folder_path = os.path.join(local_root, *folder_name.split("/"))
+    if not os.path.isdir(folder_path):
+        return None
+
+    try:
+        count = 0
+        for filename in os.listdir(folder_path):
+            if not filename.endswith(".eml"):
+                continue
+            full_path = os.path.join(folder_path, filename)
+            if os.path.isfile(full_path):
+                count += 1
+        return count
+    except OSError:
+        return None
+
+
 def main():
+    default_src_path = os.getenv("SRC_LOCAL_PATH")
+    default_dest_path = os.getenv("DEST_LOCAL_PATH")
+
+    # Phase 1: determine whether each side is local
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--src-path", default=default_src_path)
+    pre_parser.add_argument("--dest-path", default=default_dest_path)
+    pre_args, _ = pre_parser.parse_known_args()
+
+    src_requires_imap = not bool(pre_args.src_path)
+    dest_requires_imap = not bool(pre_args.dest_path)
+
+    # Phase 2: full parser with conditional requirements
     parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
+    parser.add_argument(
+        "--src-path",
+        default=default_src_path,
+        help="Source local folder (backup root). If set, IMAP source args are ignored.",
+    )
+    parser.add_argument(
+        "--dest-path",
+        default=default_dest_path,
+        help="Destination local folder (backup root). If set, IMAP destination args are ignored.",
+    )
 
     # Source args
-    parser.add_argument("--src-host", default=os.getenv("SRC_IMAP_HOST"), help="Source IMAP Server")
-    parser.add_argument("--src-user", default=os.getenv("SRC_IMAP_USERNAME"), help="Source Username")
-    parser.add_argument("--src-pass", default=os.getenv("SRC_IMAP_PASSWORD"), help="Source Password")
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=src_requires_imap and not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=src_requires_imap and not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+    src_auth_required = src_requires_imap and not bool(default_src_pass or default_src_client_id)
+    src_auth = parser.add_mutually_exclusive_group(required=src_auth_required)
+    src_auth.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    src_auth.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="Source OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="Source OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
 
     # Dest args
-    parser.add_argument("--dest-host", default=os.getenv("DEST_IMAP_HOST"), help="Destination IMAP Server")
-    parser.add_argument("--dest-user", default=os.getenv("DEST_IMAP_USERNAME"), help="Destination Username")
-    parser.add_argument("--dest-pass", default=os.getenv("DEST_IMAP_PASSWORD"), help="Destination Password")
+    default_dest_host = os.getenv("DEST_IMAP_HOST")
+    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
+    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
+    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--dest-host",
+        default=default_dest_host,
+        required=dest_requires_imap and not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--dest-user",
+        default=default_dest_user,
+        required=dest_requires_imap and not bool(default_dest_user),
+        help="Destination Username (or DEST_IMAP_USERNAME)",
+    )
+    dest_auth_required = dest_requires_imap and not bool(default_dest_pass or default_dest_client_id)
+    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
+    dest_auth.add_argument(
+        "--dest-pass",
+        default=default_dest_pass,
+        help="Destination Password (or DEST_IMAP_PASSWORD)",
+    )
+    dest_auth.add_argument(
+        "--dest-oauth2-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
+    )
+    dest_auth.add_argument(
+        "--dest-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--dest-oauth2-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--dest-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help=argparse.SUPPRESS,
+    )
 
     args = parser.parse_args()
 
-    # Assign to variables
+    src_is_local = bool(args.src_path)
+    dest_is_local = bool(args.dest_path)
+
     SRC_HOST = args.src_host
     SRC_USER = args.src_user
-    SRC_PASS = args.src_pass
     DEST_HOST = args.dest_host
     DEST_USER = args.dest_user
-    DEST_PASS = args.dest_pass
 
-    # Validation
-    missing_vars = []
-    if not SRC_HOST:
-        missing_vars.append("SRC_IMAP_HOST")
-    if not SRC_USER:
-        missing_vars.append("SRC_IMAP_USERNAME")
-    if not SRC_PASS:
-        missing_vars.append("SRC_IMAP_PASSWORD")
-    if not DEST_HOST:
-        missing_vars.append("DEST_IMAP_HOST")
-    if not DEST_USER:
-        missing_vars.append("DEST_IMAP_USERNAME")
-    if not DEST_PASS:
-        missing_vars.append("DEST_IMAP_PASSWORD")
+    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
 
-    if missing_vars:
-        print(f"Error: Missing configuration variables: {', '.join(missing_vars)}")
-        print("Please provide them via environment variables or command-line arguments.")
-        sys.exit(1)
+    # 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")
 
     print("\n--- Configuration Summary ---")
-    print(f"Source Host     : {SRC_HOST}")
-    print(f"Source User     : {SRC_USER}")
-    print(f"Destination Host: {DEST_HOST}")
-    print(f"Destination User: {DEST_USER}")
+    if src_is_local:
+        print(f"Source (Local)  : {args.src_path}")
+    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)'}"
+        )
+
+    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("-----------------------------\n")
 
     src = None
     dest = None
 
     try:
-        # Connect to Source
-        print("Connecting to Source...")
-        src = imap_common.get_imap_connection(SRC_HOST, SRC_USER, SRC_PASS)
-        if not src:
-            return
+        if not src_is_local:
+            # Connect to Source
+            print("Connecting to Source...")
+            src = imap_common.get_imap_connection(args.src_host, args.src_user, args.src_pass, src_oauth2_token)
+            if not src:
+                return
 
-        # Connect to Dest
-        print("Connecting to Destination...")
-        dest = imap_common.get_imap_connection(DEST_HOST, DEST_USER, DEST_PASS)
-        if not dest:
-            return
+        if not dest_is_local:
+            # Connect to Dest
+            print("Connecting to Destination...")
+            dest = imap_common.get_imap_connection(args.dest_host, args.dest_user, args.dest_pass, dest_oauth2_token)
+            if not dest:
+                return
 
         # List Source Folders
         print("Listing folders in Source...")
-        typ, folders = src.list()
-        if typ != "OK":
+        if src_is_local:
+            folders = list_local_folders(args.src_path)
+        else:
+            folders = imap_common.list_selectable_folders(src)
+
+        if not folders:
             print("Failed to list source folders.")
             return
 
@@ -129,12 +351,17 @@ def main():
         total_dest = 0
 
         # Iterate through Source folders
-        for folder_info in folders:
-            folder_name = imap_common.normalize_folder_name(folder_info)
-
+        for folder_name in folders:
             # Get Counts
-            src_count = get_email_count(src, folder_name)
-            dest_count = get_email_count(dest, folder_name)
+            if src_is_local:
+                src_count = get_local_email_count(args.src_path, folder_name)
+            else:
+                src_count = get_email_count(src, folder_name)
+
+            if dest_is_local:
+                dest_count = get_local_email_count(args.dest_path, folder_name)
+            else:
+                dest_count = get_email_count(dest, folder_name)
 
             # Format for display
             src_str = str(src_count) if src_count is not None else "Err"

src/count_imap_emails.py

@@ -1,44 +1,69 @@
-"""
-IMAP Email Counting Script
+"""IMAP Email Counting Script.
 
-This script connects to an IMAP server, iterates through all available folders/mailboxes,
-and counts the number of emails in each. It provides a progressive output of counts
-per folder and a grand total at the end.
+Counts emails per folder from either:
+- An IMAP account, or
+- A local backup folder created by ``backup_imap_emails.py`` (counts ``.eml`` files).
 
 Configuration (Environment Variables):
-  IMAP_HOST     : IMAP Host (e.g., imap.gmail.com)
-  IMAP_USERNAME : Username/Email
-  IMAP_PASSWORD : Password (or App Password)
+    IMAP_HOST        : IMAP Host (e.g., imap.gmail.com)
+    IMAP_USERNAME    : Username/Email
+    IMAP_PASSWORD    : Password (or App Password)
 
-Usage Example:
-  export IMAP_HOST="imap.gmail.com"
-  export IMAP_USERNAME="user@gmail.com"
-  export IMAP_PASSWORD="secretpassword"
+    OAuth2 (Optional - instead of password):
+    OAUTH2_CLIENT_ID        : OAuth2 Client ID
+    OAUTH2_CLIENT_SECRET    : OAuth2 Client Secret (required for Google)
+    SRC_OAUTH2_CLIENT_ID    : Alternate OAuth2 client ID env var
+    SRC_OAUTH2_CLIENT_SECRET: Alternate OAuth2 client secret env var
 
-  python3 count_imap_emails.py
+Local backup counting:
+    BACKUP_LOCAL_PATH : Local backup root (preferred)
+    SRC_LOCAL_PATH    : Alternate local backup root
+
+Examples:
+    # Count an IMAP account
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export IMAP_PASSWORD="secretpassword"
+    python3 count_imap_emails.py
+
+    # Count an IMAP account using OAuth2
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export OAUTH2_CLIENT_ID="your-client-id"
+    export OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
+    python3 count_imap_emails.py
+
+    # Count a local backup
+    python3 count_imap_emails.py --path "./my_backup"
+
+    # Or set a default local backup path via env var
+    export BACKUP_LOCAL_PATH="./my_backup"
+    python3 count_imap_emails.py
 """
 
 import argparse
 import imaplib
 import os
 import sys
+from typing import Optional
 
 import imap_common
+import imap_oauth2
 
 
-def count_emails(imap_server, username, password):
+def count_emails(imap_server, username, password=None, oauth2_token=None):
     try:
         # Connect to the IMAP server (using SSL)
         print(f"Connecting to {imap_server}...")
-        mail = imap_common.get_imap_connection(imap_server, username, password)
+        mail = imap_common.get_imap_connection(imap_server, username, password, oauth2_token)
         if not mail:
             return
 
         # List all mailboxes
         print("Listing mailboxes...")
-        status, folders = mail.list()
+        folders = imap_common.list_selectable_folders(mail)
 
-        if status != "OK":
+        if not folders:
             print("Failed to list mailboxes.")
             return
 
@@ -46,8 +71,7 @@ def count_emails(imap_server, username, password):
         print(f"{'Folder Name':<40} {'Count':>10}")
         print("-" * 52)
 
-        for folder_info in folders:
-            folder_name = imap_common.normalize_folder_name(folder_info)
+        for folder_name in folders:
             display_name = folder_name
 
             try:
@@ -85,32 +109,188 @@ def count_emails(imap_server, username, password):
         print(f"An error occurred: {e}")
 
 
-if __name__ == "__main__":
+def _is_ignored_local_dir(dirname: str) -> bool:
+    return dirname.startswith(".") or dirname == "__pycache__"
+
+
+def list_local_folders(local_root: str) -> list[str]:
+    """List all folders under a local backup root in IMAP-style names."""
+    folders: set[str] = set()
+
+    for dirpath, dirnames, _filenames in os.walk(local_root):
+        dirnames[:] = [d for d in dirnames if not _is_ignored_local_dir(d)]
+
+        if os.path.abspath(dirpath) == os.path.abspath(local_root):
+            continue
+
+        rel = os.path.relpath(dirpath, local_root)
+        if rel == ".":
+            continue
+
+        parts = [p for p in rel.split(os.sep) if p and not _is_ignored_local_dir(p)]
+        if not parts:
+            continue
+
+        folders.add("/".join(parts))
+
+    return sorted(folders)
+
+
+def get_local_email_count(local_root: str, folder_name: str):
+    """Return the count of .eml files in a local folder, or None if missing/unreadable."""
+    folder_path = os.path.join(local_root, *folder_name.split("/"))
+    if not os.path.isdir(folder_path):
+        return None
+
+    try:
+        count = 0
+        for filename in os.listdir(folder_path):
+            if not filename.endswith(".eml"):
+                continue
+            full_path = os.path.join(folder_path, filename)
+            if os.path.isfile(full_path):
+                count += 1
+        return count
+    except OSError:
+        return None
+
+
+def count_local_emails(local_path: str) -> None:
+    print(f"Scanning local backup: {local_path}")
+
+    folders = list_local_folders(local_path)
+    if not folders:
+        print("No folders found.")
+        return
+
+    total_all_folders = 0
+    print(f"{'Folder Name':<40} {'Count':>10}")
+    print("-" * 52)
+
+    for folder_name in folders:
+        count = get_local_email_count(local_path, folder_name)
+        if count is None:
+            print(f"{folder_name:<40} {'N/A':>10}")
+            continue
+
+        print(f"{folder_name:<40} {count:>10}")
+        total_all_folders += count
+
+    print("-" * 52)
+    print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+
+def main(argv: Optional[list[str]] = None) -> None:
+    # Phase 1: determine whether we're in local mode (--path)
+    default_path = os.getenv("BACKUP_LOCAL_PATH") or os.getenv("SRC_LOCAL_PATH")
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--path", default=default_path)
+    pre_args, _ = pre_parser.parse_known_args(argv)
+    require_imap = not bool(pre_args.path)
+
+    # Phase 2: full parser with conditional requirements
     parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
 
+    parser.add_argument(
+        "--path",
+        default=default_path,
+        help="Local backup root to count (counts .eml files per folder). If set, IMAP args are ignored.",
+    )
+
     # Try to unify var names for defaults. Priority: IMAP_* > SRC_IMAP_* > None
     default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
     default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
     default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
+    default_client_id = os.getenv("OAUTH2_CLIENT_ID") or os.getenv("SRC_OAUTH2_CLIENT_ID")
 
-    parser.add_argument("--host", default=default_host, help="IMAP Server")
-    parser.add_argument("--user", default=default_user, help="Username")
-    parser.add_argument("--pass", dest="password", default=default_pass, help="Password")
+    parser.add_argument(
+        "--host",
+        default=default_host,
+        required=require_imap and not bool(default_host),
+        help="IMAP Server (or IMAP_HOST / SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--user",
+        default=default_user,
+        required=require_imap and not bool(default_user),
+        help="Username (or IMAP_USERNAME / SRC_IMAP_USERNAME)",
+    )
 
-    args = parser.parse_args()
+    auth_required = require_imap and not bool(default_pass or default_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument(
+        "--pass", dest="password", default=default_pass, help="Password (or IMAP_PASSWORD / SRC_IMAP_PASSWORD)"
+    )
+    auth_group.add_argument(
+        "--oauth2-client-id",
+        default=default_client_id,
+        dest="client_id",
+        help="OAuth2 Client ID (or OAUTH2_CLIENT_ID / SRC_OAUTH2_CLIENT_ID)",
+    )
+    auth_group.add_argument(
+        "--client-id",
+        default=default_client_id,
+        dest="client_id",
+        help=argparse.SUPPRESS,
+    )
+
+    parser.add_argument(
+        "--oauth2-client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help="OAuth2 Client Secret (if required) (or OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    args = parser.parse_args(argv)
+
+    if args.path:
+        if not os.path.isdir(args.path):
+            print(f"Error: Local path does not exist or is not a directory: {args.path}")
+            sys.exit(1)
+
+        print("\n--- Configuration Summary ---")
+        print(f"Local Path      : {args.path}")
+        print("-----------------------------\n")
+        count_local_emails(args.path)
+        raise SystemExit(0)
 
     IMAP_SERVER = args.host
     USERNAME = args.user
     PASSWORD = args.password
 
-    if not all([IMAP_SERVER, USERNAME, PASSWORD]):
-        print("Error: Missing credentials.")
-        print("Please provide --host, --user, --pass via CLI or set IMAP_* / SRC_IMAP_* environment variables.")
-        sys.exit(1)
+    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 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("-----------------------------\n")
 
-    count_emails(IMAP_SERVER, USERNAME, PASSWORD)
+    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
+
+
+if __name__ == "__main__":
+    main()

src/imap_common.py

@@ -4,13 +4,113 @@ IMAP Common Utilities
 Shared functionality for IMAP migration, counting, and comparison scripts.
 """
 
+from __future__ import annotations
+
 import imaplib
 import os
 import re
 import sys
+from email import policy
 from email.header import decode_header
 from email.parser import BytesParser
 
+# Standard IMAP flags
+FLAG_SEEN = "\\Seen"
+FLAG_ANSWERED = "\\Answered"
+FLAG_FLAGGED = "\\Flagged"
+FLAG_DRAFT = "\\Draft"
+FLAG_DELETED = "\\Deleted"
+FLAG_DELETED_LITERAL = "(\\Deleted)"
+
+# Standard IMAP flags that can be preserved during migration
+# \Recent is session-specific and cannot be set by clients
+# \Deleted should not be preserved as it marks messages for removal
+PRESERVABLE_FLAGS = {FLAG_SEEN, FLAG_ANSWERED, FLAG_FLAGGED, FLAG_DRAFT}
+
+# Gmail constants
+GMAIL_ALL_MAIL = "[Gmail]/All Mail"
+GMAIL_TRASH = "[Gmail]/Trash"
+GMAIL_SPAM = "[Gmail]/Spam"
+GMAIL_DRAFTS = "[Gmail]/Drafts"
+GMAIL_BIN = "[Gmail]/Bin"
+GMAIL_IMPORTANT = "[Gmail]/Important"
+GMAIL_SENT = "[Gmail]/Sent Mail"
+GMAIL_STARRED = "[Gmail]/Starred"
+
+GMAIL_SYSTEM_FOLDERS = {
+    GMAIL_ALL_MAIL,
+    GMAIL_SPAM,
+    GMAIL_TRASH,
+    GMAIL_DRAFTS,
+    GMAIL_BIN,
+    GMAIL_IMPORTANT,
+}
+
+# IMAP Folder Constants
+FOLDER_INBOX = "INBOX"
+
+# Gmail-mode restore/migrate fallback folder for messages with no usable labels.
+# Keeping this as a normal folder (not [Gmail]/Drafts) avoids populating Drafts with non-drafts.
+FOLDER_RESTORED_UNLABELED = "Restored/Unlabeled"
+
+# IMAP Commands
+CMD_STORE = "store"
+CMD_SEARCH = "search"
+CMD_FETCH = "fetch"
+OP_ADD_FLAGS = "+FLAGS"
+
+
+def ensure_folder_exists(imap_conn, folder_name: str) -> None:
+    """Best-effort create of a folder if it doesn't already exist.
+
+    IMAP servers typically return an error if the mailbox exists; this helper
+    intentionally ignores those errors.
+    """
+    try:
+        if folder_name and folder_name.upper() != FOLDER_INBOX:
+            imap_conn.create(f'"{folder_name}"')
+    except Exception:
+        # Folder may already exist, or server may restrict creation.
+        pass
+
+
+def append_email(
+    imap_conn,
+    folder_name: str,
+    raw_content: bytes,
+    date_str: str,
+    flags: str | None = None,
+    *,
+    ensure_folder: bool = True,
+) -> bool:
+    """Append an email message to a folder.
+
+    This is intentionally a thin wrapper around IMAP APPEND; callers can
+    perform duplicate checks or folder selection separately.
+
+    Args:
+        flags: Optional IMAP flags. If provided, they are normalized to a
+            parenthesized list before being passed to `imaplib.IMAP4.append`.
+        ensure_folder: If True, attempts to create the folder first (best-effort).
+    """
+    try:
+        if ensure_folder:
+            ensure_folder_exists(imap_conn, folder_name)
+
+        normalized_flags = None
+        if flags:
+            stripped = str(flags).strip()
+            if stripped:
+                if stripped.startswith("(") and stripped.endswith(")"):
+                    normalized_flags = stripped
+                else:
+                    normalized_flags = f"({stripped})"
+
+        resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+        return resp == "OK"
+    except Exception:
+        return False
+
 
 def verify_env_vars(vars_list):
     """
@@ -25,18 +125,43 @@ def verify_env_vars(vars_list):
     return True
 
 
-def get_imap_connection(host, user, password):
+def get_imap_connection_from_conf(conf):
+    """
+    Establishes an IMAP connection using a conf dict.
+
+    conf dict structure:
+        {
+            "host": str,
+            "user": str,
+            "password": str or None,
+            "oauth2_token": str or None,
+            "oauth2": dict or None  # Contains provider, client_id, email, client_secret
+        }
+    """
+    return get_imap_connection(conf["host"], conf["user"], conf.get("password"), conf.get("oauth2_token"))
+
+
+def get_imap_connection(host, user, password=None, oauth2_token=None):
     """
     Establishes an SSL connection to the IMAP server and logs in.
+    Supports both basic auth (password) and OAuth 2.0 (XOAUTH2).
     Returns the connection object or None if failed.
     """
-    if not all([host, user, password]):
+    if not host or not user:
         print(f"Error: Invalid credentials for {host}")
         return None
 
+    if not password and not oauth2_token:
+        print(f"Error: Either password or oauth2_token is required for {host}")
+        return None
+
     try:
         conn = imaplib.IMAP4_SSL(host)
-        conn.login(user, password)
+        if oauth2_token:
+            auth_string = f"user={user}\x01auth=Bearer {oauth2_token}\x01\x01"
+            conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
+        else:
+            conn.login(user, password)
         return conn
     except Exception as e:
         print(f"Connection error to {host}: {e}")
@@ -64,73 +189,122 @@ def normalize_folder_name(folder_info_str):
     return folder_info_str.split()[-1].strip('"')
 
 
+def list_selectable_folders(imap_conn):
+    """
+    Lists all selectable folders (excluding \\Noselect) on the IMAP connection.
+    Returns a list of normalized folder name strings, or an empty list on failure.
+    """
+    try:
+        status, folders = imap_conn.list()
+        if status != "OK" or not folders:
+            return []
+    except Exception:
+        return []
+
+    result = []
+    for f in folders:
+        f_str = f.decode("utf-8", errors="ignore") if isinstance(f, bytes) else str(f)
+        if "\\Noselect" in f_str:
+            continue
+        result.append(normalize_folder_name(f))
+    return result
+
+
 def decode_mime_header(header_value):
     """
-    Decodes MIME encoded headers (Subject, etc.) to a unicode (str) string.
+    Decodes MIME encoded headers (Subject, etc.) to a unicode string.
     """
     if not header_value:
         return "(No Subject)"
     try:
         decoded_list = decode_header(header_value)
-        default_charset = "utf-8"
         text_parts = []
-        for bytes_data, encoding in decoded_list:
-            if isinstance(bytes_data, bytes):
-                if encoding:
-                    try:
-                        text_parts.append(bytes_data.decode(encoding, errors="ignore"))
-                    except LookupError:
-                        text_parts.append(bytes_data.decode(default_charset, errors="ignore"))
-                else:
-                    text_parts.append(bytes_data.decode(default_charset, errors="ignore"))
+        for data, encoding in decoded_list:
+            if isinstance(data, bytes):
+                charset = encoding or "utf-8"
+                try:
+                    text_parts.append(data.decode(charset, errors="ignore"))
+                except LookupError:
+                    text_parts.append(data.decode("utf-8", errors="ignore"))
             else:
-                text_parts.append(str(bytes_data))
+                text_parts.append(str(data))
         return "".join(text_parts)
     except Exception:
         return str(header_value)
 
 
-def get_msg_details(imap_conn, uid):
+def decode_message_id(msg_id):
     """
-    Fetches simplified message details (Message-ID, Size, Subject) for a given UID.
-    Returns (msg_id, size, subject) tuple.
+    Decodes a Message-ID header value by unfolding continuation lines.
+    Returns the stripped Message-ID string or None if empty.
     """
+    if not msg_id:
+        return None
+    # Unfold any header continuation lines (CRLF/LF + whitespace)
+    return re.sub(r"\r?\n[ \t]+", " ", str(msg_id)).strip() or None
+
+
+def extract_message_id(header_data):
+    """
+    Extracts the Message-ID from header bytes or string using BytesParser.
+    Returns the stripped Message-ID string or None if not found.
+    """
+    if not header_data:
+        return None
+
     try:
-        resp, data = imap_conn.uid("fetch", uid, "(RFC822.SIZE BODY.PEEK[HEADER.FIELDS (MESSAGE-ID SUBJECT)])")
+        # Use compat32 to preserve raw header with continuation lines
+        # (policy.default's _MessageIDHeader parser truncates at newlines)
+        parser = BytesParser(policy=policy.compat32)
+        if isinstance(header_data, str):
+            header_data = header_data.encode("utf-8", errors="ignore")
+
+        msg = parser.parsebytes(header_data, headersonly=True)
+        return decode_message_id(msg.get("Message-ID"))
     except Exception:
-        return None, None, None
-
-    if resp != "OK":
-        return None, None, None
-
-    msg_id = None
-    subject = "(No Subject)"
-    size = 0
-
-    for item in data:
-        if isinstance(item, tuple):
-            content = item[0].decode("utf-8", errors="ignore")
-
-            # Parse Size
-            size_match = re.search(r"RFC822\.SIZE\s+(\d+)", content)
-            if size_match:
-                size = int(size_match.group(1))
-
-            # Parse Headers
-            msg_bytes = item[1]
-            parser = BytesParser()
-            email_obj = parser.parsebytes(msg_bytes)
-            msg_id = email_obj.get("Message-ID")
-            raw_subject = email_obj.get("Subject")
-            if raw_subject:
-                subject = decode_mime_header(raw_subject)
-
-    return msg_id, size, subject
+        pass
+    return None
 
 
-def message_exists_in_folder(dest_conn, msg_id, src_size):
+def parse_message_id_from_bytes(raw_message):
+    """Parse Message-ID from RFC822 bytes.
+
+    Parses the full message bytes to extract the Message-ID header.
     """
-    Checks if a message with the given Message-ID and RFC822.SIZE exists in the CURRENTLY SELECTED folder of dest_conn.
+    if not raw_message:
+        return None
+    try:
+        parser = BytesParser()
+        msg = parser.parsebytes(raw_message)
+        return decode_message_id(msg.get("Message-ID"))
+    except Exception:
+        return None
+
+
+def parse_message_id_and_subject_from_bytes(raw_message):
+    """Parse Message-ID and decoded Subject from RFC822 bytes.
+
+    Uses header-only parsing to avoid walking large message bodies.
+    Returns a tuple: (message_id, subject).
+    """
+    if not raw_message:
+        return None, "(No Subject)"
+
+    try:
+        # Use compat32 to preserve raw headers with continuation lines
+        parser = BytesParser(policy=policy.compat32)
+        email_obj = parser.parsebytes(raw_message, headersonly=True)
+        msg_id = decode_message_id(email_obj.get("Message-ID"))
+        raw_subject = email_obj.get("Subject")
+        subject = decode_mime_header(raw_subject) if raw_subject else "(No Subject)"
+        return msg_id, subject
+    except Exception:
+        return None, "(No Subject)"
+
+
+def message_exists_in_folder(dest_conn, msg_id):
+    """
+    Checks if a message with the given Message-ID exists in the CURRENTLY SELECTED folder of dest_conn.
     Returns True if found, False otherwise.
     """
     if not msg_id:
@@ -143,23 +317,66 @@ def message_exists_in_folder(dest_conn, msg_id, src_size):
             return False
 
         dest_ids = data[0].split()
-        if not dest_ids:
-            return False
-
-        for did in dest_ids:
-            resp, items = dest_conn.fetch(did, "(RFC822.SIZE)")
-            if resp == "OK":
-                for item in items:
-                    if isinstance(item, bytes):
-                        content = item.decode("utf-8", errors="ignore")
-                    else:
-                        content = item[0].decode("utf-8", errors="ignore")
-                    size_match = re.search(r"RFC822\.SIZE\s+(\d+)", content)
-                    if size_match and int(size_match.group(1)) == src_size:
-                        return True
+        return len(dest_ids) > 0
     except Exception:
         return False
-    return False
+
+
+def get_message_ids_in_folder(imap_conn):
+    """
+    Fetches all Message-IDs from the currently selected folder.
+    Returns a dict mapping UID (bytes) -> Message-ID (str).
+    UIDs without a Message-ID are not included in the result.
+    Use set(result.values()) to get just the Message-ID set.
+    """
+    try:
+        resp, data = imap_conn.uid("search", None, "ALL")
+        if resp != "OK" or not data[0].strip():
+            return {}
+    except Exception:
+        return {}
+
+    uids = data[0].split()
+    return get_uid_to_message_id_map(imap_conn, uids)
+
+
+def get_uid_to_message_id_map(imap_conn, uids):
+    """
+    Fetches Message-IDs for a list of UIDs from the currently selected folder.
+    Returns a dict mapping UID (bytes) -> Message-ID (str).
+    UIDs without a Message-ID are not included in the result.
+    """
+    uid_to_msgid = {}
+    if not uids:
+        return uid_to_msgid
+
+    FETCH_BATCH = 500
+    for i in range(0, len(uids), FETCH_BATCH):
+        batch = uids[i : i + FETCH_BATCH]
+        uid_range = b",".join(batch)
+        try:
+            resp, fetch_data = imap_conn.uid("fetch", uid_range, "(UID BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
+            if resp != "OK":
+                continue
+            for item in fetch_data:
+                if isinstance(item, tuple) and len(item) >= 2:
+                    # Extract UID from response like b'1 (UID 12345 BODY[HEADER.FIELDS ...]'
+                    meta = item[0]
+                    if isinstance(meta, bytes):
+                        meta = meta.decode("utf-8", errors="ignore")
+                    uid_match = re.search(r"UID\s+(\d+)", meta)
+                    if not uid_match:
+                        continue
+                    uid = uid_match.group(1).encode()
+
+                    # Extract Message-ID
+                    mid = extract_message_id(item[1])
+                    if mid:
+                        uid_to_msgid[uid] = mid
+        except Exception:
+            continue
+
+    return uid_to_msgid
 
 
 def sanitize_filename(filename):

src/imap_oauth2.py

@@ -0,0 +1,258 @@
+"""
+IMAP OAuth2 Authentication
+
+OAuth2 token acquisition and refresh for Microsoft and Google IMAP providers.
+Supports device code flow (Microsoft) and installed app flow (Google),
+with in-memory caching for silent token refresh.
+
+Notes:
+- Microsoft OAuth2 requires the `msal` package and uses device code flow.
+- Google OAuth2 requires the `google-auth-oauthlib` package and uses an installed-app flow
+    (opens a browser and listens on a local HTTP redirect).
+- Provider is auto-detected from the IMAP host string.
+"""
+
+import http.client
+import json
+import re
+import ssl
+import sys
+import threading
+import urllib.parse
+
+# Module-level caches for OAuth2 token refresh
+_msal_app_cache = {}  # (client_id, tenant_id) -> PublicClientApplication
+_google_creds_cache = {}  # (client_id, client_secret) -> credentials
+_token_refresh_lock = threading.Lock()
+
+
+def _fetch_json_https(host, path, timeout=10):
+    if not host or any(ch in host for ch in "\r\n"):
+        raise ValueError("Invalid host")
+    if not path.startswith("/"):
+        path = f"/{path}"
+
+    context = ssl.create_default_context()
+    conn = http.client.HTTPSConnection(host, timeout=timeout, context=context)
+    try:
+        conn.request("GET", path, headers={"Accept": "application/json"})
+        response = conn.getresponse()
+        body = response.read()
+    finally:
+        conn.close()
+
+    if response.status != 200:
+        raise RuntimeError(f"Unexpected HTTP status {response.status}")
+    return json.loads(body.decode("utf-8"))
+
+
+def detect_oauth2_provider(host):
+    """
+    Detects the OAuth2 provider from the IMAP host.
+    Returns "microsoft", "google", or None if unrecognized.
+    """
+    host_lower = host.lower()
+    if "outlook" in host_lower or "office365" in host_lower or "microsoft" in host_lower:
+        return "microsoft"
+    if "gmail" in host_lower or "google" in host_lower:
+        return "google"
+    return None
+
+
+def discover_microsoft_tenant(email):
+    """
+    Auto-discovers the Microsoft tenant ID from an email address domain.
+    Uses the OpenID Connect discovery endpoint (no authentication required).
+    Returns the tenant ID string or None if discovery fails.
+    """
+    domain = email.split("@")[-1].strip()
+    if not domain:
+        print("Error: Could not discover Microsoft tenant: missing email domain")
+        return None
+
+    domain_quoted = urllib.parse.quote(domain, safe=".-")
+    path = f"/{domain_quoted}/.well-known/openid-configuration"
+
+    try:
+        data = _fetch_json_https("login.microsoftonline.com", path, timeout=10)
+    except (OSError, http.client.HTTPException, json.JSONDecodeError, RuntimeError, ValueError) as e:
+        print(f"Error: Could not discover Microsoft tenant for domain '{domain}': {e}")
+        return None
+
+    issuer = data.get("issuer", "")
+    match = re.search(r"/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})", issuer)
+    if match:
+        return match.group(1)
+
+    print(f"Error: Could not extract tenant ID from issuer: {issuer}")
+    return None
+
+
+def acquire_microsoft_oauth2_token(client_id, email):
+    """
+    Acquires a Microsoft OAuth2 access token using the MSAL device code flow.
+    Auto-discovers tenant ID from the email domain.
+    Requires the 'msal' package: pip install msal
+
+    On subsequent calls, silently refreshes the token using the cached MSAL app
+    (which holds the refresh token in its in-memory cache).
+    """
+    tenant_id = discover_microsoft_tenant(email)
+    if not tenant_id:
+        return None
+
+    try:
+        import msal
+    except ImportError:
+        print("Error: 'msal' package is required for Microsoft OAuth2. Install it with: pip install msal")
+        sys.exit(1)
+
+    authority = f"https://login.microsoftonline.com/{tenant_id}"
+    scopes = ["https://outlook.office365.com/IMAP.AccessAsUser.All"]
+
+    # Reuse cached MSAL app so acquire_token_silent can access refresh tokens
+    cache_key = (client_id, tenant_id)
+    if cache_key in _msal_app_cache:
+        app = _msal_app_cache[cache_key]
+    else:
+        print(f"Discovered Microsoft tenant: {tenant_id}")
+        app = msal.PublicClientApplication(client_id, authority=authority)
+        _msal_app_cache[cache_key] = app
+
+    # Try cached/refreshed token first (handles refresh tokens automatically)
+    accounts = app.get_accounts()
+    if accounts:
+        result = app.acquire_token_silent(scopes, account=accounts[0])
+        if result and "access_token" in result:
+            return result["access_token"]
+
+    # Fall back to device code flow (first call or if refresh fails)
+    flow = app.initiate_device_flow(scopes=scopes)
+    if "user_code" not in flow:
+        print(f"Error: Could not initiate device flow: {flow.get('error_description', 'Unknown error')}")
+        return None
+
+    print(flow["message"])
+    result = app.acquire_token_by_device_flow(flow)
+
+    if "access_token" in result:
+        return result["access_token"]
+
+    print(f"Error: Could not acquire token: {result.get('error_description', 'Unknown error')}")
+    return None
+
+
+def acquire_google_oauth2_token(client_id, client_secret):
+    """
+    Acquires a Google OAuth2 access token using the installed app flow.
+    Opens a browser for user consent and runs a local HTTP server for the redirect.
+    Requires the 'google-auth-oauthlib' package: pip install google-auth-oauthlib
+
+    On subsequent calls, silently refreshes the token using the cached credentials
+    object (which holds the refresh token). No browser interaction needed for refresh.
+    """
+    # Try refreshing cached credentials first (no browser needed)
+    cache_key = (client_id, client_secret)
+    if cache_key in _google_creds_cache:
+        creds = _google_creds_cache[cache_key]
+        if creds and creds.refresh_token:
+            try:
+                import google.auth.transport.requests
+
+                creds.refresh(google.auth.transport.requests.Request())
+                if creds.token:
+                    return creds.token
+            except Exception:
+                pass  # Fall through to full auth flow
+
+    try:
+        from google_auth_oauthlib.flow import InstalledAppFlow
+    except ImportError:
+        print("Error: 'google-auth-oauthlib' package is required for Google OAuth2.")
+        print("Install it with: pip install google-auth-oauthlib")
+        sys.exit(1)
+
+    client_config = {
+        "installed": {
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": "https://oauth2.googleapis.com/token",
+            "redirect_uris": ["http://localhost"],
+        }
+    }
+
+    flow = InstalledAppFlow.from_client_config(client_config, scopes=["https://mail.google.com/"])
+
+    print("Opening browser for Google authentication...")
+    print("If the browser does not open, check the terminal for a URL to visit.")
+
+    credentials = flow.run_local_server(port=0)
+
+    if credentials and credentials.token:
+        _google_creds_cache[cache_key] = credentials
+        return credentials.token
+
+    print("Error: Could not acquire Google OAuth2 token.")
+    return None
+
+
+def acquire_oauth2_token_for_provider(provider, client_id, email, client_secret=None):
+    """
+    Acquires an OAuth2 token for the specified provider.
+
+    Args:
+        provider: "microsoft" or "google"
+        client_id: OAuth2 client ID
+        email: User's email address (used for Microsoft tenant discovery)
+        client_secret: Required for Google, not needed for Microsoft
+    """
+    if provider == "microsoft":
+        return acquire_microsoft_oauth2_token(client_id, email)
+    elif provider == "google":
+        if not client_secret:
+            print(
+                "Error: OAuth2 client secret is required for Google OAuth2. "
+                "Provide --oauth2-client-secret (or --src-oauth2-client-secret / --dest-oauth2-client-secret), "
+                "or set OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET / DEST_OAUTH2_CLIENT_SECRET."
+            )
+            return None
+        return acquire_google_oauth2_token(client_id, client_secret)
+    else:
+        print(f"Error: Unknown OAuth2 provider: {provider}")
+        return None
+
+
+def refresh_oauth2_token(conf, old_token):
+    """
+    Thread-safe OAuth2 token refresh using double-checked locking.
+
+    Multiple threads may detect an expired token simultaneously. This function
+    ensures only one thread performs the actual refresh. Other threads waiting
+    on the lock will see that conf["oauth2_token"] has already been updated and
+    skip the redundant refresh.
+
+    Args:
+        conf: Mutable dict with keys:
+            - host, user, password, oauth2_token: connection credentials
+            - oauth2: dict with provider, client_id, email, client_secret
+        old_token: The expired token that triggered this refresh (for comparison)
+
+    Returns:
+        The new token string, or None if refresh failed.
+    """
+    oauth2 = conf.get("oauth2")
+    if not oauth2:
+        return None
+
+    with _token_refresh_lock:
+        # Another thread may have already refreshed while we were waiting
+        if conf["oauth2_token"] != old_token:
+            return conf["oauth2_token"]
+
+        new_token = acquire_oauth2_token_for_provider(
+            oauth2["provider"], oauth2["client_id"], oauth2["email"], oauth2.get("client_secret")
+        )
+        if new_token:
+            conf["oauth2_token"] = new_token
+        return new_token

src/restore_cache.py

@@ -0,0 +1,254 @@
+"""Restore progress cache.
+
+This module supports faster incremental restores by persisting, per destination+folder,
+the set of Message-IDs already seen/processed by this tool.
+
+The caller decides where the cache file lives by passing a cache directory/root.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import os
+import re
+import threading
+import time
+from collections.abc import Callable
+
+RESTORE_CACHE_VERSION = 1
+
+# Throttle disk writes so we can update frequently without rewriting a large JSON file
+# on every single message.
+_MIN_SECONDS_BETWEEN_SAVES = 2.0
+_MIN_PENDING_UPDATES_BEFORE_SAVE = 50
+
+
+def _safe_cache_component(value: str) -> str:
+    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value or "")
+
+
+def _prepare_cache_for_json(cache_data: dict) -> dict:
+    """Create a deep copy of cache_data suitable for JSON serialization.
+    Removes in-memory helper fields like _ids_set that should not be persisted.
+    """
+    snapshot = copy.deepcopy(cache_data)
+    folders = snapshot.get("folders", {})
+    if isinstance(folders, dict):
+        for folder_entry in folders.values():
+            if isinstance(folder_entry, dict):
+                # Remove the in-memory set cache
+                folder_entry.pop("_ids_set", None)
+    return snapshot
+
+
+def get_dest_index_cache_path(cache_root: str, dest_host: str, dest_user: str) -> str:
+    safe_host = _safe_cache_component(dest_host)
+    safe_user = _safe_cache_component(dest_user)
+    return os.path.join(cache_root, f"restore_cache_{safe_host}_{safe_user}.json")
+
+
+def load_dest_index_cache(cache_path: str) -> dict:
+    try:
+        with open(cache_path, encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, dict):
+            return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
+        if data.get("version") != RESTORE_CACHE_VERSION:
+            return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
+        if not isinstance(data.get("folders"), dict):
+            data["folders"] = {}
+        if not isinstance(data.get("_meta"), dict):
+            data["_meta"] = {}
+        return data
+    except FileNotFoundError:
+        return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
+    except Exception:
+        return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}}
+
+
+def save_dest_index_cache(cache_path: str, cache_data: dict) -> None:
+    try:
+        tmp_path = f"{cache_path}.tmp"
+        with open(tmp_path, "w", encoding="utf-8") as f:
+            json.dump(cache_data, f, ensure_ascii=False)
+        os.replace(tmp_path, cache_path)
+        return True
+    except Exception:
+        # Cache is best-effort; do not fail restore.
+        return False
+
+
+def _ensure_dest(cache_data: dict, dest_host: str, dest_user: str) -> None:
+    cache_dest = cache_data.get("dest")
+    if not isinstance(cache_dest, dict) or cache_dest.get("host") != dest_host or cache_dest.get("user") != dest_user:
+        cache_data.clear()
+        cache_data.update(
+            {
+                "version": RESTORE_CACHE_VERSION,
+                "dest": {"host": dest_host, "user": dest_user},
+                "folders": {},
+                "_meta": {},
+            }
+        )
+
+
+def get_cached_message_ids(
+    cache_data: dict,
+    cache_lock: threading.Lock,
+    dest_host: str,
+    dest_user: str,
+    folder_name: str,
+) -> set[str]:
+    """Return Message-IDs we've already seen/processed for this destination+folder."""
+    with cache_lock:
+        _ensure_dest(cache_data, dest_host, dest_user)
+        folders = cache_data.setdefault("folders", {})
+        entry = folders.get(folder_name)
+        if not isinstance(entry, dict):
+            return set()
+        ids = entry.get("message_ids")
+        if not isinstance(ids, list):
+            return set()
+        return {str(x) for x in ids if x}
+
+
+def add_cached_message_id(
+    cache_data: dict,
+    cache_lock: threading.Lock,
+    dest_host: str,
+    dest_user: str,
+    folder_name: str,
+    message_id: str,
+) -> bool:
+    """Add a Message-ID to the cache. Returns True if it was newly added."""
+    if not message_id:
+        return False
+    msg_id = str(message_id).strip()
+    if not msg_id:
+        return False
+
+    with cache_lock:
+        _ensure_dest(cache_data, dest_host, dest_user)
+        folders = cache_data.setdefault("folders", {})
+        entry = folders.setdefault(folder_name, {})
+        if not isinstance(entry, dict):
+            folders[folder_name] = {}
+            entry = folders[folder_name]
+
+        ids = entry.get("message_ids")
+        if not isinstance(ids, list):
+            ids = []
+            entry["message_ids"] = ids
+
+        # Maintain entry-level in-memory set for O(1) lookups (not persisted to JSON)
+        ids_set = entry.get("_ids_set")
+        if ids_set is None:
+            ids_set = set(ids)
+            entry["_ids_set"] = ids_set
+
+        if msg_id in ids_set:
+            return False
+
+        ids.append(msg_id)
+        ids_set.add(msg_id)
+
+        meta = cache_data.setdefault("_meta", {})
+        if not isinstance(meta, dict):
+            meta = {}
+            cache_data["_meta"] = meta
+        meta["pending_updates"] = int(meta.get("pending_updates") or 0) + 1
+        return True
+
+
+def maybe_save_dest_index_cache(
+    cache_path: str,
+    cache_data: dict,
+    cache_lock: threading.Lock,
+    *,
+    force: bool = False,
+    log_fn: Callable[[str], None] | None = None,
+) -> bool:
+    """Persist cache to disk if enough updates/time has accumulated."""
+    now = time.time()
+    pending = 0
+    with cache_lock:
+        meta = cache_data.setdefault("_meta", {})
+        if not isinstance(meta, dict):
+            meta = {}
+            cache_data["_meta"] = meta
+
+        pending = int(meta.get("pending_updates") or 0)
+        last_saved = float(meta.get("last_saved_ts") or 0.0)
+
+        should_save = force or (
+            pending > 0
+            and (pending >= _MIN_PENDING_UPDATES_BEFORE_SAVE or (now - last_saved) >= _MIN_SECONDS_BETWEEN_SAVES)
+        )
+        if not should_save:
+            return False
+
+        # Update meta before writing so the on-disk file reflects the flush decision.
+        meta["pending_updates"] = 0
+        meta["last_saved_ts"] = now
+
+        # Prepare a clean snapshot for JSON serialization (removes in-memory helper fields)
+        snapshot = _prepare_cache_for_json(cache_data)
+
+    did_write = save_dest_index_cache(cache_path, snapshot)
+    if did_write and log_fn is not None:
+        log_fn(f"Wrote restore cache ({pending} updates): {cache_path}")
+    return did_write
+
+
+def record_progress(
+    *,
+    message_id: str | None,
+    folder_name: str,
+    existing_dest_msg_ids: set[str] | None,
+    existing_dest_msg_ids_lock: threading.Lock | None,
+    progress_cache_path: str | None,
+    progress_cache_data: dict | None,
+    progress_cache_lock: threading.Lock | None,
+    dest_host: str | None,
+    dest_user: str | None,
+    log_fn: Callable[[str], None] | None = None,
+) -> None:
+    """Record a processed Message-ID for fast skipping on future incremental runs.
+
+    Updates both:
+    - the in-memory set used by the current run, and
+    - the persisted progress cache on disk (throttled writes).
+    """
+    if not message_id:
+        return
+
+    msg_id = str(message_id).strip()
+    if not msg_id:
+        return
+
+    if existing_dest_msg_ids is not None and existing_dest_msg_ids_lock is not None:
+        with existing_dest_msg_ids_lock:
+            existing_dest_msg_ids.add(msg_id)
+
+    if (
+        progress_cache_path
+        and progress_cache_data is not None
+        and progress_cache_lock is not None
+        and dest_host
+        and dest_user
+    ):
+        add_cached_message_id(
+            progress_cache_data,
+            progress_cache_lock,
+            dest_host,
+            dest_user,
+            folder_name,
+            msg_id,
+        )
+        maybe_save_dest_index_cache(
+            progress_cache_path,
+            progress_cache_data,
+            progress_cache_lock,
+            log_fn=log_fn,
+        )

src/restore_imap_emails.py

@@ -7,12 +7,18 @@ Reads .eml files from a local directory and uploads them to the destination serv
 Features:
 - Folder Restoration: Recreates the folder structure from the backup.
 - Gmail Labels Restoration: Uses labels_manifest.json to apply Gmail labels.
-- Incremental Restore: Skips emails that already exist (based on Message-ID and size).
+- Incremental Restore: Skips emails that already exist (based on Message-ID).
 - Parallel Processing: Uses multithreading for fast uploads.
 - Date Preservation: Restores emails with their original dates.
 
 Configuration (Environment Variables):
-  DEST_IMAP_HOST, DEST_IMAP_USERNAME, DEST_IMAP_PASSWORD: Destination credentials.
+    DEST_IMAP_HOST, DEST_IMAP_USERNAME: Destination credentials.
+    DEST_IMAP_PASSWORD: Destination password (or App Password).
+
+    OAuth2 (Optional - instead of password):
+    DEST_OAUTH2_CLIENT_ID: OAuth2 Client ID
+    DEST_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
+
   BACKUP_LOCAL_PATH: Source local directory containing the backup.
   MAX_WORKERS: Number of concurrent threads (default: 4).
   BATCH_SIZE: Number of emails to process per batch (default: 10).
@@ -23,10 +29,19 @@ Configuration (Environment Variables):
               Default is "false".
 
 Usage:
-  python3 restore_imap_emails.py --src-path "./my_backup" --dest-host "imap.gmail.com"
+    python3 restore_imap_emails.py \
+        --src-path "./my_backup" \
+        --dest-host "imap.gmail.com" \
+        --dest-user "you@gmail.com" \
+        --dest-pass "your-app-password"
 
 Gmail Labels Restoration:
-  python3 restore_imap_emails.py --src-path "./gmail_backup" --dest-host "imap.gmail.com" --apply-labels
+    python3 restore_imap_emails.py \
+        --src-path "./gmail_backup" \
+        --dest-host "imap.gmail.com" \
+        --dest-user "you@gmail.com" \
+        --dest-pass "your-app-password" \
+        --apply-labels
   This uploads emails and applies labels from labels_manifest.json to recreate
   the original Gmail label structure.
 """
@@ -41,8 +56,11 @@ import time
 from email import policy
 from email.parser import BytesParser
 from email.utils import parsedate_to_datetime
+from typing import Optional
 
 import imap_common
+import imap_oauth2
+import restore_cache
 
 # Defaults
 MAX_WORKERS = 4  # Lower default for restore to avoid rate limits
@@ -63,12 +81,17 @@ def safe_print(message):
 def get_thread_connection(dest_conf):
     """Get or create a thread-local IMAP connection."""
     if not hasattr(thread_local, "dest") or thread_local.dest is None:
-        thread_local.dest = imap_common.get_imap_connection(*dest_conf)
+        thread_local.dest = imap_common.get_imap_connection_from_conf(dest_conf)
     try:
         if thread_local.dest:
             thread_local.dest.noop()
     except Exception:
-        thread_local.dest = imap_common.get_imap_connection(*dest_conf)
+        thread_local.dest = imap_common.get_imap_connection_from_conf(dest_conf)
+        # If reconnection failed (possibly expired token), try refreshing
+        if thread_local.dest is None and dest_conf.get("oauth2"):
+            old_token = dest_conf["oauth2_token"]
+            imap_oauth2.refresh_oauth2_token(dest_conf, old_token)
+            thread_local.dest = imap_common.get_imap_connection_from_conf(dest_conf)
     return thread_local.dest
 
 
@@ -185,6 +208,22 @@ def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size):
         safe_print(f"  -> Error syncing flags: {e}")
 
 
+def extract_message_id_from_eml(file_path):
+    """
+    Extract just the Message-ID from an .eml file efficiently.
+    Returns the Message-ID string or None on error.
+    """
+    try:
+        with open(file_path, "rb") as f:
+            # Read just the first 64KB to get headers
+            header_bytes = f.read(65536)
+
+        msg_id = imap_common.extract_message_id(header_bytes)
+        return msg_id
+    except Exception:
+        return None
+
+
 def parse_eml_file(file_path):
     """
     Parse an .eml file and extract metadata.
@@ -194,11 +233,13 @@ def parse_eml_file(file_path):
         with open(file_path, "rb") as f:
             raw_content = f.read()
 
-        parser = BytesParser(policy=policy.default)
-        msg = parser.parsebytes(raw_content)
+        # Use compat32 to preserve raw headers with continuation lines
+        parser = BytesParser(policy=policy.compat32)
+        msg = parser.parsebytes(raw_content, headersonly=True)
 
-        message_id = msg.get("Message-ID", "").strip()
-        subject = msg.get("Subject", "(No Subject)")
+        message_id = imap_common.decode_message_id(msg.get("Message-ID"))
+        raw_subject = msg.get("Subject")
+        subject = imap_common.decode_mime_header(raw_subject) if raw_subject else "(No Subject)"
         date_header = msg.get("Date")
 
         # Parse date for IMAP INTERNALDATE
@@ -238,46 +279,48 @@ def get_eml_files(folder_path):
     return eml_files
 
 
-def email_exists_in_folder(imap_conn, message_id, size):
+def email_exists_in_folder(imap_conn, message_id):
     """
-    Check if an email with the given Message-ID and size exists in the currently selected folder.
+    Check if an email with the given Message-ID exists in the currently selected folder.
     """
     if not message_id:
         return False
 
     try:
-        return imap_common.message_exists_in_folder(imap_conn, message_id, size)
+        return imap_common.message_exists_in_folder(imap_conn, message_id)
     except Exception:
         return False
 
 
-def upload_email(dest, folder_name, raw_content, date_str, message_id, subject, flags=None):
+def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True):
     """
     Upload a single email to the destination folder.
-    Returns True on success, False on failure.
+    Returns True on success, False if duplicate or on failure.
 
     Args:
         flags: Optional string of IMAP flags like "\\Seen" for read emails.
+        check_duplicate: Whether to check for duplicates before uploading.
     """
     try:
         # Ensure folder exists
-        if folder_name.upper() != "INBOX":
-            try:
-                dest.create(f'"{folder_name}"')
-            except Exception:
-                pass  # Folder may already exist
+        imap_common.ensure_folder_exists(dest, folder_name)
 
         # Select folder
         dest.select(f'"{folder_name}"')
 
-        # Check for duplicates
-        size = len(raw_content)
-        if message_id and email_exists_in_folder(dest, message_id, size):
+        # Check for duplicates if requested
+        if check_duplicate and message_id and email_exists_in_folder(dest, message_id):
             return False  # Already exists
 
         # Upload with original date and flags
-        resp, _ = dest.append(f'"{folder_name}"', flags, date_str, raw_content)
-        return resp == "OK"
+        return imap_common.append_email(
+            dest,
+            folder_name,
+            raw_content,
+            date_str,
+            flags,
+            ensure_folder=False,
+        )
 
     except Exception as e:
         safe_print(f"Error uploading to {folder_name}: {e}")
@@ -304,15 +347,30 @@ def label_to_folder(label):
     """
     Convert a Gmail label name to an IMAP folder path.
     """
-    if label == "INBOX":
-        return "INBOX"
+    if label == imap_common.FOLDER_INBOX:
+        return imap_common.FOLDER_INBOX
     elif label in ("Sent Mail", "Starred", "Drafts", "Important"):
         return f"[Gmail]/{label}"
     else:
         return label
 
 
-def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_labels, apply_flags):
+def process_restore_batch(
+    eml_files,
+    folder_name,
+    dest_conf,
+    manifest,
+    apply_labels,
+    apply_flags,
+    full_restore: bool = False,
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+    progress_cache_path: Optional[str] = None,
+    dest_host: Optional[str] = None,
+    dest_user: Optional[str] = None,
+):
     """
     Process a batch of .eml files for restoration.
 
@@ -367,21 +425,70 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab
                     else:
                         remaining_labels.append(label)
 
-                # If no valid label found, use Drafts as fallback (won't appear in INBOX)
+                # If no valid label found, use a dedicated label folder as fallback.
+                # This avoids placing non-draft messages into Gmail Drafts.
                 if target_folder is None:
-                    target_folder = "[Gmail]/Drafts"
+                    target_folder = imap_common.FOLDER_RESTORED_UNLABELED
                     remaining_labels = []
             else:
                 target_folder = folder_name
                 remaining_labels = labels
 
-            # Upload to target folder (or check if exists)
-            uploaded = upload_email(dest, target_folder, raw_content, date_str, message_id, display_subject, flags)
+            existing_dest_msg_ids: Optional[set[str]] = None
+            if existing_dest_msg_ids_by_folder is not None:
+                if existing_dest_msg_ids_lock is None:
+                    existing_dest_msg_ids_lock = threading.Lock()
+
+                with existing_dest_msg_ids_lock:
+                    existing_dest_msg_ids = existing_dest_msg_ids_by_folder.get(target_folder)
+
+                # Lazy-load per-folder progress cache (no destination scan).
+                if existing_dest_msg_ids is None:
+                    built: set[str] = set()
+                    if progress_cache_data is not None and progress_cache_lock is not None and dest_host and dest_user:
+                        built = restore_cache.get_cached_message_ids(
+                            progress_cache_data,
+                            progress_cache_lock,
+                            dest_host,
+                            dest_user,
+                            target_folder,
+                        )
+                    with existing_dest_msg_ids_lock:
+                        existing_dest_msg_ids_by_folder.setdefault(target_folder, built)
+                        existing_dest_msg_ids = existing_dest_msg_ids_by_folder[target_folder]
+
+            # Incremental default: if we already know we processed it before, skip entirely.
+            if (
+                not full_restore
+                and message_id
+                and existing_dest_msg_ids is not None
+                and message_id in existing_dest_msg_ids
+            ):
+                safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
+                continue
+
+            # Upload to target folder.
+            # Keep server-side duplicate checks enabled to avoid creating duplicates for emails
+            # that exist on the destination but are not in our local progress cache.
+            uploaded = upload_email(dest, target_folder, raw_content, date_str, message_id, flags)
+
+            restore_cache.record_progress(
+                message_id=message_id,
+                folder_name=target_folder,
+                existing_dest_msg_ids=existing_dest_msg_ids,
+                existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
+                progress_cache_path=progress_cache_path,
+                progress_cache_data=progress_cache_data,
+                progress_cache_lock=progress_cache_lock,
+                dest_host=dest_host,
+                dest_user=dest_user,
+                log_fn=safe_print,
+            )
 
             if not uploaded:
                 safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}")
-                # Even if skipped, sync flags on existing email if requested
-                if apply_flags and flags and message_id:
+                # Full restore preserves legacy behavior: sync flags on existing email if requested
+                if full_restore and apply_flags and flags and message_id:
                     sync_flags_on_existing(dest, target_folder, message_id, flags, size)
             else:
                 safe_print(f"[{target_folder}] UPLOADED      | {size_str:<8} | {display_subject}")
@@ -390,9 +497,10 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab
                     for flag in flags.split():
                         safe_print(f"  -> Applied flag: {flag}")
 
-            # Apply remaining Gmail labels (always, whether uploaded or skipped)
-            # This ensures labels are synced even for existing emails
-            if apply_labels and remaining_labels:
+            # Apply remaining Gmail labels:
+            # - Full restore: apply/sync labels even for existing emails
+            # - Incremental (default): apply labels only for newly uploaded emails
+            if apply_labels and remaining_labels and (uploaded or full_restore):
                 for label in remaining_labels:
                     label_folder = label_to_folder(label)
 
@@ -401,24 +509,42 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab
                         continue
 
                     # Skip system folders we can't upload to
-                    if label_folder in ("[Gmail]/All Mail", "[Gmail]/Spam", "[Gmail]/Trash"):
+                    if label_folder in (
+                        imap_common.GMAIL_ALL_MAIL,
+                        imap_common.GMAIL_SPAM,
+                        imap_common.GMAIL_TRASH,
+                    ):
                         continue
 
                     try:
                         # Ensure label folder exists
-                        if label_folder.upper() != "INBOX":
-                            try:
-                                dest.create(f'"{label_folder}"')
-                            except Exception:
-                                pass
+                        imap_common.ensure_folder_exists(dest, label_folder)
 
                         # Select and check for duplicate
                         dest.select(f'"{label_folder}"')
-                        if not email_exists_in_folder(dest, message_id, size):
-                            dest.append(f'"{label_folder}"', flags, date_str, raw_content)
+                        if not email_exists_in_folder(dest, message_id):
+                            imap_common.append_email(
+                                dest,
+                                label_folder,
+                                raw_content,
+                                date_str,
+                                flags,
+                                ensure_folder=False,
+                            )
+                            restore_cache.record_progress(
+                                label_folder,
+                                message_id,
+                                existing_dest_msg_ids=existing_dest_msg_ids,
+                                existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
+                                progress_cache_path=progress_cache_path,
+                                progress_cache_data=progress_cache_data,
+                                progress_cache_lock=progress_cache_lock,
+                                dest_host=dest_host,
+                                dest_user=dest_user,
+                            )
                             safe_print(f"  -> Applied label: {label}")
-                        # If email exists in this label folder, sync flags
-                        elif apply_flags and flags:
+                        # If email exists in this label folder, sync flags (full restore only)
+                        elif full_restore and apply_flags and flags:
                             sync_flags_on_existing(dest, label_folder, message_id, flags, size)
                     except Exception as e:
                         safe_print(f"  -> Error applying label {label}: {e}")
@@ -483,14 +609,7 @@ def delete_orphan_emails_from_dest(imap_conn, folder_name, local_msg_ids):
                         uid = uid_match.group(1)
 
                         # Extract Message-ID
-                        header_data = item[1]
-                        msg_id = None
-                        if isinstance(header_data, bytes):
-                            header_str = header_data.decode("utf-8", errors="ignore")
-                            for line in header_str.split("\n"):
-                                if line.lower().startswith("message-id:"):
-                                    msg_id = line.split(":", 1)[1].strip()
-                                    break
+                        msg_id = imap_common.extract_message_id(item[1])
 
                         # If not in local backup, mark for deletion
                         if msg_id and msg_id not in local_msg_ids:
@@ -517,7 +636,17 @@ def delete_orphan_emails_from_dest(imap_conn, folder_name, local_msg_ids):
     return deleted_count
 
 
-def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_labels, apply_flags, dest_delete=False):
+def restore_folder(
+    folder_name,
+    local_folder_path,
+    dest_conf,
+    manifest,
+    apply_labels,
+    apply_flags,
+    dest_delete=False,
+    full_restore: bool = False,
+    cache_root: Optional[str] = None,
+):
     """
     Restore all emails from a local folder to the destination IMAP server.
     """
@@ -528,7 +657,7 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
         safe_print(f"No .eml files found in {folder_name}")
         # Even if empty, check for orphans to delete
         if dest_delete:
-            dest = imap_common.get_imap_connection(*dest_conf)
+            dest = imap_common.get_imap_connection_from_conf(dest_conf)
             if dest:
                 delete_orphan_emails_from_dest(dest, folder_name, set())
                 dest.logout()
@@ -536,6 +665,30 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
 
     safe_print(f"Found {len(eml_files)} emails to restore.")
 
+    cache_root = cache_root or local_folder_path
+    cache_path = restore_cache.get_dest_index_cache_path(cache_root, dest_conf["host"], dest_conf["user"])
+    cache_data: dict = restore_cache.load_dest_index_cache(cache_path)
+    cache_lock = threading.Lock()
+
+    safe_print(f"Using progress cache: {cache_path}")
+
+    # Incremental mode uses cached Message-IDs to skip already-processed emails.
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {folder_name: set()}
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
+    try:
+        existing_dest_msg_ids_by_folder[folder_name] = restore_cache.get_cached_message_ids(
+            cache_data,
+            cache_lock,
+            dest_conf["host"],
+            dest_conf["user"],
+            folder_name,
+        )
+        safe_print(f"Cache has {len(existing_dest_msg_ids_by_folder[folder_name])} Message-IDs for this folder.")
+    except Exception as e:
+        # Fall back to an empty cache for this folder if reading cached Message-IDs fails.
+        safe_print(f"Warning: Failed to load cached Message-IDs for folder '{folder_name}': {e}")
+        existing_dest_msg_ids_by_folder[folder_name] = set()
+
     # If dest_delete enabled, get local Message-IDs for comparison
     local_msg_ids = None
     if dest_delete:
@@ -543,8 +696,50 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
         local_msg_ids = get_local_message_ids(local_folder_path)
         safe_print(f"Found {len(local_msg_ids)} unique Message-IDs in local backup.")
 
+    # Pre-fetch destination Message-IDs and filter duplicates (non-Gmail mode only)
+    gmail_mode = folder_name == "__GMAIL_MODE__"
+    files_to_restore = eml_files
+
+    if not gmail_mode:
+        dest_msg_ids = set()
+        try:
+            dest_tmp = imap_common.get_imap_connection_from_conf(dest_conf)
+            if dest_tmp:
+                # Ensure folder exists before selecting
+                if folder_name.upper() != "INBOX":
+                    try:
+                        dest_tmp.create(f'"{folder_name}"')
+                    except Exception:
+                        pass
+                dest_tmp.select(f'"{folder_name}"')
+                dest_msg_ids = set(imap_common.get_message_ids_in_folder(dest_tmp).values())
+                dest_tmp.logout()
+        except Exception:
+            dest_msg_ids = set()
+
+        safe_print(f"{len(dest_msg_ids)} existing messages in destination.")
+
+        # Pre-filter files to skip duplicates
+        if dest_msg_ids:
+            safe_print("Pre-filtering duplicates...")
+            files_to_restore = []
+            skipped = 0
+            for file_path, filename in eml_files:
+                msg_id = extract_message_id_from_eml(file_path)
+                if msg_id and msg_id in dest_msg_ids:
+                    skipped += 1
+                else:
+                    files_to_restore.append((file_path, filename))
+            safe_print(f"Skipping {skipped} duplicates, {len(files_to_restore)} to restore.")
+
+    if not files_to_restore:
+        safe_print("No new emails to restore.")
+        return
+
+    safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...")
+
     # Create batches
-    batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)]
+    batches = [files_to_restore[i : i + BATCH_SIZE] for i in range(0, len(files_to_restore), BATCH_SIZE)]
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
         futures = []
@@ -558,6 +753,14 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
                     manifest,
                     apply_labels,
                     apply_flags,
+                    full_restore,
+                    existing_dest_msg_ids_by_folder,
+                    existing_dest_msg_ids_lock,
+                    cache_data,
+                    cache_lock,
+                    cache_path,
+                    dest_conf.get("host"),
+                    dest_conf.get("user"),
                 )
             )
 
@@ -570,13 +773,16 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la
     # Delete orphan emails from destination if enabled
     if dest_delete and local_msg_ids is not None:
         safe_print("Syncing destination: removing emails not in local backup...")
-        dest = imap_common.get_imap_connection(*dest_conf)
+        dest = imap_common.get_imap_connection_from_conf(dest_conf)
         if dest:
             delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids)
             dest.logout()
 
+    # Force-flush progress cache at end.
+    restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
 
-def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
+
+def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore: bool = False):
     """
     Special restoration mode for Gmail: Upload emails to their first label folder
     and then apply additional labels from the manifest.
@@ -607,6 +813,19 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
 
     batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)]
 
+    cache_path = restore_cache.get_dest_index_cache_path(local_path, dest_conf["host"], dest_conf["user"])
+    progress_cache_data: dict = restore_cache.load_dest_index_cache(cache_path)
+    progress_cache_lock = threading.Lock()
+
+    safe_print(f"Using progress cache: {cache_path}")
+    safe_print(
+        "Cache will be populated as restore runs (no up-front destination indexing). "
+        "First run may still do per-message duplicate checks; subsequent runs will skip quickly."
+    )
+
+    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {}  # lazily loaded per folder
+    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
+
     with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
         futures = []
         for batch in batches:
@@ -622,6 +841,14 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
                     manifest,
                     True,  # apply_labels
                     apply_flags,  # apply_flags
+                    full_restore,
+                    existing_dest_msg_ids_by_folder,
+                    existing_dest_msg_ids_lock,
+                    progress_cache_data,
+                    progress_cache_lock,
+                    cache_path,
+                    dest_conf.get("host"),
+                    dest_conf.get("user"),
                 )
             )
 
@@ -639,6 +866,9 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags):
     elapsed = time.time() - start_time
     safe_print(f"\nRestore completed in {elapsed:.1f}s")
 
+    # Force-flush progress cache at end.
+    restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True)
+
 
 def get_backup_folders(local_path):
     """
@@ -675,23 +905,62 @@ def main():
 
     # Source (Local Path)
     env_path = os.getenv("BACKUP_LOCAL_PATH")
-    parser.add_argument("--src-path", default=env_path, help="Local source path containing backup")
+    parser.add_argument(
+        "--src-path",
+        default=env_path,
+        required=not bool(env_path),
+        help="Local source path containing backup (or BACKUP_LOCAL_PATH)",
+    )
 
     # Destination
+    default_dest_host = os.getenv("DEST_IMAP_HOST")
+    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
+    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
+    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
+
     parser.add_argument(
         "--dest-host",
-        default=os.getenv("DEST_IMAP_HOST"),
-        help="Destination IMAP Server",
+        default=default_dest_host,
+        required=not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
     )
     parser.add_argument(
         "--dest-user",
-        default=os.getenv("DEST_IMAP_USERNAME"),
-        help="Destination Username",
+        default=default_dest_user,
+        required=not bool(default_dest_user),
+        help="Destination Username (or DEST_IMAP_USERNAME)",
+    )
+
+    dest_auth_required = not bool(default_dest_pass or default_dest_client_id)
+    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
+    dest_auth.add_argument(
+        "--dest-pass",
+        default=default_dest_pass,
+        help="Destination Password (or DEST_IMAP_PASSWORD)",
+    )
+    dest_auth.add_argument(
+        "--dest-oauth2-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
+    )
+    dest_auth.add_argument(
+        "--dest-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help=argparse.SUPPRESS,
     )
     parser.add_argument(
-        "--dest-pass",
-        default=os.getenv("DEST_IMAP_PASSWORD"),
-        help="Destination Password",
+        "--dest-oauth2-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--dest-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help=argparse.SUPPRESS,
     )
 
     # Config
@@ -731,6 +1000,14 @@ def main():
         help="Gmail restore mode: Upload to INBOX and apply labels + flags from manifest",
     )
 
+    env_full_restore = os.getenv("FULL_RESTORE", "false").lower() == "true"
+    parser.add_argument(
+        "--full-restore",
+        action="store_true",
+        default=env_full_restore,
+        help="Force full restore (legacy): process all emails and sync labels/flags for already-present messages.",
+    )
+
     # Sync mode: delete from dest emails not in local backup
     env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
     parser.add_argument(
@@ -745,29 +1022,44 @@ def main():
 
     args = parser.parse_args()
 
-    # Validate
-    missing = []
-    if not args.dest_host:
-        missing.append("DEST_IMAP_HOST")
-    if not args.dest_user:
-        missing.append("DEST_IMAP_USERNAME")
-    if not args.dest_pass:
-        missing.append("DEST_IMAP_PASSWORD")
-
-    if missing:
-        print(f"Error: Missing credentials: {', '.join(missing)}")
-        sys.exit(1)
-
-    if not args.src_path:
-        print("Error: Source path is required.")
-        print("Please provide --src-path or set environment variable BACKUP_LOCAL_PATH.")
-        sys.exit(1)
+    dest_use_oauth2 = bool(args.dest_client_id)
 
     global MAX_WORKERS, BATCH_SIZE
     MAX_WORKERS = args.workers
     BATCH_SIZE = args.batch
 
-    dest_conf = (args.dest_host, args.dest_user, args.dest_pass)
+    # 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,
+    }
 
     # Expand path
     local_path = os.path.expanduser(args.src_path)
@@ -795,6 +1087,9 @@ 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"Workers         : {args.workers}")
     if args.gmail_mode:
         print("Mode            : Gmail Restore with Labels + Flags")
@@ -806,11 +1101,14 @@ def main():
         print("Apply Flags     : Yes (read/starred/answered/draft)")
     if args.dest_delete:
         print("Dest Delete     : Yes (remove orphans from destination)")
+    print(
+        f"Restore Mode    : {'Full (all emails)' if args.full_restore else 'Incremental (new emails only, use --full-restore to restore all)'}"
+    )
     print("-----------------------------\n")
 
     try:
         # Test connection
-        dest = imap_common.get_imap_connection(*dest_conf)
+        dest = imap_common.get_imap_connection_from_conf(dest_conf)
         if not dest:
             print("Error: Could not connect to destination server.")
             sys.exit(1)
@@ -818,14 +1116,24 @@ def main():
 
         if args.gmail_mode:
             # Special Gmail mode
-            restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags)
+            restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore=args.full_restore)
         elif args.folder:
             # Restore specific folder
             folder_path = os.path.join(local_path, args.folder.replace("/", os.sep))
             if not os.path.exists(folder_path):
                 print(f"Error: Folder not found: {folder_path}")
                 sys.exit(1)
-            restore_folder(args.folder, folder_path, dest_conf, manifest, apply_labels, apply_flags, args.dest_delete)
+            restore_folder(
+                args.folder,
+                folder_path,
+                dest_conf,
+                manifest,
+                apply_labels,
+                apply_flags,
+                args.dest_delete,
+                full_restore=args.full_restore,
+                cache_root=local_path,
+            )
         else:
             # Restore all folders
             folders = get_backup_folders(local_path)
@@ -846,6 +1154,8 @@ def main():
                     apply_labels,
                     apply_flags,
                     args.dest_delete,
+                    full_restore=args.full_restore,
+                    cache_root=local_path,
                 )
 
         print("\nRestore completed successfully.")

test/conftest.py

@@ -103,7 +103,7 @@ def make_mock_connection(src_port, dest_port, src_user="src_user", dest_user="de
     Creates a mock connection function that routes to the correct server based on username.
     """
 
-    def mock_conn(host, user, pwd):
+    def mock_conn(host, user, pwd, oauth2_token=None):
         if user == src_user:
             port = src_port
         elif user == dest_user:
@@ -122,7 +122,7 @@ def make_single_mock_connection(port):
     Creates a mock connection function for a single server.
     """
 
-    def mock_conn(host, user, pwd):
+    def mock_conn(host, user, pwd, oauth2_token=None):
         c = imaplib.IMAP4("localhost", port)
         c.login(user, pwd)
         return c

test/test_backup_imap_emails.py

@@ -18,6 +18,7 @@ import pytest
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
 import backup_imap_emails
+import imap_common
 from conftest import make_single_mock_connection
 
 
@@ -27,7 +28,7 @@ class TestBackupBasic:
     def test_single_email_backup(self, single_mock_server, monkeypatch, tmp_path):
         """Test backing up a single email to local directory."""
         src_data = {"INBOX": [b"Subject: Test Email\r\nMessage-ID: <1@test>\r\n\r\nBody content"]}
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -60,7 +61,7 @@ class TestBackupBasic:
                 b"Subject: Email 3\r\nMessage-ID: <3@test>\r\n\r\nBody 3",
             ]
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -89,7 +90,7 @@ class TestIncrementalBackup:
                 b"Subject: Email 2\r\nMessage-ID: <2@test>\r\n\r\nBody 2",
             ]
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         # Create pre-existing file
         inbox_path = tmp_path / "INBOX"
@@ -126,7 +127,7 @@ class TestMultipleFolderBackup:
             "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
             "Archive": [b"Subject: Archive\r\nMessage-ID: <3@test>\r\n\r\nC"],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -149,7 +150,7 @@ class TestMultipleFolderBackup:
             "INBOX": [b"Subject: Inbox\r\nMessage-ID: <1@test>\r\n\r\nC"],
             "Sent": [b"Subject: Sent\r\nMessage-ID: <2@test>\r\n\r\nC"],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -173,7 +174,7 @@ class TestEmptyFolderHandling:
     def test_empty_folder(self, single_mock_server, monkeypatch, tmp_path):
         """Test handling of empty folders."""
         src_data = {"INBOX": [], "Empty": []}
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -191,16 +192,16 @@ class TestEmptyFolderHandling:
 class TestConfigValidation:
     """Tests for configuration validation."""
 
-    def test_missing_credentials(self, monkeypatch, capsys):
+    def test_missing_credentials(self, monkeypatch, capsys, tmp_path):
         """Test that missing credentials cause exit."""
         env = {}
         monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", "/tmp"])
+        monkeypatch.setattr(sys, "argv", ["backup_imap_emails.py", "--dest-path", str(tmp_path)])
 
         with pytest.raises(SystemExit) as exc_info:
             backup_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
     def test_missing_dest_path(self, monkeypatch, capsys):
         """Test that missing destination path causes exit."""
@@ -215,7 +216,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             backup_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
 
 class TestGetExistingUids:
@@ -264,30 +265,28 @@ class TestGetExistingUids:
 class TestBackupErrorHandling:
     """Tests for error handling scenarios in backup."""
 
-    def test_connection_error_in_worker(self, monkeypatch):
+    def test_connection_error_in_worker(self, monkeypatch, tmp_path):
         """Test worker handles connection failure gracefully."""
         # Mock get_imap_connection to fail
         monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: None)
 
         # Should return None/Exit without crashing
-        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), "/tmp")
+        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path))
 
-    def test_select_error_in_worker(self, monkeypatch):
+    def test_select_error_in_worker(self, monkeypatch, tmp_path):
         """Test worker handles SELECT failure."""
         mock_conn = MagicMock()
         mock_conn.select.side_effect = Exception("Select error")
         monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn)
 
         # Should log error and return
-        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), "/tmp")
+        backup_imap_emails.process_batch([], "INBOX", ("h", "u", "p"), str(tmp_path))
         mock_conn.select.assert_called()
 
     def test_fetch_body_error(self, monkeypatch, tmp_path):
         """Test handling of fetch body failure."""
         mock_conn = MagicMock()
         mock_conn.select.return_value = "OK"
-        # get_msg_details calls fetch headers, make it work
-        monkeypatch.setattr("imap_common.get_msg_details", lambda conn, uid: (uid, 100, "Subject"))
 
         # Fetch body fails
         mock_conn.uid.return_value = ("NO", [None])
@@ -306,7 +305,6 @@ class TestBackupErrorHandling:
         monkeypatch.setattr("imap_common.get_imap_connection", lambda *args: mock_conn)
 
         # Mock fetched data
-        monkeypatch.setattr("imap_common.get_msg_details", lambda conn, uid: (uid, 100, "Subject"))
         mock_conn.uid.return_value = ("OK", [(b"1 (RFC822 {10}", b"Content")])
 
         # Mock open to fail
@@ -317,7 +315,7 @@ class TestBackupErrorHandling:
 
         backup_imap_emails.process_batch([b"1"], "INBOX", ("h", "u", "p"), str(tmp_path))
 
-    def test_folder_creation_error(self, monkeypatch):
+    def test_folder_creation_error(self, monkeypatch, tmp_path):
         """Test handling failure to create local folder."""
 
         def mock_makedirs(path, exist_ok=False):
@@ -328,25 +326,25 @@ class TestBackupErrorHandling:
         mock_conn = MagicMock()
 
         # backup_folder should return early
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", "/tmp", ("h", "u", "p"))
+        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
         mock_conn.select.assert_not_called()
 
-    def test_select_folder_error(self, monkeypatch):
+    def test_select_folder_error(self, monkeypatch, tmp_path):
         """Test handling of select folder failure in main loop."""
         mock_conn = MagicMock()
         mock_conn.select.side_effect = Exception("Select failed")
         monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None)
 
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", "/tmp", ("h", "u", "p"))
+        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
         mock_conn.uid.assert_not_called()
 
-    def test_search_error(self, monkeypatch):
+    def test_search_error(self, monkeypatch, tmp_path):
         """Test handling of search failure."""
         mock_conn = MagicMock()
         mock_conn.uid.return_value = ("NO", [])
         monkeypatch.setattr(os, "makedirs", lambda p, exist_ok: None)
 
-        backup_imap_emails.backup_folder(mock_conn, "INBOX", "/tmp", ("h", "u", "p"))
+        backup_imap_emails.backup_folder(mock_conn, "INBOX", str(tmp_path), ("h", "u", "p"))
 
     def test_main_makedirs_error(self, monkeypatch, capsys):
         """Test failure to create main backup directory."""
@@ -596,7 +594,7 @@ class TestGmailLabelsPreservation:
             "Work": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
             "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -627,7 +625,7 @@ class TestGmailLabelsPreservation:
                 b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody",
             ],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -668,7 +666,7 @@ class TestGmailLabelsPreservation:
             "[Gmail]/Bin",
             "[Gmail]/Important",
         }
-        assert backup_imap_emails.GMAIL_SYSTEM_FOLDERS == expected
+        assert imap_common.GMAIL_SYSTEM_FOLDERS == expected
 
     def test_gmail_mode_flag(self, single_mock_server, monkeypatch, tmp_path):
         """Test --gmail-mode flag backs up All Mail and creates manifest."""
@@ -677,7 +675,7 @@ class TestGmailLabelsPreservation:
             "Work": [b"Subject: Work Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
             "[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",

test/test_compare_imap_folders.py

@@ -20,7 +20,7 @@ import pytest
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
 import compare_imap_folders
-from conftest import make_mock_connection
+from conftest import make_mock_connection, make_single_mock_connection
 
 
 class TestFolderComparison:
@@ -32,7 +32,7 @@ class TestFolderComparison:
             "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
             "Sent": [b"Subject: 3\r\n\r\nB"],
         }
-        src_server, dest_server, p1, p2 = mock_server_factory(data, data.copy())
+        _, _, p1, p2 = mock_server_factory(data, data.copy())
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -60,7 +60,7 @@ class TestFolderComparison:
         dest_data = {
             "INBOX": [b"Subject: 1\r\n\r\nB"],
         }
-        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -90,7 +90,7 @@ class TestFolderComparison:
         dest_data = {
             "INBOX": [b"Subject: 1\r\n\r\nB"],
         }
-        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -120,7 +120,7 @@ class TestEmptyFolders:
         src_data = {"INBOX": [], "Empty": []}
         dest_data = {"INBOX": [], "Empty": []}
 
-        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -147,7 +147,7 @@ class TestGetEmailCount:
     def test_successful_count(self, single_mock_server):
         """Test successful email count."""
         data = {"INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"]}
-        server, port = single_mock_server(data)
+        _, port = single_mock_server(data)
 
         conn = imaplib.IMAP4("localhost", port)
         conn.login("user", "pass")
@@ -160,7 +160,7 @@ class TestGetEmailCount:
     def test_nonexistent_folder(self, single_mock_server):
         """Test count for non-existent folder returns None."""
         data = {"INBOX": []}
-        server, port = single_mock_server(data)
+        _, port = single_mock_server(data)
 
         conn = imaplib.IMAP4("localhost", port)
         conn.login("user", "pass")
@@ -227,7 +227,9 @@ class TestCompareFoldersErrorHandling:
         mock_src.list.return_value = ("NO", [])
         mock_dest = MagicMock()
 
-        monkeypatch.setattr("imap_common.get_imap_connection", lambda h, u, p: mock_src if u == "s" else mock_dest)
+        monkeypatch.setattr(
+            "imap_common.get_imap_connection", lambda h, u, p, oauth2_token=None: mock_src if u == "s" else mock_dest
+        )
 
         env = {
             "SRC_IMAP_HOST": "h",
@@ -248,7 +250,7 @@ class TestCompareFoldersErrorHandling:
     def test_select_failure(self, single_mock_server):
         """Test get_email_count handles select failure."""
         data = {"INBOX": []}
-        server, port = single_mock_server(data)
+        _, port = single_mock_server(data)
         conn = imaplib.IMAP4("localhost", port)
         conn.login("user", "pass")
 
@@ -294,7 +296,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             compare_imap_folders.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
     def test_missing_dest_credentials(self, monkeypatch, capsys):
         """Test that missing destination credentials cause exit."""
@@ -309,7 +311,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             compare_imap_folders.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
 
 class TestTotals:
@@ -325,7 +327,7 @@ class TestTotals:
             "INBOX": [b"Subject: 1\r\n\r\nB"],
             "Sent": [b"Subject: 3\r\n\r\nB"],
         }
-        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        _, _, p1, p2 = mock_server_factory(src_data, dest_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",
@@ -346,3 +348,78 @@ class TestTotals:
         assert "TOTAL" in captured.out
         assert "5" in captured.out
         assert "2" in captured.out
+
+
+class TestLocalFolderComparison:
+    """Tests for comparing IMAP folders to local backup folders."""
+
+    def test_local_source_to_imap_dest(self, single_mock_server, monkeypatch, tmp_path, capsys):
+        """Local source folder list drives the comparison; destination is IMAP."""
+        # Local source: INBOX has 2 emails, Archive has 1
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
+
+        archive_path = tmp_path / "Archive"
+        archive_path.mkdir()
+        (archive_path / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
+
+        # Destination IMAP: INBOX has 1, Archive is missing
+        dest_data = {"INBOX": [b"Subject: 1\r\n\r\nB"]}
+        _, port = single_mock_server(dest_data)
+
+        env = {
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "dest_user",
+            "DEST_IMAP_PASSWORD": "p",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr(
+            sys,
+            "argv",
+            ["compare_imap_folders.py", "--src-path", str(tmp_path)],
+        )
+        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
+
+        compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Archive" in captured.out
+        # Destination should show N/A for missing Archive
+        assert "N/A" in captured.out
+
+    def test_imap_source_to_local_dest(self, single_mock_server, monkeypatch, tmp_path, capsys):
+        """IMAP source folder list drives the comparison; destination is local."""
+        # Source IMAP: INBOX has 2, Sent has 1
+        src_data = {
+            "INBOX": [b"Subject: 1\r\n\r\nB", b"Subject: 2\r\n\r\nB"],
+            "Sent": [b"Subject: 3\r\n\r\nB"],
+        }
+        _, port = single_mock_server(src_data)
+
+        # Local dest: INBOX has 1, Sent missing
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "src_user",
+            "SRC_IMAP_PASSWORD": "p",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr(
+            sys,
+            "argv",
+            ["compare_imap_folders.py", "--dest-path", str(tmp_path)],
+        )
+        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
+
+        compare_imap_folders.main()
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "Sent" in captured.out
+        assert "N/A" in captured.out

test/test_count_imap_emails.py

@@ -34,7 +34,7 @@ class TestEmailCounting:
                 b"Subject: Email 3\r\n\r\nBody",
             ]
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
 
@@ -51,7 +51,7 @@ class TestEmailCounting:
             "Sent": [b"Subject: 3\r\n\r\nB"],
             "Archive": [b"Subject: 4\r\n\r\nB", b"Subject: 5\r\n\r\nB", b"Subject: 6\r\n\r\nB"],
         }
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
 
@@ -67,7 +67,7 @@ class TestEmailCounting:
     def test_empty_folder(self, single_mock_server, monkeypatch, capsys):
         """Test counting in empty folders."""
         src_data = {"INBOX": [], "Empty": []}
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
 
@@ -77,6 +77,28 @@ class TestEmailCounting:
         assert "0" in captured.out
 
 
+class TestLocalEmailCounting:
+    """Tests for counting emails from a local backup folder."""
+
+    def test_count_local_folders(self, tmp_path, capsys):
+        inbox_path = tmp_path / "INBOX"
+        inbox_path.mkdir()
+        (inbox_path / "1_a.eml").write_bytes(b"Subject: A\r\n\r\nBody")
+        (inbox_path / "2_b.eml").write_bytes(b"Subject: B\r\n\r\nBody")
+
+        gmail_all_mail = tmp_path / "[Gmail]" / "All Mail"
+        gmail_all_mail.mkdir(parents=True)
+        (gmail_all_mail / "1_c.eml").write_bytes(b"Subject: C\r\n\r\nBody")
+
+        count_imap_emails.count_local_emails(str(tmp_path))
+
+        captured = capsys.readouterr()
+        assert "INBOX" in captured.out
+        assert "[Gmail]/All Mail" in captured.out
+        assert "TOTAL" in captured.out
+        assert "3" in captured.out
+
+
 class TestEmailCountingErrors:
     """Tests for error handling in email counting."""
 
@@ -145,7 +167,7 @@ class TestEmailCountingErrors:
         count_imap_emails.count_emails("host", "user", "pass")
 
         captured = capsys.readouterr()
-        assert "IMAP Error: Crash listing" in captured.out
+        assert "Failed to list mailboxes" in captured.out
 
     def test_imap_exception_during_select(self, monkeypatch, capsys):
         """Test handling of IMAP4 exception during folder selection."""
@@ -180,7 +202,7 @@ class TestMainFunction:
     def test_main_with_env_vars(self, single_mock_server, monkeypatch, capsys):
         """Test main function with environment variables."""
         src_data = {"INBOX": [b"Subject: Test\r\n\r\nBody"]}
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "IMAP_HOST": "localhost",
@@ -198,25 +220,14 @@ class TestMainFunction:
         assert "INBOX" in captured.out
 
     def test_missing_credentials(self, monkeypatch, capsys):
-        """Test that missing credentials cause exit."""
-        env = {}
-        monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["count_imap_emails.py"])
+        """Test that missing auth is rejected by argparse (neither password nor OAuth2 client-id)."""
+        monkeypatch.setattr(os, "environ", {})
+        monkeypatch.setattr(sys, "argv", ["count_imap_emails.py", "--host", "localhost", "--user", "user"])
 
-        # Since the script uses if __name__ == "__main__", we need to test differently
-        # Test the validation path directly
-        with pytest.raises(SystemExit):
-            # Simulate running main
-            import argparse
+        with pytest.raises(SystemExit) as exc_info:
+            count_imap_emails.main()
 
-            parser = argparse.ArgumentParser()
-            parser.add_argument("--host", default=None)
-            parser.add_argument("--user", default=None)
-            parser.add_argument("--pass", dest="password", default=None)
-            args = parser.parse_args([])
-
-            if not all([args.host, args.user, args.password]):
-                sys.exit(1)
+        assert exc_info.value.code == 2
 
 
 class TestSrcImapFallback:
@@ -225,7 +236,7 @@ class TestSrcImapFallback:
     def test_src_imap_vars_fallback(self, single_mock_server, monkeypatch, capsys):
         """Test that SRC_IMAP_* vars work as fallback."""
         src_data = {"INBOX": [b"Subject: Test\r\n\r\nBody"]}
-        server, port = single_mock_server(src_data)
+        _, port = single_mock_server(src_data)
 
         env = {
             "SRC_IMAP_HOST": "localhost",

test/test_imap_common.py

@@ -274,7 +274,7 @@ class TestMessageExistsInFolder:
     def test_no_message_id(self):
         """Test returns False when message_id is None."""
         mock_conn = Mock()
-        result = imap_common.message_exists_in_folder(mock_conn, None, 100)
+        result = imap_common.message_exists_in_folder(mock_conn, None)
         assert result is False
 
     def test_search_fails(self):
@@ -282,7 +282,7 @@ class TestMessageExistsInFolder:
         mock_conn = Mock()
         mock_conn.search.return_value = ("NO", [])
 
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>", 100)
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
         assert result is False
 
     def test_no_matches(self):
@@ -290,47 +290,448 @@ class TestMessageExistsInFolder:
         mock_conn = Mock()
         mock_conn.search.return_value = ("OK", [b""])
 
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>", 100)
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
         assert result is False
 
-    def test_match_found_same_size(self):
-        """Test returns True when message with same ID and size found."""
+    def test_match_found(self):
+        """Test returns True when message with same ID found."""
         mock_conn = Mock()
         mock_conn.search.return_value = ("OK", [b"1"])
-        mock_conn.fetch.return_value = ("OK", [b"1 (RFC822.SIZE 100)"])
 
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>", 100)
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
         assert result is True
 
-    def test_match_found_different_size(self):
-        """Test returns False when message with same ID but different size found."""
+    def test_search_exception(self):
+        """Test returns False when search raises an exception."""
         mock_conn = Mock()
-        mock_conn.search.return_value = ("OK", [b"1"])
-        mock_conn.fetch.return_value = ("OK", [b"1 (RFC822.SIZE 200)"])
+        mock_conn.search.side_effect = Exception("Connection error")
 
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>", 100)
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
         assert result is False
 
 
-class TestGetMsgDetails:
-    """Tests for get_msg_details function."""
+class TestExtractMessageId:
+    """Tests for extract_message_id function."""
 
-    def test_fetch_error(self):
-        """Test returns None tuple on fetch error."""
+    def test_standard_header(self):
+        """Test extracting standard Message-ID."""
+        header = b"Message-ID: <test@example.com>\r\nSubject: Test"
+        result = imap_common.extract_message_id(header)
+        assert result == "<test@example.com>"
+
+    def test_lowercase_header(self):
+        """Test extracting Message-ID with lowercase header name."""
+        header = b"message-id: <lower@example.com>\r\nSubject: Test"
+        result = imap_common.extract_message_id(header)
+        assert result == "<lower@example.com>"
+
+    def test_folded_header(self):
+        """Test extracting folded Message-ID."""
+        # BytesParser unfolds values (replacing newline+indent with space)
+        header = b"Message-ID: <part1\r\n part2@example.com>\r\nSubject: Test"
+        result = imap_common.extract_message_id(header)
+        assert result == "<part1 part2@example.com>"
+
+    def test_string_input(self):
+        """Test with string input."""
+        header = "Message-ID: <str@example.com>\nSubject: Test"
+        result = imap_common.extract_message_id(header)
+        assert result == "<str@example.com>"
+
+    def test_no_message_id(self):
+        """Test header without Message-ID."""
+        header = b"Subject: Test\r\nFrom: sender"
+        result = imap_common.extract_message_id(header)
+        assert result is None
+
+    def test_empty_input(self):
+        """Test empty input."""
+        assert imap_common.extract_message_id(None) is None
+        assert imap_common.extract_message_id(b"") is None
+
+
+class TestEnsureFolderExists:
+    """Tests for ensure_folder_exists function."""
+
+    def test_creates_folder_successfully(self):
+        """Test folder is created when it doesn't exist."""
         mock_conn = Mock()
-        mock_conn.uid.side_effect = Exception("Fetch error")
+        mock_conn.create.return_value = ("OK", [])
 
-        msg_id, size, subject = imap_common.get_msg_details(mock_conn, b"1")
-        assert msg_id is None
-        assert size is None
-        assert subject is None
+        # Should not raise any exception
+        imap_common.ensure_folder_exists(mock_conn, "TestFolder")
+        mock_conn.create.assert_called_once_with('"TestFolder"')
 
-    def test_not_ok_response(self):
-        """Test returns None tuple on non-OK response."""
+    def test_folder_already_exists(self):
+        """Test exception is suppressed when folder already exists."""
         mock_conn = Mock()
-        mock_conn.uid.return_value = ("NO", None)
+        mock_conn.create.side_effect = Exception("Folder already exists")
 
-        msg_id, size, subject = imap_common.get_msg_details(mock_conn, b"1")
-        assert msg_id is None
-        assert size is None
-        assert subject is None
+        # Should not raise any exception
+        imap_common.ensure_folder_exists(mock_conn, "ExistingFolder")
+        mock_conn.create.assert_called_once_with('"ExistingFolder"')
+
+    def test_inbox_folder_skipped(self):
+        """Test INBOX folder is not created."""
+        mock_conn = Mock()
+
+        imap_common.ensure_folder_exists(mock_conn, "INBOX")
+        mock_conn.create.assert_not_called()
+
+        # Also test lowercase
+        imap_common.ensure_folder_exists(mock_conn, "inbox")
+        mock_conn.create.assert_not_called()
+
+    def test_empty_folder_name(self):
+        """Test empty folder name is skipped."""
+        mock_conn = Mock()
+
+        imap_common.ensure_folder_exists(mock_conn, "")
+        mock_conn.create.assert_not_called()
+
+    def test_none_folder_name(self):
+        """Test None folder name is skipped."""
+        mock_conn = Mock()
+
+        imap_common.ensure_folder_exists(mock_conn, None)
+        mock_conn.create.assert_not_called()
+
+    def test_server_restriction_error_suppressed(self):
+        """Test server restriction errors are suppressed."""
+        mock_conn = Mock()
+        mock_conn.create.side_effect = Exception("Permission denied")
+
+        # Should not raise any exception
+        imap_common.ensure_folder_exists(mock_conn, "RestrictedFolder")
+        mock_conn.create.assert_called_once()
+
+
+class TestAppendEmail:
+    """Tests for append_email function."""
+
+    def test_successful_append(self):
+        """Test successful email append returns True."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            None,
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_flags_with_parentheses(self):
+        """Test flag normalization when flags already have parentheses."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags="(\\Seen \\Flagged)",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        # Flags should remain unchanged
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            "(\\Seen \\Flagged)",
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_flags_without_parentheses(self):
+        """Test flag normalization when flags lack parentheses."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags="\\Seen",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        # Flags should be wrapped in parentheses
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            "(\\Seen)",
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_none_flags(self):
+        """Test append with None flags."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags=None,
+            ensure_folder=False,
+        )
+
+        assert result is True
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            None,
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_empty_flags(self):
+        """Test append with empty string flags."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags="",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        # Empty flags should result in None
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            None,
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_whitespace_only_flags(self):
+        """Test append with whitespace-only flags."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags="   ",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        # Whitespace-only flags should result in None
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            None,
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+    def test_append_with_ensure_folder_enabled(self):
+        """Test append with ensure_folder enabled."""
+        mock_conn = Mock()
+        mock_conn.create.return_value = ("OK", [])
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=True,
+        )
+
+        assert result is True
+        mock_conn.create.assert_called_once_with('"TestFolder"')
+        mock_conn.append.assert_called_once()
+
+    def test_append_with_ensure_folder_disabled(self):
+        """Test append with ensure_folder disabled."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        mock_conn.create.assert_not_called()
+        mock_conn.append.assert_called_once()
+
+    def test_append_failure_returns_false(self):
+        """Test append returns False on failure."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("NO", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=False,
+        )
+
+        assert result is False
+
+    def test_append_exception_returns_false(self):
+        """Test append returns False when exception occurs."""
+        mock_conn = Mock()
+        mock_conn.append.side_effect = Exception("Connection error")
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            ensure_folder=False,
+        )
+
+        assert result is False
+
+    def test_append_with_multiple_flags(self):
+        """Test append with multiple flags."""
+        mock_conn = Mock()
+        mock_conn.append.return_value = ("OK", [])
+
+        result = imap_common.append_email(
+            mock_conn,
+            "TestFolder",
+            b"email content",
+            "01-Jan-2024 12:00:00 +0000",
+            flags="\\Seen \\Flagged \\Answered",
+            ensure_folder=False,
+        )
+
+        assert result is True
+        # Multiple flags without parentheses should be wrapped
+        mock_conn.append.assert_called_once_with(
+            '"TestFolder"',
+            "(\\Seen \\Flagged \\Answered)",
+            "01-Jan-2024 12:00:00 +0000",
+            b"email content",
+        )
+
+
+class TestGetMessageIdsInFolder:
+    """Tests for get_message_ids_in_folder function."""
+
+    def test_empty_folder(self):
+        """Test returns empty dict for folder with no messages."""
+        mock_conn = Mock()
+        mock_conn.uid.return_value = ("OK", [b""])
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {}
+
+    def test_search_fails(self):
+        """Test returns empty dict when search fails."""
+        mock_conn = Mock()
+        mock_conn.uid.return_value = ("NO", [])
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {}
+
+    def test_search_exception(self):
+        """Test returns empty dict when search raises exception."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = Exception("Connection error")
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {}
+
+    def test_single_message(self):
+        """Test fetching single message ID."""
+        mock_conn = Mock()
+        # First call: search returns one UID
+        # Second call: fetch returns the Message-ID header (with UID in response)
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1"]),
+            ("OK", [(b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <test@example.com>\r\n"), b")"]),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {b"1": "<test@example.com>"}
+        assert set(result.values()) == {"<test@example.com>"}
+
+    def test_multiple_messages(self):
+        """Test fetching multiple message IDs."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1 2 3"]),
+            (
+                "OK",
+                [
+                    (b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <msg1@example.com>\r\n"),
+                    b")",
+                    (b"2 (UID 2 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <msg2@example.com>\r\n"),
+                    b")",
+                    (b"3 (UID 3 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <msg3@example.com>\r\n"),
+                    b")",
+                ],
+            ),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert set(result.values()) == {"<msg1@example.com>", "<msg2@example.com>", "<msg3@example.com>"}
+
+    def test_fetch_fails_for_batch(self):
+        """Test continues when fetch fails for a batch."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1"]),
+            ("NO", []),  # Fetch fails
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {}
+
+    def test_fetch_exception_for_batch(self):
+        """Test continues when fetch raises exception for a batch."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1"]),
+            Exception("Fetch error"),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert result == {}
+
+    def test_skips_empty_message_id(self):
+        """Test that empty message IDs are not added to dict."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1 2"]),
+            (
+                "OK",
+                [
+                    (b"1 (UID 1 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: <valid@example.com>\r\n"),
+                    b")",
+                    (b"2 (UID 2 BODY[HEADER.FIELDS (MESSAGE-ID)] {50}", b"Message-ID: \r\n"),  # Empty
+                    b")",
+                ],
+            ),
+        ]
+
+        result = imap_common.get_message_ids_in_folder(mock_conn)
+        assert set(result.values()) == {"<valid@example.com>"}

test/test_imap_oauth2.py

@@ -0,0 +1,489 @@
+"""
+Tests for imap_oauth2.py
+
+Tests cover:
+- OAuth2 provider detection
+- Microsoft tenant discovery
+- Microsoft token acquisition and caching
+- Google token acquisition and caching
+- Provider dispatch
+- Token refresh (caching and silent refresh)
+- Thread-safe token refresh
+"""
+
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_oauth2
+
+
+@pytest.fixture(autouse=True)
+def clear_oauth2_caches():
+    """Clear module-level OAuth2 caches between tests."""
+    imap_oauth2._msal_app_cache.clear()
+    imap_oauth2._google_creds_cache.clear()
+    yield
+    imap_oauth2._msal_app_cache.clear()
+    imap_oauth2._google_creds_cache.clear()
+
+
+class TestDetectOauth2Provider:
+    """Tests for detect_oauth2_provider function."""
+
+    def test_microsoft_outlook(self):
+        """Test detects Microsoft from outlook host."""
+        assert imap_oauth2.detect_oauth2_provider("outlook.office365.com") == "microsoft"
+
+    def test_microsoft_office365(self):
+        """Test detects Microsoft from office365 host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.office365.com") == "microsoft"
+
+    def test_microsoft_mixed_case(self):
+        """Test detects Microsoft case-insensitively."""
+        assert imap_oauth2.detect_oauth2_provider("Outlook.Office365.COM") == "microsoft"
+
+    def test_google_gmail(self):
+        """Test detects Google from gmail host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.gmail.com") == "google"
+
+    def test_google_googlemail(self):
+        """Test detects Google from google host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.google.com") == "google"
+
+    def test_unknown_provider(self):
+        """Test returns None for unrecognized host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.example.com") is None
+
+    def test_unknown_yahoo(self):
+        """Test returns None for Yahoo host."""
+        assert imap_oauth2.detect_oauth2_provider("imap.mail.yahoo.com") is None
+
+
+class TestDiscoverMicrosoftTenant:
+    """Tests for discover_microsoft_tenant function."""
+
+    def test_successful_discovery(self):
+        """Test successful tenant ID extraction from OpenID config."""
+        tenant_id = "12345678-abcd-ef01-2345-67890abcdef0"
+        data = {
+            "issuer": f"https://sts.windows.net/{tenant_id}/",
+            "authorization_endpoint": f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize",
+        }
+
+        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data):
+            result = imap_oauth2.discover_microsoft_tenant("user@contoso.com")
+
+        assert result == tenant_id
+
+    def test_domain_extraction(self):
+        """Test that domain is correctly extracted from email."""
+        data = {"issuer": "https://sts.windows.net/abcdef01-2345-6789-abcd-ef0123456789/"}
+
+        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data) as mock_fetch:
+            imap_oauth2.discover_microsoft_tenant("user@example.org")
+            host, path = mock_fetch.call_args[0][0], mock_fetch.call_args[0][1]
+            assert host == "login.microsoftonline.com"
+            assert "example.org" in path
+
+    def test_network_error(self, capsys):
+        """Test returns None on network error."""
+        with patch.object(imap_oauth2, "_fetch_json_https", side_effect=OSError("Connection refused")):
+            result = imap_oauth2.discover_microsoft_tenant("user@invalid.example")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Could not discover" in captured.out
+
+    def test_invalid_json(self, capsys):
+        """Test returns None on invalid JSON response."""
+        with patch.object(
+            imap_oauth2,
+            "_fetch_json_https",
+            side_effect=json.JSONDecodeError("Expecting value", "not json", 0),
+        ):
+            result = imap_oauth2.discover_microsoft_tenant("user@test.com")
+
+        assert result is None
+
+    def test_no_tenant_in_issuer(self, capsys):
+        """Test returns None when issuer has no tenant GUID."""
+        data = {"issuer": "https://sts.windows.net/not-a-guid/"}
+
+        with patch.object(imap_oauth2, "_fetch_json_https", return_value=data):
+            result = imap_oauth2.discover_microsoft_tenant("user@test.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Could not extract tenant ID" in captured.out
+
+
+class TestAcquireMicrosoftOauth2Token:
+    """Tests for acquire_microsoft_oauth2_token function."""
+
+    def test_successful_token(self):
+        """Test successful token acquisition with auto-discovery."""
+        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC123", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "test_token"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+            assert result == "test_token"
+
+    def test_tenant_discovery_failure(self, capsys):
+        """Test returns None when tenant discovery fails."""
+        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value=None):
+            result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+        assert result is None
+
+    def test_cached_token(self):
+        """Test returns cached token when available."""
+        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_account = {"username": "user@test.com"}
+            mock_app.get_accounts.return_value = [mock_account]
+            mock_app.acquire_token_silent.return_value = {"access_token": "cached_token"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+            assert result == "cached_token"
+
+
+class TestAcquireGoogleOauth2Token:
+    """Tests for acquire_google_oauth2_token function."""
+
+    def test_successful_token(self):
+        """Test successful Google token acquisition."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = "google_test_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+        assert result == "google_test_token"
+
+    def test_missing_library(self):
+        """Test exits when google-auth-oauthlib is not installed."""
+        with patch.dict("sys.modules", {"google_auth_oauthlib": None, "google_auth_oauthlib.flow": None}):
+            with pytest.raises(SystemExit):
+                imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+    def test_no_token_returned(self):
+        """Test returns None when credentials have no token."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = None
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+        assert result is None
+
+
+class TestAcquireOauth2TokenForProvider:
+    """Tests for acquire_oauth2_token_for_provider dispatch function."""
+
+    def test_dispatch_to_microsoft(self):
+        """Test dispatches to Microsoft when provider is 'microsoft'."""
+        with patch.object(imap_oauth2, "acquire_microsoft_oauth2_token", return_value="ms_token") as mock_ms:
+            result = imap_oauth2.acquire_oauth2_token_for_provider("microsoft", "cid", "user@test.com")
+
+        assert result == "ms_token"
+        mock_ms.assert_called_once_with("cid", "user@test.com")
+
+    def test_dispatch_to_google(self):
+        """Test dispatches to Google when provider is 'google'."""
+        with patch.object(imap_oauth2, "acquire_google_oauth2_token", return_value="g_token") as mock_g:
+            result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com", "secret")
+
+        assert result == "g_token"
+        mock_g.assert_called_once_with("cid", "secret")
+
+    def test_google_requires_client_secret(self, capsys):
+        """Test returns None when Google is selected without client_secret."""
+        result = imap_oauth2.acquire_oauth2_token_for_provider("google", "cid", "user@gmail.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "--oauth2-client-secret" in captured.out
+
+    def test_unknown_provider(self, capsys):
+        """Test returns None for unknown provider."""
+        result = imap_oauth2.acquire_oauth2_token_for_provider("yahoo", "cid", "user@yahoo.com")
+
+        assert result is None
+        captured = capsys.readouterr()
+        assert "Unknown OAuth2 provider" in captured.out
+
+
+class TestMicrosoftTokenRefresh:
+    """Tests for Microsoft OAuth2 token caching and refresh."""
+
+    def test_msal_app_cached_on_first_call(self):
+        """Test MSAL app is cached after first call."""
+        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+            assert ("client-id", "tenant-123") in imap_oauth2._msal_app_cache
+
+    def test_cached_app_reused_on_second_call(self):
+        """Test second call reuses cached MSAL app instead of creating new one."""
+        with patch.object(imap_oauth2, "discover_microsoft_tenant", return_value="tenant-123"):
+            mock_msal = MagicMock()
+            mock_app = MagicMock()
+            mock_app.get_accounts.return_value = []
+            mock_app.initiate_device_flow.return_value = {"user_code": "ABC", "message": "Go to..."}
+            mock_app.acquire_token_by_device_flow.return_value = {"access_token": "token1"}
+            mock_msal.PublicClientApplication.return_value = mock_app
+
+            with patch.dict("sys.modules", {"msal": mock_msal}):
+                imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+                # Second call — simulate cached token available (refresh token worked)
+                mock_account = {"username": "user@test.com"}
+                mock_app.get_accounts.return_value = [mock_account]
+                mock_app.acquire_token_silent.return_value = {"access_token": "refreshed_token"}
+
+                result = imap_oauth2.acquire_microsoft_oauth2_token("client-id", "user@test.com")
+
+            assert result == "refreshed_token"
+            # PublicClientApplication should only have been called once (first call)
+            assert mock_msal.PublicClientApplication.call_count == 1
+
+
+class TestGoogleTokenRefresh:
+    """Tests for Google OAuth2 token caching and refresh."""
+
+    def test_credentials_cached_on_first_call(self):
+        """Test Google credentials are cached after first call."""
+        mock_credentials = MagicMock()
+        mock_credentials.token = "google_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+        assert ("client-id", "client-secret") in imap_oauth2._google_creds_cache
+
+    def test_cached_credentials_refreshed_on_second_call(self):
+        """Test second call refreshes cached credentials without opening browser."""
+        # Pre-populate cache with credentials that have a refresh token
+        mock_creds = MagicMock()
+        mock_creds.refresh_token = "refresh_tok"
+        mock_creds.token = "refreshed_google_token"
+        imap_oauth2._google_creds_cache[("client-id", "client-secret")] = mock_creds
+
+        mock_request_module = MagicMock()
+        with patch.dict(
+            "sys.modules",
+            {
+                "google": MagicMock(),
+                "google.auth": MagicMock(),
+                "google.auth.transport": MagicMock(),
+                "google.auth.transport.requests": mock_request_module,
+            },
+        ):
+            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+        assert result == "refreshed_google_token"
+        # Verify refresh was called
+        mock_creds.refresh.assert_called_once()
+
+    def test_falls_back_to_browser_if_refresh_fails(self):
+        """Test falls back to full auth flow if cached token refresh fails."""
+        # Pre-populate cache with credentials whose refresh fails
+        mock_creds = MagicMock()
+        mock_creds.refresh_token = "refresh_tok"
+        mock_creds.refresh.side_effect = Exception("Refresh failed")
+        imap_oauth2._google_creds_cache[("client-id", "client-secret")] = mock_creds
+
+        # Set up the full auth flow
+        mock_credentials = MagicMock()
+        mock_credentials.token = "new_browser_token"
+
+        mock_flow = MagicMock()
+        mock_flow.run_local_server.return_value = mock_credentials
+
+        mock_installed_app_flow = MagicMock()
+        mock_installed_app_flow.from_client_config.return_value = mock_flow
+
+        mock_module = MagicMock()
+        mock_module.InstalledAppFlow = mock_installed_app_flow
+
+        with patch.dict("sys.modules", {"google_auth_oauthlib": MagicMock(), "google_auth_oauthlib.flow": mock_module}):
+            result = imap_oauth2.acquire_google_oauth2_token("client-id", "client-secret")
+
+        assert result == "new_browser_token"
+
+
+class TestRefreshOauth2Token:
+    """Tests for thread-safe refresh_oauth2_token function."""
+
+    def test_refreshes_token_and_updates_conf(self):
+        """Test that a new token is acquired and conf["oauth2_token"] is updated."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value="new_token") as mock_acquire:
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result == "new_token"
+        assert conf["oauth2_token"] == "new_token"
+        mock_acquire.assert_called_once_with("microsoft", "client-id", "user@test.com", None)
+
+    def test_skips_refresh_when_token_already_updated(self):
+        """Test that refresh is skipped if another thread already updated the token."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "already_refreshed_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider") as mock_acquire:
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result == "already_refreshed_token"
+        assert conf["oauth2_token"] == "already_refreshed_token"
+        mock_acquire.assert_not_called()
+
+    def test_returns_none_on_refresh_failure(self):
+        """Test returns None and leaves conf unchanged when refresh fails."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "google",
+                "client_id": "client-id",
+                "email": "user@gmail.com",
+                "client_secret": "secret",
+            },
+        }
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", return_value=None):
+            result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result is None
+        assert conf["oauth2_token"] == "old_token"
+
+    def test_returns_none_when_no_oauth2_context(self):
+        """Test returns None when oauth2 context is not set."""
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": None,
+        }
+
+        result = imap_oauth2.refresh_oauth2_token(conf, "old_token")
+
+        assert result is None
+        assert conf["oauth2_token"] == "old_token"
+
+    def test_concurrent_threads_only_one_refreshes(self):
+        """Test that only one thread performs the refresh when multiple threads compete."""
+        import threading
+        import time
+
+        conf = {
+            "host": "host",
+            "user": "user",
+            "password": "pass",
+            "oauth2_token": "expired_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@test.com",
+                "client_secret": None,
+            },
+        }
+        call_count = {"value": 0}
+        barrier = threading.Barrier(3)  # 3 threads
+
+        def slow_acquire(provider, client_id, email, client_secret):
+            call_count["value"] += 1
+            time.sleep(0.05)  # Simulate network delay
+            return "fresh_token"
+
+        def thread_func():
+            barrier.wait()  # Ensure all threads start at the same time
+            imap_oauth2.refresh_oauth2_token(conf, "expired_token")
+
+        with patch.object(imap_oauth2, "acquire_oauth2_token_for_provider", side_effect=slow_acquire):
+            threads = [threading.Thread(target=thread_func) for _ in range(3)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+
+        # Only one thread should have called acquire (the first to get the lock).
+        # The other two should see conf["oauth2_token"] changed and skip.
+        assert call_count["value"] == 1
+        assert conf["oauth2_token"] == "fresh_token"

test/test_migrate_imap_emails.py

@@ -269,6 +269,94 @@ class TestFolderHandling:
         migrate_imap_emails.main()
 
 
+class TestPreserveFlags:
+    def test_preserve_flags_applied_on_copy(self, mock_server_factory, monkeypatch):
+        msg = b"Subject: Flagged\r\nMessage-ID: <f1@test>\r\n\r\nBody"
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": []}
+
+        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        src_server.folders["INBOX"][0]["flags"].add("\\Seen")
+
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "src_user",
+            "SRC_IMAP_PASSWORD": "p",
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "dest_user",
+            "DEST_IMAP_PASSWORD": "p",
+            "MAX_WORKERS": "1",
+            "PRESERVE_FLAGS": "true",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        migrate_imap_emails.main()
+
+        assert len(dest_server.folders["INBOX"]) == 1
+        assert "\\Seen" in dest_server.folders["INBOX"][0]["flags"]
+
+    def test_preserve_flags_syncs_on_duplicate(self, mock_server_factory, monkeypatch):
+        msg = b"Subject: DupFlags\r\nMessage-ID: <df1@test>\r\n\r\nBody"
+        src_data = {"INBOX": [msg]}
+        dest_data = {"INBOX": [msg]}
+
+        src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+        src_server.folders["INBOX"][0]["flags"].add("\\Seen")
+
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "src_user",
+            "SRC_IMAP_PASSWORD": "p",
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "dest_user",
+            "DEST_IMAP_PASSWORD": "p",
+            "MAX_WORKERS": "1",
+            "PRESERVE_FLAGS": "true",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        migrate_imap_emails.main()
+
+        assert len(dest_server.folders["INBOX"]) == 1
+        assert "\\Seen" in dest_server.folders["INBOX"][0]["flags"]
+
+
+class TestGmailModeLabels:
+    def test_gmail_mode_migrates_all_mail_and_applies_labels(self, mock_server_factory, monkeypatch):
+        msg1 = b"Subject: One\r\nMessage-ID: <gm1@test>\r\n\r\nBody 1"
+        msg2 = b"Subject: Two\r\nMessage-ID: <gm2@test>\r\n\r\nBody 2"
+
+        src_data = {
+            "[Gmail]/All Mail": [msg1, msg2],
+            "INBOX": [msg1],
+            "Work": [msg1, msg2],
+        }
+        dest_data = {"INBOX": []}
+
+        _, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
+
+        env = {
+            "SRC_IMAP_HOST": "localhost",
+            "SRC_IMAP_USERNAME": "src_user",
+            "SRC_IMAP_PASSWORD": "p",
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "dest_user",
+            "DEST_IMAP_PASSWORD": "p",
+            "MAX_WORKERS": "1",
+            "GMAIL_MODE": "true",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr("imap_common.get_imap_connection", make_mock_connection(p1, p2))
+
+        migrate_imap_emails.main()
+
+        assert len(dest_server.folders["INBOX"]) == 1
+        assert "Work" in dest_server.folders
+        assert len(dest_server.folders["Work"]) == 2
+
+
 class TestConfigValidation:
     """Tests for configuration validation."""
 
@@ -284,7 +372,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             migrate_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
     def test_missing_dest_credentials(self, monkeypatch, capsys):
         """Test that missing destination credentials cause exit."""
@@ -298,7 +386,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             migrate_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
 
 class TestMigrateErrorHandling:
@@ -347,7 +435,7 @@ class TestMigrateErrorHandling:
 
         def side_effect_select(self, mailbox, readonly=False):
             if threading.current_thread() is not threading.main_thread():
-                raise Exception("Select failed")
+                raise RuntimeError("Select failed")
             return original_select(self, mailbox, readonly)
 
         monkeypatch.setattr(imaplib.IMAP4, "select", side_effect_select)
@@ -378,7 +466,7 @@ class TestMigrateErrorHandling:
 
         def side_effect_uid(self, command, *args):
             if command == "fetch" and threading.current_thread() is not threading.main_thread():
-                raise Exception("Fetch failed")
+                raise RuntimeError("Fetch failed")
             return original_uid(self, command, *args)
 
         monkeypatch.setattr(imaplib.IMAP4, "uid", side_effect_uid)
@@ -645,41 +733,3 @@ class TestDestDeleteFunctionality:
 
         # All dest emails should be deleted
         assert len(dest_server.folders["INBOX"]) == 0
-
-    def test_get_message_ids_in_folder(self, mock_server_factory, monkeypatch):
-        """Test get_message_ids_in_folder returns correct Message-IDs."""
-        data = {
-            "INBOX": [
-                b"Subject: Email 1\r\nMessage-ID: <msg1@test>\r\n\r\nBody",
-                b"Subject: Email 2\r\nMessage-ID: <msg2@test>\r\n\r\nBody",
-                b"Subject: Email 3\r\nMessage-ID: <msg3@test>\r\n\r\nBody",
-            ]
-        }
-
-        server, port = None, None
-        # Use single_mock_server fixture pattern
-        import time
-
-        from conftest import get_free_port, start_server_thread
-
-        port = get_free_port()
-        thread, server = start_server_thread(port, data)
-        time.sleep(0.3)
-
-        try:
-            import imaplib
-
-            conn = imaplib.IMAP4("localhost", port)
-            conn.login("user", "pass")
-
-            msg_ids = migrate_imap_emails.get_message_ids_in_folder(conn, "INBOX")
-
-            assert "<msg1@test>" in msg_ids
-            assert "<msg2@test>" in msg_ids
-            assert "<msg3@test>" in msg_ids
-            assert len(msg_ids) == 3
-
-            conn.logout()
-        finally:
-            server.shutdown()
-            thread.join(timeout=2)

test/test_restore_imap_emails.py

@@ -18,6 +18,8 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
+import imap_common
+import restore_cache
 import restore_imap_emails
 from conftest import make_single_mock_connection
 
@@ -135,7 +137,7 @@ Body content.
 
         message_id, date_str, raw_content, subject = restore_imap_emails.parse_eml_file(str(eml_file))
 
-        assert message_id == ""
+        assert message_id is None
         assert raw_content is not None
 
     def test_parse_nonexistent_file(self):
@@ -216,16 +218,16 @@ class TestGetBackupFolders:
 class TestConfigValidation:
     """Tests for configuration validation."""
 
-    def test_missing_credentials(self, monkeypatch, capsys):
+    def test_missing_credentials(self, monkeypatch, capsys, tmp_path):
         """Test that missing credentials cause exit."""
         env = {}
         monkeypatch.setattr(os, "environ", env)
-        monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", "/tmp"])
+        monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", str(tmp_path)])
 
         with pytest.raises(SystemExit) as exc_info:
             restore_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
     def test_missing_src_path(self, monkeypatch, capsys):
         """Test that missing source path causes exit."""
@@ -240,7 +242,7 @@ class TestConfigValidation:
         with pytest.raises(SystemExit) as exc_info:
             restore_imap_emails.main()
 
-        assert exc_info.value.code == 1
+        assert exc_info.value.code == 2
 
     def test_nonexistent_src_path(self, monkeypatch, capsys):
         """Test that non-existent source path causes exit."""
@@ -272,8 +274,8 @@ class TestUploadEmail:
         mock_conn.select.return_value = ("OK", [b"1"])
         mock_conn.append.return_value = ("OK", [])
 
-        # Mock message_exists_in_folder to return False (not a duplicate)
-        monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False)
+        # Mock email_exists_in_folder to return False (not a duplicate)
+        monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False)
 
         result = restore_imap_emails.upload_email(
             mock_conn,
@@ -281,19 +283,18 @@ class TestUploadEmail:
             b"raw email content",
             '"15-Jan-2024 10:30:00 +0000"',
             "<test@test.com>",
-            "Test Subject",
         )
 
         assert result is True
         mock_conn.append.assert_called_once()
 
     def test_upload_email_duplicate(self, monkeypatch):
-        """Test upload skips duplicate."""
+        """Test upload returns False when message exists."""
         mock_conn = MagicMock()
         mock_conn.select.return_value = ("OK", [b"1"])
 
-        # Mock message_exists_in_folder to return True (is a duplicate)
-        monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: True)
+        # Mock email_exists_in_folder to return True (is a duplicate)
+        monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: True)
 
         result = restore_imap_emails.upload_email(
             mock_conn,
@@ -301,7 +302,7 @@ class TestUploadEmail:
             b"raw email content",
             '"15-Jan-2024 10:30:00 +0000"',
             "<test@test.com>",
-            "Test Subject",
+            check_duplicate=True,
         )
 
         assert result is False
@@ -314,8 +315,8 @@ class TestUploadEmail:
         mock_conn.select.return_value = ("OK", [b"1"])
         mock_conn.append.return_value = ("OK", [])
 
-        # Mock message_exists_in_folder to return False (not a duplicate)
-        monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False)
+        # Mock email_exists_in_folder to return False (not a duplicate)
+        monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False)
 
         result = restore_imap_emails.upload_email(
             mock_conn,
@@ -323,14 +324,13 @@ class TestUploadEmail:
             b"raw email content",
             '"15-Jan-2024 10:30:00 +0000"',
             "<test@test.com>",
-            "Test Subject",
             flags="\\Seen",  # Mark as read
         )
 
         assert result is True
         # Check that append was called with the \\Seen flag
         call_args = mock_conn.append.call_args
-        assert call_args[0][1] == "\\Seen"
+        assert call_args[0][1] == "(\\Seen)"
 
 
 class TestRestoreIntegration:
@@ -372,6 +372,39 @@ Body content.
         # Run restore
         restore_imap_emails.main()
 
+    def test_restore_single_folder_full_restore_flag(self, single_mock_server, monkeypatch, tmp_path):
+        """Smoke test: --full-restore flag is accepted."""
+        inbox = tmp_path / "INBOX"
+        inbox.mkdir()
+
+        eml_content = b"""From: sender@test.com
+To: recipient@test.com
+Subject: Test Email
+Message-ID: <test123@test.com>
+Date: Mon, 15 Jan 2024 10:30:00 +0000
+
+Body content.
+"""
+        (inbox / "1_Test_Email.eml").write_bytes(eml_content)
+
+        src_data = {"INBOX": []}
+        _server, port = single_mock_server(src_data)
+
+        env = {
+            "DEST_IMAP_HOST": "localhost",
+            "DEST_IMAP_USERNAME": "user",
+            "DEST_IMAP_PASSWORD": "pass",
+        }
+        monkeypatch.setattr(os, "environ", env)
+        monkeypatch.setattr(
+            sys,
+            "argv",
+            ["restore_imap_emails.py", "--src-path", str(tmp_path), "--full-restore", "INBOX"],
+        )
+        monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
+
+        restore_imap_emails.main()
+
     def test_restore_with_labels_manifest(self, tmp_path):
         """Test that labels manifest is loaded correctly."""
         # Create manifest
@@ -396,7 +429,7 @@ class TestEmailExistsInFolder:
         mock_conn = MagicMock()
         monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: True)
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>", 1000)
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
         assert result is True
 
     def test_email_exists_false(self, monkeypatch):
@@ -404,28 +437,69 @@ class TestEmailExistsInFolder:
         mock_conn = MagicMock()
         monkeypatch.setattr("imap_common.message_exists_in_folder", lambda *args: False)
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>", 1000)
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
         assert result is False
 
     def test_email_exists_no_message_id(self, monkeypatch):
         """Test with no message ID."""
         mock_conn = MagicMock()
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, None, 1000)
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, None)
         assert result is False
 
     def test_email_exists_exception(self, monkeypatch):
         """Test handling exception."""
         mock_conn = MagicMock()
+
         monkeypatch.setattr(
             "imap_common.message_exists_in_folder",
             lambda *args: (_ for _ in ()).throw(Exception("Error")),
         )
 
-        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>", 1000)
+        result = restore_imap_emails.email_exists_in_folder(mock_conn, "<test@test.com>")
         assert result is False
 
 
+class TestRestoreProgressCache:
+    def test_progress_cache_add_get_persist(self, tmp_path):
+        import threading
+
+        cache_path = restore_cache.get_dest_index_cache_path(str(tmp_path), "imap.example.com", "user@example.com")
+        cache_data = restore_cache.load_dest_index_cache(cache_path)
+        lock = threading.Lock()
+
+        assert (
+            restore_cache.get_cached_message_ids(cache_data, lock, "imap.example.com", "user@example.com", "INBOX")
+            == set()
+        )
+
+        assert (
+            restore_cache.add_cached_message_id(
+                cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<a@test>"
+            )
+            is True
+        )
+        assert (
+            restore_cache.add_cached_message_id(
+                cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<a@test>"
+            )
+            is False
+        )
+        assert (
+            restore_cache.add_cached_message_id(
+                cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<b@test>"
+            )
+            is True
+        )
+
+        restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, lock, force=True)
+
+        cache_data2 = restore_cache.load_dest_index_cache(cache_path)
+        ids = restore_cache.get_cached_message_ids(cache_data2, lock, "imap.example.com", "user@example.com", "INBOX")
+        assert "<a@test>" in ids
+        assert "<b@test>" in ids
+
+
 class TestGetLabelsFromManifest:
     """Tests for get_labels_from_manifest function."""
 
@@ -497,6 +571,41 @@ class TestLabelToFolder:
         assert result == "Projects/2024"
 
 
+class TestGmailModeDraftsFallbackRegression:
+    def test_gmail_mode_no_labels_does_not_upload_to_drafts(self, monkeypatch):
+        """Regression: messages with no usable labels must not be uploaded to Gmail Drafts."""
+
+        captured = {}
+
+        def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True):
+            captured["folder_name"] = folder_name
+            return True
+
+        def fake_parse_eml_file(_path):
+            return (
+                "<no-labels@test>",
+                '"01-Jan-2024 00:00:00 +0000"',
+                b"Subject: X\r\nMessage-ID: <no-labels@test>\r\n\r\nBody",
+                "X",
+            )
+
+        monkeypatch.setattr(restore_imap_emails, "get_thread_connection", lambda _conf: MagicMock())
+        monkeypatch.setattr(restore_imap_emails, "upload_email", fake_upload_email)
+        monkeypatch.setattr(restore_imap_emails, "parse_eml_file", fake_parse_eml_file)
+
+        restore_imap_emails.process_restore_batch(
+            eml_files=[("/does/not/matter.eml", "x.eml")],
+            folder_name="__GMAIL_MODE__",
+            dest_conf=("host", "user", "pass"),
+            manifest={},
+            apply_labels=True,
+            apply_flags=False,
+        )
+
+        assert captured["folder_name"] == imap_common.FOLDER_RESTORED_UNLABELED
+        assert captured["folder_name"] != "[Gmail]/Drafts"
+
+
 class TestSyncFlagsOnExisting:
     """Tests for sync_flags_on_existing function."""
 

tools/mock_imap_server.py

@@ -2,6 +2,9 @@ import re
 import socketserver
 import threading
 
+RESPONSE_SEARCH_COMPLETED = "OK SEARCH completed"
+RESPONSE_SELECT_FIRST = "NO Select first"
+
 
 class MockIMAPHandler(socketserver.StreamRequestHandler):
     """
@@ -84,12 +87,11 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                         # We use a dictionary-like structure implicitly via dicts in list now
                         # Check "flags" set in msg object
                         new_msgs = [m for m in msgs if "\\Deleted" not in m["flags"]]
-                        removed_count = len(msgs) - len(new_msgs)
                         self.current_folders[self.selected_folder] = new_msgs
                         # According to IMAP, we should send Expunge indices but we'll skip for mock
                         self.send_response(tag, "OK EXPUNGE completed")
                     else:
-                        self.send_response(tag, "NO Select first")
+                        self.send_response(tag, RESPONSE_SELECT_FIRST)
 
                 elif cmd == "UID":
                     sub_parts = args.split(" ", 1)
@@ -99,21 +101,32 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                     if sub_cmd == "SEARCH":
                         sub_args = sub_rest
                         if not self.selected_folder:
-                            self.send_response(tag, "NO Select first")
+                            self.send_response(tag, RESPONSE_SELECT_FIRST)
                             continue
                         msgs = self.current_folders[self.selected_folder]
 
+                        # Support UID range searches like: UID SEARCH UID 101:*
+                        uid_min = None
+                        try:
+                            m = re.search(r"UID\s+(\d+)", sub_args, re.IGNORECASE)
+                            if m:
+                                uid_min = int(m.group(1))
+                        except Exception:
+                            uid_min = None
+
                         # Handle UNDELETED
                         # If args has UNDELETED, filter flags
                         valid_uids = []
                         for m in msgs:
+                            if uid_min is not None and m["uid"] < uid_min:
+                                continue
                             if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]:
                                 continue
                             valid_uids.append(str(m["uid"]))
 
                         uids_str = " ".join(valid_uids)
                         self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
-                        self.send_response(tag, "OK SEARCH completed")
+                        self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
 
                     elif sub_cmd == "STORE":
                         # UID STORE <uid> +FLAGS (\Deleted)
@@ -148,7 +161,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                             else:
                                 self.send_response(tag, "NO UID not found")
                         else:
-                            self.send_response(tag, "NO Select first")
+                            self.send_response(tag, RESPONSE_SELECT_FIRST)
 
                     elif sub_cmd == "FETCH":
                         # sub_rest is e.g. "1 (RFC822.SIZE BODY.PEEK[...])"
@@ -157,7 +170,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                         opts = parts[1].upper() if len(parts) > 1 else ""
 
                         if not self.selected_folder:
-                            self.send_response(tag, "NO Select first")
+                            self.send_response(tag, RESPONSE_SELECT_FIRST)
                             continue
 
                         msgs = self.current_folders[self.selected_folder]
@@ -172,13 +185,13 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                             try:
                                 uid_list = [int(u) for u in uid_set.split(",")]
                                 target_msgs = [m for m in msgs if m["uid"] in uid_list]
-                            except:
+                            except Exception:
                                 pass
                         else:
                             try:
                                 t_uid = int(uid_set)
                                 target_msgs = [m for m in msgs if m["uid"] == t_uid]
-                            except:
+                            except Exception:
                                 pass
 
                         for m in target_msgs:
@@ -252,13 +265,22 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                             self.wfile.write(b"+ Ready\r\n")
                             data = self.rfile.read(size)
 
+                            # Args format (typical): <folder> (<flags>) "<internaldate>" {<size>}
+                            # We only care about folder + flags for tests.
                             folder_arg = args.split(" ")[0].strip().strip('"')
+                            flags_match = re.search(r"\(([^)]*)\)", args)
+                            flags_set = set()
+                            if flags_match:
+                                flags_str = flags_match.group(1).strip()
+                                if flags_str:
+                                    flags_set = {f.strip() for f in flags_str.split() if f.strip()}
+
                             if folder_arg in self.current_folders:
                                 dest_msgs = self.current_folders[folder_arg]
                                 max_uid = max([m["uid"] for m in dest_msgs], default=0)
 
                                 # IMPORTANT: Reset pointer or copy
-                                new_msg = {"uid": max_uid + 1, "flags": set(), "content": data}
+                                new_msg = {"uid": max_uid + 1, "flags": flags_set, "content": data}
                                 dest_msgs.append(new_msg)
                                 print(f"MOCK APPEND SUCCESS: {folder_arg} now has {len(dest_msgs)}")
                                 self.send_response(tag, "OK APPEND completed")
@@ -278,7 +300,7 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                     if not self.selected_folder:
                         self.wfile.write(b"* SEARCH\r\n")
-                        self.send_response(tag, "OK SEARCH completed")
+                        self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
                         continue
 
                     msgs = self.current_folders[self.selected_folder]
@@ -312,12 +334,50 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                     indices_str = " ".join(found_indices)
                     self.wfile.write(f"* SEARCH {indices_str}\r\n".encode())
-                    self.send_response(tag, "OK SEARCH completed")
+                    self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
+
+                elif cmd == "STORE":
+                    # STORE <msg_set> +FLAGS (\Seen)  (non-UID; uses message sequence numbers)
+                    if not self.selected_folder:
+                        self.send_response(tag, RESPONSE_SELECT_FIRST)
+                        continue
+
+                    try:
+                        store_parts = args.split(" ", 2)
+                        msg_set = store_parts[0]
+                        action = store_parts[1].upper()  # +FLAGS or -FLAGS
+                        flags_str = store_parts[2].strip().strip("()")
+                        flags_list = {f.strip() for f in flags_str.split() if f.strip()}
+
+                        msgs = self.current_folders[self.selected_folder]
+
+                        # Only implement single message number for now (sufficient for tests)
+                        try:
+                            msg_num = int(msg_set)
+                        except Exception:
+                            self.send_response(tag, "BAD STORE")
+                            continue
+
+                        if msg_num < 1 or msg_num > len(msgs):
+                            self.send_response(tag, "NO Message not found")
+                            continue
+
+                        m = msgs[msg_num - 1]
+                        if action == "+FLAGS":
+                            m["flags"].update(flags_list)
+                        elif action == "-FLAGS":
+                            m["flags"].difference_update(flags_list)
+
+                        flag_output = " ".join(m["flags"])
+                        self.wfile.write(f"* {msg_num} FETCH (FLAGS ({flag_output}))\r\n".encode())
+                        self.send_response(tag, "OK STORE completed")
+                    except Exception:
+                        self.send_response(tag, "BAD STORE")
 
                 elif cmd == "FETCH":
                     # Non-UID FETCH - args is e.g. "1 (RFC822.SIZE)"
                     if not self.selected_folder:
-                        self.send_response(tag, "NO Select first")
+                        self.send_response(tag, RESPONSE_SELECT_FIRST)
                         continue
 
                     parts = args.split(" ", 1)
@@ -362,8 +422,8 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
     allow_reuse_address = True
     daemon_threads = True
 
-    def __init__(self, server_address, RequestHandlerClass, initial_folders=None):
-        super().__init__(server_address, RequestHandlerClass)
+    def __init__(self, server_address, request_handler_class, initial_folders=None):
+        super().__init__(server_address, request_handler_class)
         self.folders = {}
         if initial_folders:
             for fname, contents in initial_folders.items():