Fix import path for CI build check (#41)

* Add comprehensive tests for IMAP email restoration functionality

- Implement tests for loading labels manifest, including both old and new formats.
- Add tests for extracting flags from the manifest.
- Create tests for parsing .eml files, including edge cases.
- Develop tests for retrieving .eml files from a directory.
- Validate configuration settings and handle missing credentials.
- Integrate tests for restoring emails from a backup, including Gmail mode and deletion of orphaned emails.
- Test progress caching and backup folder discovery.
- Ensure robust handling of email labels and flags during restoration.
- Verify that the restore process correctly interacts with a mock IMAP server.

* Add comprehensive tests for IMAP utilities and compression support

- Introduced tests for `imap_common.py` covering environment variable verification, IMAP connection handling, folder name normalization, MIME header decoding, message details extraction, duplicate detection, filename sanitization, and trash folder detection.
- Added tests for `imap_compress.py` to validate the functionality of the `_CompressedSocket` wrapper, ensuring data compression and decompression works as expected.
- Implemented integration tests to verify that compression is enabled during IMAP connection setup.

.github/workflows/ci.yml

@@ -153,5 +153,5 @@ jobs:
 
       - name: Check imports
         run: |
-          cd src && python -c "import imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
+          cd src && python -c "import utils.imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
           echo "All imports successful"

.github/workflows/publish.yml

@@ -15,6 +15,14 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: "3.x"
+    
+    - name: Update version to match release
+      run: |
+        VERSION=${GITHUB_REF_NAME#v}
+        echo "Updating version to $VERSION"
+        sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml
+        grep "^version =" pyproject.toml
+        
     - name: Install pypa/build
       run: >-
         python3 -m pip install build --user

README.md

@@ -16,7 +16,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### The Scripts
 
-1. **`migrate_imap_emails.py`** (The Solution)
+1. **`imap_migrate.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 skips it if found.
@@ -28,17 +28,17 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Incremental Cache**: Supports `--migrate-cache <path>` to store a local map of processed emails, drastically speeding up subsequent runs by skipping known messages.
    - **Configurable**: Adjustable concurrency and batch sizes to respect server rate limits.
 
-2. **`compare_imap_folders.py`** (The Validator)
+2. **`imap_compare.py`** (The Validator)
    - Connects to both Source and Destination accounts.
    - Prints a side-by-side comparison table of message counts for every folder.
   - 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)
+3. **`imap_count.py`** (The Investigator)
   - 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)
+4. **`imap_backup.py`** (The Backup)
    - Downloads emails from an IMAP account to a local disk.
    - **Format**: Saves emails as individual `.eml` files (RFC 5322), compatible with Outlook, Thunderbird, and Apple Mail.
    - **Structure**: Replicates the IMAP folder hierarchy locally.
@@ -46,7 +46,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Sync Mode**: Optionally deletes local `.eml` files that no longer exist on the server (`--dest-delete`).
    - **Gmail Labels Preservation**: Creates a `labels_manifest.json` file mapping each email's Message-ID to its Gmail labels, enabling proper restoration with labels intact.
 
-5. **`restore_imap_emails.py`** (The Restore)
+5. **`imap_restore.py`** (The Restore)
    - 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.
@@ -67,7 +67,36 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### 2. Installation
 
-Clone this repository or download the script files to your local machine.
+#### Installation via PyPI (Recommended)
+You can install the tools directly from PyPI.
+
+**For macOS Users (and other externally managed environments):**
+Due to recent changes in Python environments on macOS, it is recommended to use `pipx` to install the tools globally without conflicting with system packages.
+
+1. Install `pipx` (if not already installed):
+   ```bash
+   brew install pipx
+   pipx ensurepath
+   ```
+2. Install the tools:
+   ```bash
+   pipx install imap-migration-tools
+   ```
+
+**For Standard Environments:**
+```bash
+pip install imap-migration-tools
+```
+
+Once installed via PyPI, the following commands are globally available in your terminal:
+- `imap-backup`
+- `imap-compare`
+- `imap-count`
+- `imap-migrate`
+- `imap-restore`
+
+#### Installation from Source
+Alternatively, clone this repository or download the script files to your local machine.
 
 #### macOS
 macOS often comes with Python, but it's best to install the latest version.
@@ -129,7 +158,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```bash
-   python3 migrate_imap_emails.py
+   python3 imap_migrate.py
    ```
 
 #### Windows (PowerShell)
@@ -146,7 +175,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```powershell
-   python migrate_imap_emails.py
+   python imap_migrate.py
    ```
 
 ### Method 2: Command Line Arguments (Overrides)
@@ -154,7 +183,7 @@ All scripts support command-line arguments which take precedence over environmen
 
 **Migration:**
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -167,7 +196,7 @@ python3 migrate_imap_emails.py \
 
 **Comparison:**
 ```bash
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -176,14 +205,14 @@ python3 compare_imap_folders.py \
   --dest-pass "your-app-password"
 
 # Compare IMAP source to a local backup folder
-python3 compare_imap_folders.py \
+python3 imap_compare.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 \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -192,23 +221,23 @@ python3 compare_imap_folders.py \
 
 **Counting:**
 ```bash
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --pass "secret"
 
 # Count a local backup folder
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --path "./my_backup"
 
 # Or via environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 count_imap_emails.py
+python3 imap_count.py
 ```
 
 **Counting (OAuth2):**
 ```bash
-python3 count_imap_emails.py \
+python3 imap_count.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --oauth2-client-id "id" \
@@ -219,12 +248,12 @@ 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
+python3 imap_count.py
 ```
 
 **Backup:**
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -236,7 +265,7 @@ python3 backup_imap_emails.py \
 ### 1. Full Migration
 Migrate all folders from Source to Destination.
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -249,7 +278,7 @@ python3 migrate_imap_emails.py \
 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 \
+python3 imap_migrate.py \
   --gmail-mode \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
@@ -263,7 +292,7 @@ python3 migrate_imap_emails.py \
 Use a local cache file to remember processed emails. Use this for large migrations that may be interrupted or need to run multiple times.
 
 ```bash
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.source.com" \
   --src-user "source" \
   --src-pass "pass" \
@@ -281,7 +310,7 @@ Preserve IMAP flags (`\Seen`, `\Flagged`, `\Answered`, `\Draft`) during migratio
 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 \
+python3 imap_migrate.py \
   --preserve-flags \
   --src-host "imap.example.com" \
   --src-user "source@example.com" \
@@ -294,8 +323,8 @@ python3 migrate_imap_emails.py \
 ### 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 \
+# Syntax: python3 imap_migrate.py "[Folder Name]"
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -309,7 +338,7 @@ python3 migrate_imap_emails.py \
 Migrate and **delete** from source immediately after verifying the copy.
 ```bash
 # Using flag
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -319,7 +348,7 @@ python3 migrate_imap_emails.py \
   --src-delete
 
 # Or specific folder with delete
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -336,7 +365,7 @@ 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 \
+python3 imap_migrate.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -346,7 +375,7 @@ python3 migrate_imap_emails.py \
   --dest-delete
 
 # Backup: Delete local .eml files not found on server
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -354,7 +383,7 @@ python3 backup_imap_emails.py \
   --dest-delete
 
 # Restore: Delete server emails not found in local backup
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -367,7 +396,7 @@ python3 restore_imap_emails.py \
 ### 5. Verify Migration
 Compare counts between source and destination.
 ```bash
-python3 compare_imap_folders.py \
+python3 imap_compare.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -376,14 +405,14 @@ python3 compare_imap_folders.py \
   --dest-pass "dest-app-password"
 
 # IMAP source -> local backup destination
-python3 compare_imap_folders.py \
+python3 imap_compare.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 \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -401,21 +430,21 @@ INBOX                   | 1250         | 1250       | MATCH
 Download all your emails to your computer as `.eml` files.
 ```bash
 # Backup all folders from an IMAP account to a local folder
-python3 backup_imap_emails.py \
+python3 imap_backup.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 \
+python3 imap_backup.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 \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -424,18 +453,18 @@ python3 backup_imap_emails.py \
 ```
 
 ### 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`.
+Use `imap_compare.py` to validate an IMAP account against a local backup created by `imap_backup.py`.
 
 ```bash
 # Option 1: IMAP source -> local destination
-python3 compare_imap_folders.py \
+python3 imap_compare.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 \
+python3 imap_compare.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -450,15 +479,15 @@ 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`.
+Use `imap_count.py` to get per-folder counts from a local backup created by `imap_backup.py`.
 
 ```bash
 # Option 1: explicit path
-python3 count_imap_emails.py --path "./my_backup"
+python3 imap_count.py --path "./my_backup"
 
 # Option 2: environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 count_imap_emails.py
+python3 imap_count.py
 ```
 
 ### 7. Gmail Backup with Labels Preservation
@@ -466,7 +495,7 @@ When backing up a Gmail account, use `--gmail-mode` for the recommended workflow
 
 ```bash
 # Recommended: Use --gmail-mode for simplest workflow
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -476,7 +505,7 @@ python3 backup_imap_emails.py \
 
 This is equivalent to the more verbose:
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -488,7 +517,7 @@ python3 backup_imap_emails.py \
 **For large accounts (100K+ emails)**, you can build the manifest first to test:
 ```bash
 # Step 1: Build manifest only (fast, no download)
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -496,7 +525,7 @@ python3 backup_imap_emails.py \
   --manifest-only
 
 # Step 2: Download emails (can run later, manifest already exists)
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -542,7 +571,7 @@ python3 backup_imap_emails.py \
 For non-Gmail servers, you can preserve read/starred status with `--preserve-flags`:
 
 ```bash
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.example.com" \
   --src-user "you@example.com" \
   --src-pass "your-password" \
@@ -561,14 +590,14 @@ Use `--full-restore` to force the legacy behavior (process all emails and re-syn
 
 ```bash
 # Restore all folders from backup
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
   --dest-pass "your-app-password"
 
 # Force full restore (legacy behavior)
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -576,7 +605,7 @@ python3 restore_imap_emails.py \
   --full-restore
 
 # Restore with flags (read/starred status)
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.example.com" \
   --dest-user "you@example.com" \
@@ -584,7 +613,7 @@ python3 restore_imap_emails.py \
   --apply-flags
 
 # Restore a specific folder
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -596,7 +625,7 @@ python3 restore_imap_emails.py \
 Restore a Gmail backup with full label structure using `--gmail-mode`:
 
 ```bash
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -613,7 +642,7 @@ python3 restore_imap_emails.py \
 
 **Alternatively**, restore folders individually with labels and flags applied:
 ```bash
-python3 restore_imap_emails.py \
+python3 imap_restore.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -635,7 +664,7 @@ To use OAuth2, pass `--oauth2-client-id` (or `--src-oauth2-client-id`/`--dest-oa
 
 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`).
+- Single-account scripts (like `imap_count.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
 
 ### Microsoft (Outlook / Office 365)
 
@@ -674,7 +703,7 @@ Requires the `msal` package (`pip install msal`). Uses the **device code flow**
 pip install msal
 
 # Migration with Microsoft OAuth2 on source
-python3 migrate_imap_emails.py \
+python3 imap_migrate.py \
   --src-host "outlook.office365.com" \
   --src-user "user@contoso.com" \
   --src-oauth2-client-id "your-azure-app-client-id" \
@@ -694,7 +723,7 @@ Requires the `google-auth-oauthlib` package (`pip install google-auth-oauthlib`)
 pip install google-auth-oauthlib
 
 # Backup with Google OAuth2
-python3 backup_imap_emails.py \
+python3 imap_backup.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-oauth2-client-id "your-google-client-id" \
@@ -707,7 +736,7 @@ The script will open your default browser for Google sign-in. After authorizing,
 ## Troubleshooting
 
 - **"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.
+  IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `imap_migrate.py` uses multiple threads, you may hit this limit.
   **Solution**: Reduce `MAX_WORKERS` to `4` or `2` using the environment variable.
 
 - **Authentication Errors**:
@@ -760,7 +789,7 @@ PYTHONPATH=src pytest test/ -v
 make coverage
 
 # Run a specific test file
-PYTHONPATH=src pytest test/test_migrate_imap_emails.py -v
+PYTHONPATH=src pytest test/test_imap_migrate.py -v
 
 # Run a specific test
 PYTHONPATH=src pytest test/test_imap_common.py::TestNormalizeFolderName -v
@@ -792,11 +821,11 @@ make ci
 
 | Test File | Description |
 |-----------|-------------|
-| `test_migrate_imap_emails.py` | Email migration tests (basic, duplicates, deletion, folders) |
-| `test_backup_imap_emails.py` | Backup functionality tests |
-| `test_restore_imap_emails.py` | Restore functionality tests |
-| `test_count_imap_emails.py` | Email counting tests |
-| `test_compare_imap_folders.py` | Folder comparison tests |
+| `test_imap_migrate.py` | Email migration tests (basic, duplicates, deletion, folders) |
+| `test_imap_backup.py` | Backup functionality tests |
+| `test_imap_restore.py` | Restore functionality tests |
+| `test_imap_count.py` | Email counting tests |
+| `test_imap_compare.py` | Folder comparison tests |
 | `test_imap_common.py` | Shared utility function tests |
 
 ### Continuous Integration

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imap-migration-tools"
-version = "1.0.0"
+version = "1.1.0"
 authors = [
   { name="Javier Callico", email="jcallico@callicode.com" },
   { name="Nathan Moinvaziri", email="nathan@nathanm.com" }
@@ -30,11 +30,11 @@ dependencies = [
 "Bug Tracker" = "https://github.com/jcallico/imap-migration-tools/issues"
 
 [project.scripts]
-imap-backup = "backup_imap_emails:main"
-imap-restore = "restore_imap_emails:main"
-imap-migrate = "migrate_imap_emails:main"
-imap-count = "count_imap_emails:main"
-imap-compare = "compare_imap_folders:main"
+imap-backup = "imap_backup:main"
+imap-restore = "imap_restore:main"
+imap-migrate = "imap_migrate:main"
+imap-count = "imap_count:main"
+imap-compare = "imap_compare:main"
 
 [tool.setuptools.package-dir]
 "" = "src"

src/auth/imap_oauth2.py

@@ -10,8 +10,7 @@ Provider is auto-detected from the IMAP host string.
 import sys
 import threading
 
-import oauth2_google
-import oauth2_microsoft
+from auth import oauth2_google, oauth2_microsoft
 
 # Re-export caches for test access
 _msal_app_cache = oauth2_microsoft._msal_app_cache

src/backup_imap_emails.py

@@ -1,895 +1,24 @@
+#!/usr/bin/env python3
 """
-IMAP Email Backup Script
-
-Backs up emails from an IMAP account to a local directory.
-Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
-most email clients (Thunderbird, Apple Mail, Outlook, etc.).
-
-Features:
-- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
-- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
-- Folder Replication: Recreates the IMAP folder structure locally.
-- Parallel Processing: Uses multithreading for fast downloads.
-- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
-
-Configuration (Environment Variables):
-    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).
-  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
-  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
-  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
-  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
-  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
-              Default is "false".
-
-Usage:
-    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 \
-        --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.
+DEPRECATED: This script has been renamed.
+Please use imap_backup.py instead.
 """
 
-import argparse
-import concurrent.futures
-import json
-import os
 import sys
-import threading
 
-import imap_common
-import imap_oauth2
-import imap_session
-import provider_exchange
-import provider_gmail
+from imap_backup import main
 
-# Defaults
-MAX_WORKERS = 10
-BATCH_SIZE = 10
-MANIFEST_FILENAME = "labels_manifest.json"
-
-# Thread-local storage
-thread_local = threading.local()
-safe_print = imap_common.safe_print
-
-
-def process_single_uid(src, uid, folder_name, local_folder_path):
-    """
-    Fetch and save a single email by UID.
-
-    Args:
-        src: IMAP connection
-        uid: Message UID to fetch
-        folder_name: Current folder name (for logging)
-        local_folder_path: Local directory to save email
-
-    Returns:
-        Tuple of (success: bool, connection):
-        - (True, src): UID processed successfully (or skipped)
-        - (False, src): Auth error occurred, caller should retry after reconnect
-    """
-    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-
-    try:
-        resp, data = src.uid("fetch", uid, "(RFC822)")
-        if resp != "OK" or not data or data[0] is None:
-            safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
-            return (True, src)  # Don't retry, move on
-
-        raw_email = None
-        for item in data:
-            if isinstance(item, tuple):
-                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)
-            clean_subject = clean_subject[:100]
-
-        filename = f"{uid_str}_{clean_subject}.eml"
-        full_path = os.path.join(local_folder_path, filename)
-
-        if os.path.exists(full_path):
-            return (True, src)  # Already exists
-
-        if raw_email:
-            try:
-                with open(full_path, "wb") as f:
-                    f.write(raw_email)
-                safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
-            except OSError as e:
-                safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
-        else:
-            safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
-
-        return (True, src)
-
-    except Exception as e:
-        if imap_oauth2.is_auth_error(e):
-            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
-            return (False, src)  # Signal retry needed
-        else:
-            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
-            return (True, src)  # Don't retry other errors
-
-
-def process_batch(uids, folder_name, src_conf, local_folder_path):
-    src = imap_session.get_thread_connection(thread_local, "src", src_conf)
-    if not src:
-        safe_print("Error: Could not establish connection for batch.")
-        return
-
-    try:
-        src.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
-        return
-
-    for uid in uids:
-        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-        max_retries = 2
-
-        for attempt in range(max_retries):
-            src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
-            thread_local.src = src
-            if not ok:
-                safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
-                return
-
-            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
-            thread_local.src = src
-
-            if success:
-                break
-            if attempt < max_retries - 1:
-                src = None
-                thread_local.src = None
-
-
-def get_existing_uids(local_path):
-    """
-    Scans the local directory for files matching pattern matches {UID}_*.eml
-    Returns a set of UIDs (as strings).
-    """
-    existing = set()
-    if not os.path.exists(local_path):
-        return existing
-
-    try:
-        for filename in os.listdir(local_path):
-            if filename.endswith(".eml") and "_" in filename:
-                # Expecting UID_Subject.eml
-                parts = filename.split("_", 1)
-                if parts[0].isdigit():
-                    existing.add(parts[0])
-    except Exception:
-        pass
-    return existing
-
-
-# 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 = imap_common.PRESERVABLE_FLAGS
-
-
-def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
-    """
-    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
-    with OAuth2 session management.
-
-    Args:
-        imap_conn: IMAP connection
-        folder_name: Folder to scan
-        src_conf: Connection config dict for OAuth2 refresh
-        progress_callback: Optional callback(current, total) for progress reporting
-
-    Returns:
-        Tuple of (message_info, imap_conn) where message_info is:
-        { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
-        The returned imap_conn may be different if reconnection occurred.
-    """
-    message_info = {}
-
-    # Ensure connection is healthy (refresh OAuth2 token if needed) before initial select
-    if src_conf:
-        imap_conn = imap_session.ensure_connection(imap_conn, src_conf)
-        if not imap_conn:
-            safe_print(f"Could not establish connection for folder {folder_name}")
-            return (message_info, None)
-
-    try:
-        imap_conn.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Could not select folder {folder_name}: {e}")
-        return (message_info, imap_conn)
-
-    try:
-        resp, data = imap_conn.uid("search", None, "ALL")
-        if resp != "OK" or not data or not data[0]:
-            return (message_info, imap_conn)
-
-        uids = data[0].split()
-        if not uids:
-            return (message_info, imap_conn)
-
-        total_uids = len(uids)
-
-        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
-        batch_size = 200
-        for i in range(0, len(uids), batch_size):
-            batch = uids[i : i + batch_size]
-            uid_range = b",".join(batch)
-
-            # Report progress
-            if progress_callback:
-                progress_callback(min(i + batch_size, total_uids), total_uids)
-
-            # Proactively refresh token and ensure folder is selected
-            if src_conf:
-                imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
-                if not ok:
-                    safe_print(f"ERROR: Connection/folder lost in {folder_name}")
-                    break
-
-            try:
-                # Fetch both Message-ID header and FLAGS
-                resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
-                if resp != "OK":
-                    continue
-
-                # Parse response - items come in pairs for each message
-                for item in items:
-                    if isinstance(item, tuple) and len(item) >= 2:
-                        # First element contains UID and FLAGS info
-                        meta_str = (
-                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
-                        )
-
-                        # Extract all preservable flags from the metadata
-                        flags = []
-                        for flag in imap_common.PRESERVABLE_FLAGS:
-                            if flag in meta_str:
-                                flags.append(flag)
-
-                        # Second element contains the header
-                        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}")
-                continue
-
-    except Exception as e:
-        safe_print(f"Error searching folder {folder_name}: {e}")
-
-    return (message_info, imap_conn)
-
-
-def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
-    """
-    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
-
-    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
-    Optional progress_callback(current, total) for progress reporting.
-    """
-    message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
-    return message_info
-
-
-def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
-    """
-    Returns a set of Message-IDs for all emails in a given folder,
-    with OAuth2 session management.
-
-    Args:
-        imap_conn: IMAP connection
-        folder_name: Folder to scan
-        src_conf: Connection config dict for OAuth2 refresh
-        progress_callback: Optional callback(current, total) for progress reporting
-
-    Returns:
-        Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
-        The returned imap_conn may be different if reconnection occurred.
-    """
-    info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
-    return (set(info.keys()), imap_conn)
-
-
-def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
-    """
-    Returns a set of Message-IDs for all emails in a given folder.
-    This is a convenience wrapper around get_message_info_in_folder.
-    Optional progress_callback(current, total) for progress reporting.
-    """
-    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
-    return set(info.keys())
-
-
-def build_labels_manifest(imap_conn, local_path, src_conf=None):
-    """
-    Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
-    Scans all folders (labels) in the account and records which Message-IDs
-    appear in each label, plus their flags from [Gmail]/All Mail.
-
-    Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
-    Saves the manifest to labels_manifest.json in the backup directory.
-    Optional src_conf for automatic token refresh on expiration.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-    total_emails_scanned = 0
-
-    safe_print("--- Building Gmail Labels Manifest ---")
-
-    # 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
-    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
-    all_mail_start = time.time()
-
-    def all_mail_progress_cb(current, total):
-        elapsed = time.time() - start_time
-        print(
-            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-            end="",
-            flush=True,
-        )
-
-    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
-        imap_conn, provider_gmail.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
+if __name__ == "__main__":
+    print(
+        "WARNING: 'backup_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_backup.py'.",
+        file=sys.stderr,
     )
-    print()  # New line after progress
-
-    # Initialize manifest with flags from All Mail
-    for msg_id, info in all_mail_info.items():
-        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
-
-    all_mail_elapsed = time.time() - all_mail_start
-    # Count flag statistics
-    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")
-
-    # Keep connection alive
+    print("Redirecting to 'imap_backup.py'...", file=sys.stderr)
     try:
-        imap_conn.noop()
-    except Exception:
-        pass
-
-    # Parse folder names and filter to label folders
-    label_folders = [f for f in all_folders if provider_gmail.is_label_folder(f)]
-
-    total_folders = len(label_folders)
-    safe_print(f"Found {total_folders} label folders to scan.\n")
-
-    # Scan each label folder
-    for folder_idx, folder_name in enumerate(label_folders, 1):
-        folder_start = time.time()
-
-        # Progress callback for this folder
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
-        print()  # New line after progress
-
-        # Keep connection alive between folders
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-        # Determine the label name to store
-        # For [Gmail]/Sent Mail -> "Sent Mail"
-        # For [Gmail]/Starred -> "Starred"
-        # For INBOX -> "INBOX"
-        # For user folders -> folder name as-is
-        if folder_name.startswith("[Gmail]/"):
-            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
-        else:
-            label_name = folder_name
-
-        for msg_id in message_ids:
-            if msg_id not in manifest:
-                # Email not in All Mail (rare, but handle it)
-                manifest[msg_id] = {"labels": [], "flags": []}
-            if label_name not in manifest[msg_id]["labels"]:
-                manifest[msg_id]["labels"].append(label_name)
-
-        folder_elapsed = time.time() - folder_start
-        total_emails_scanned += len(message_ids)
-        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    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}")
-    safe_print(f"  - Unique emails: {len(manifest)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    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)
-        safe_print(f"\nLabels manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
-    """
-    Builds a manifest mapping Message-IDs to their IMAP flags.
-    For non-Gmail servers, scans specified folders (or all folders if not specified).
-
-    Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
-    Saves the manifest to flags_manifest.json in the backup directory.
-    Optional src_conf for automatic token refresh on expiration.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-
-    safe_print("--- Building Flags Manifest ---")
-
-    # Get folders to scan
-    if folders_to_scan:
-        all_folders = folders_to_scan
-    else:
-        try:
-            typ, folders = imap_conn.list()
-            if typ != "OK":
-                safe_print("Error: Could not list folders.")
-                return manifest
-            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
-        except Exception as e:
-            safe_print(f"Error listing folders: {e}")
-            return manifest
-
-    total_folders = len(all_folders)
-    safe_print(f"Found {total_folders} folders to scan.\n")
-
-    # Scan each folder
-    for folder_idx, folder_name in enumerate(all_folders, 1):
-        folder_start = time.time()
-
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
-        print()  # New line after progress
-
-        # Merge into manifest (keep first occurrence of flags)
-        for msg_id, info in folder_info.items():
-            if msg_id not in manifest:
-                manifest[msg_id] = {"flags": info.get("flags", [])}
-
-        folder_elapsed = time.time() - folder_start
-        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
-
-        # Keep connection alive
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    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)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    manifest_path = os.path.join(local_path, "flags_manifest.json")
-    try:
-        with open(manifest_path, "w", encoding="utf-8") as f:
-            json.dump(manifest, f, indent=2, ensure_ascii=False)
-        safe_print(f"\nFlags manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def delete_orphan_local_files(local_folder_path, server_uids):
-    """
-    Delete local .eml files that no longer exist on the server.
-    Args:
-        local_folder_path: Path to local folder containing .eml files
-        server_uids: Set of UID strings currently on server
-    Returns:
-        Count of deleted files
-    """
-    deleted_count = 0
-    if not os.path.exists(local_folder_path):
-        return deleted_count
-
-    try:
-        for filename in os.listdir(local_folder_path):
-            if not filename.endswith(".eml") or "_" not in filename:
-                continue
-
-            # Extract UID from filename (format: {UID}_{Subject}.eml)
-            parts = filename.split("_", 1)
-            if not parts[0].isdigit():
-                continue
-
-            local_uid = parts[0]
-            if local_uid not in server_uids:
-                file_path = os.path.join(local_folder_path, filename)
-                try:
-                    os.remove(file_path)
-                    safe_print(f"  -> Deleted orphan: {filename}")
-                    deleted_count += 1
-                except Exception as e:
-                    safe_print(f"  -> Error deleting {filename}: {e}")
-    except Exception as e:
-        safe_print(f"Error scanning for orphan files: {e}")
-
-    return deleted_count
-
-
-def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
-    safe_print(f"--- Processing Folder: {folder_name} ---")
-
-    # create local path
-    # Handle folder separators. IMAP output might be "Parent/Child"
-    # We rely on OS to handle "Parent/Child" as subdirectories using join
-    # But clean the segments
-    cleaned_name = folder_name.replace("/", os.sep)
-    local_folder_path = os.path.join(local_base_path, cleaned_name)
-
-    try:
-        os.makedirs(local_folder_path, exist_ok=True)
-    except Exception as e:
-        safe_print(f"Error creating directory {local_folder_path}: {e}")
-        return
-
-    # Select IMAP folder
-    try:
-        src_main.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Skipping {folder_name}: {e}")
-        return
-
-    # Search all
-    resp, data = src_main.uid("search", None, "ALL")
-    if resp != "OK":
-        return
-
-    uids = data[0].split()
-    total_on_server = len(uids)
-
-    # Build set of server UIDs for comparison
-    server_uid_set = set()
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        server_uid_set.add(u_str)
-
-    if total_on_server == 0:
-        safe_print(f"Folder {folder_name} is empty.")
-        # If dest_delete enabled, delete all local files
-        if dest_delete:
-            deleted = delete_orphan_local_files(local_folder_path, set())
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-        return
-
-    # Incremental Optimization
-    # Read local directory to find UIDs we already have
-    existing_uids = get_existing_uids(local_folder_path)
-
-    # Delete orphan local files if dest_delete is enabled
-    if dest_delete:
-        orphan_uids = existing_uids - server_uid_set
-        if orphan_uids:
-            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
-            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-
-    # Filter UIDs
-    # decode uid first if bytes
-    uids_to_download = []
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        if u_str not in existing_uids:
-            uids_to_download.append(u)
-
-    skipped = total_on_server - len(uids_to_download)
-    if skipped > 0:
-        safe_print(f"Skipping {skipped} emails (already exist locally).")
-
-    if not uids_to_download:
-        safe_print(f"Folder {folder_name} is up to date.")
-        return
-
-    safe_print(f"Downloading {len(uids_to_download)} new emails...")
-
-    # Create batches
-    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
-
-    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
-    try:
-        futures = []
-        for batch in uid_batches:
-            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
-
-        for future in concurrent.futures.as_completed(futures):
-            future.result()
-
+        main()
     except KeyboardInterrupt:
-        executor.shutdown(wait=False, cancel_futures=True)
-        raise
-    finally:
-        executor.shutdown(wait=True)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
-
-    # Source
-    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,
-        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")
-    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
-
-    # Gmail Labels
-    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-labels",
-        action="store_true",
-        default=env_preserve_labels,
-        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
-    )
-    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-flags",
-        action="store_true",
-        default=env_preserve_flags,
-        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
-    )
-    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
-    parser.add_argument(
-        "--manifest-only",
-        action="store_true",
-        default=env_manifest_only,
-        help="Gmail only: Build the labels manifest and exit without downloading emails",
-    )
-    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
-    parser.add_argument(
-        "--gmail-mode",
-        action="store_true",
-        default=env_gmail_mode,
-        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
-    )
-
-    # Sync mode: delete local files not on server
-    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
-    parser.add_argument(
-        "--dest-delete",
-        action="store_true",
-        default=env_dest_delete,
-        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
-    )
-
-    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
-
-    args = parser.parse_args()
-
-    global MAX_WORKERS, BATCH_SIZE
-    MAX_WORKERS = args.workers
-    BATCH_SIZE = args.batch
-
-    # Build connection config (acquires OAuth2 token if configured)
-    src_conf = imap_session.build_imap_conf(
-        args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
-    )
-
-    # Expand path (~/...)
-    local_path = os.path.expanduser(args.dest_path)
-    if not os.path.exists(local_path):
-        try:
-            os.makedirs(local_path)
-            print(f"Created backup directory: {local_path}")
-        except Exception as e:
-            print(f"Error creating backup directory: {e}")
-            sys.exit(1)
-
-    print("\n--- Configuration Summary ---")
-    print(f"Source Host     : {args.src_host}")
-    print(f"Source User     : {args.src_user}")
-    print(f"Auth Method     : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
-    print(f"Destination Path: {local_path}")
-    if args.gmail_mode:
-        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
-    elif args.manifest_only:
-        print("Mode            : Manifest Only (no email download)")
-    elif args.folder:
-        print(f"Target Folder   : {args.folder}")
-    if args.preserve_labels or args.manifest_only or args.gmail_mode:
-        print("Preserve Labels : Yes (Gmail)")
-    if args.preserve_flags or args.gmail_mode:
-        print("Preserve Flags  : Yes (read/starred/answered/draft)")
-    if args.dest_delete:
-        print("Dest Delete     : Yes (remove local orphans)")
-    print("-----------------------------\n")
-
-    try:
-        src = imap_common.get_imap_connection_from_conf(src_conf)
-        if not src:
-            sys.exit(1)
-
-        # Build labels manifest BEFORE backing up emails (Gmail mode)
-        # This way we capture the label state at backup time
-        if args.preserve_labels or args.manifest_only or args.gmail_mode:
-            print("Building Gmail labels manifest...")
-            print("This scans all folders to map Message-IDs to labels and flags.\n")
-            build_labels_manifest(src, local_path, src_conf)
-            print("")  # Blank line after manifest building
-        # Build flags-only manifest for non-Gmail servers
-        elif args.preserve_flags:
-            print("Building flags manifest...")
-            print("This scans folders to capture read/starred/etc status.\n")
-            # Get folders to scan
-            folders_to_scan = [args.folder] if args.folder else None
-            build_flags_manifest(src, local_path, folders_to_scan, src_conf)
-            print("")  # Blank line after manifest building
-
-        # If manifest-only mode, we're done
-        if args.manifest_only:
-            try:
-                src.logout()
-            except Exception:
-                pass  # Connection may already be closed
-            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:")
-            print(f'  python3 backup_imap_emails.py --dest-path "{local_path}" "[Gmail]/All Mail"')
-            sys.exit(0)
-
-        # Gmail mode: backup only [Gmail]/All Mail
-        if args.gmail_mode:
-            backup_folder(src, provider_gmail.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:
-            # Reconnect after potentially long manifest building
-            src = imap_session.ensure_connection(src, src_conf)
-            if not src:
-                print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
-                sys.exit(0)
-            folders = imap_common.list_selectable_folders(src)
-            for name in folders:
-                if provider_exchange.is_special_folder(name):
-                    print(f"Skipping Exchange system folder: {name}")
-                    continue
-                src = imap_session.ensure_connection(src, src_conf)
-                if not src:
-                    print("Fatal: Could not reconnect to IMAP server. Aborting.")
-                    sys.exit(1)
-                backup_folder(src, name, local_path, src_conf, args.dest_delete)
-
-        try:
-            src.logout()
-        except Exception:
-            pass  # Connection may already be closed
-        print("\nBackup completed successfully.")
-
-        if args.preserve_labels or args.gmail_mode:
-            manifest_path = os.path.join(local_path, "labels_manifest.json")
-            print(f"\nGmail labels manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply labels and flags to emails.")
-        elif args.preserve_flags:
-            manifest_path = os.path.join(local_path, "flags_manifest.json")
-            print(f"\nFlags manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply read/starred status to emails.")
-
-    except KeyboardInterrupt:
-        print("\nBackup interrupted by user.")
+        print("\nProcess terminated by user.")
         sys.exit(0)
     except Exception as e:
         print(f"Fatal Error: {e}")
-
-
-if __name__ == "__main__":
-    main()
+        sys.exit(1)

src/compare_imap_folders.py

@@ -1,329 +1,24 @@
+#!/usr/bin/env python3
 """
-IMAP Folder Comparison Script
-
-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.
-
-Configuration (Environment Variables):
-  Source Account:
-    SRC_IMAP_HOST       : Source IMAP Host
-    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"
+DEPRECATED: This script has been renamed.
+Please use imap_compare.py instead.
 """
 
-import argparse
-import os
-
-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
-        typ, data = conn.select(f'"{folder_name}"', readonly=True)
-        if typ != "OK":
-            return None
-
-        # SELECT command returns the number of messages in data[0]
-        # data[0] is bytes, e.g. b'123'
-        if data and data[0]:
-            return int(data[0])
-        return 0
-
-    except Exception:
-        # print(f"Error checking {folder_name}: {e}")
-        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
-    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
-    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()
-
-    src_is_local = bool(args.src_path)
-    dest_is_local = bool(args.dest_path)
-
-    SRC_HOST = args.src_host
-    SRC_USER = args.src_user
-    DEST_HOST = args.dest_host
-    DEST_USER = args.dest_user
-
-    # Acquire OAuth2 tokens if configured
-    src_oauth2_token = None
-    src_oauth2_provider = None
-    if not src_is_local and args.src_client_id:
-        src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
-            SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
-        )
-
-    dest_oauth2_token = None
-    dest_oauth2_provider = None
-    if not dest_is_local and args.dest_client_id:
-        dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
-            DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
-        )
-
-    print("\n--- Configuration Summary ---")
-    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     : {imap_oauth2.auth_description(src_oauth2_provider)}")
-
-    if dest_is_local:
-        print(f"Destination (Local): {args.dest_path}")
-    else:
-        print(f"Destination Host: {args.dest_host}")
-        print(f"Destination User: {args.dest_user}")
-        print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
-    print("-----------------------------\n")
-
-    src = None
-    dest = None
-
-    try:
-        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
-
-        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...")
-        if src_is_local:
-            folders = imap_common.list_local_folders(args.src_path)
-        else:
-            folders = imap_common.list_selectable_folders(src)
-
-        if not folders:
-            print("Failed to list source folders.")
-            return
-
-        # Prepare Table Header
-        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
-        print("-" * len(header))
-        print(header)
-        print("-" * len(header))
-
-        total_src = 0
-        total_dest = 0
-
-        # Iterate through Source folders
-        for folder_name in folders:
-            # Get Counts
-            if src_is_local:
-                src_count = imap_common.get_local_email_count(args.src_path, folder_name)
-            else:
-                src_count = get_email_count(src, folder_name)
-
-            if dest_is_local:
-                dest_count = imap_common.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"
-            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
-
-            diff_str = ""
-            if src_count is not None and dest_count is not None:
-                diff = src_count - dest_count
-                diff_str = str(diff)
-                total_src += src_count
-                total_dest += dest_count
-            elif src_count is not None:
-                total_src += src_count
-
-            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
-
-        print("-" * len(header))
-        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
-
-    except Exception as e:
-        print(f"\nFatal Error: {e}")
-        if "Too many simultaneous connections" in str(e):
-            print("Tip: Wait a few minutes for previous connections to timeout or check other active scripts.")
-
-    finally:
-        # Check source connection state and logout if possible
-        if src:
-            try:
-                src.logout()
-            except Exception:
-                pass
-
-        # Check dest connection state and logout if possible
-        if dest:
-            try:
-                dest.logout()
-            except Exception:
-                pass
+import sys
 
+from imap_compare import main
 
 if __name__ == "__main__":
-    main()
+    print(
+        "WARNING: 'compare_imap_folders.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_compare.py'.",
+        file=sys.stderr,
+    )
+    print("Redirecting to 'imap_compare.py'...", file=sys.stderr)
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/core/imap_session.py

@@ -5,8 +5,8 @@ Connection and session management for IMAP operations with OAuth2 support.
 Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh).
 """
 
-import imap_common
-import imap_oauth2
+from auth import imap_oauth2
+from utils import imap_common
 
 
 def build_imap_conf(host, user, password, client_id=None, client_secret=None, label=None):

src/count_imap_emails.py

@@ -1,239 +1,24 @@
-"""IMAP Email Counting Script.
-
-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)
-
-    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
-
-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
+#!/usr/bin/env python3
+"""
+DEPRECATED: This script has been renamed.
+Please use imap_count.py instead.
 """
 
-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=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, oauth2_token)
-        if not mail:
-            return
-
-        # List all mailboxes
-        print("Listing mailboxes...")
-        folders = imap_common.list_selectable_folders(mail)
-
-        if not folders:
-            print("Failed to list mailboxes.")
-            return
-
-        total_all_folders = 0
-        print(f"{'Folder Name':<40} {'Count':>10}")
-        print("-" * 52)
-
-        for folder_name in folders:
-            display_name = folder_name
-
-            try:
-                # Select the mailbox (read-only is sufficient for counting)
-                # folder_name extracted from list usually handles quotes correctly for select
-                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
-                if rv != "OK":
-                    print(f"{display_name:<40} {'Skipped':>10}")
-                    continue
-
-                # Search for all emails
-                status, data = mail.search(None, "ALL")
-
-                if status == "OK":
-                    # data[0] is space separated IDs
-                    email_ids = data[0].split()
-                    count = len(email_ids)
-                    print(f"{display_name:<40} {count:>10}")
-                    total_all_folders += count
-                else:
-                    print(f"{display_name:<40} {'Error':>10}")
-
-            except imaplib.IMAP4.error:
-                print(f"{display_name:<40} {'Error':>10}")
-
-        print("-" * 52)
-        print(f"{'TOTAL':<40} {total_all_folders:>10}")
-
-        # Logout
-        mail.logout()
-
-    except imaplib.IMAP4.error as e:
-        print(f"IMAP Error: {e}")
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-
-def count_local_emails(local_path: str) -> None:
-    print(f"Scanning local backup: {local_path}")
-
-    folders = imap_common.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 = imap_common.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,
-        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)",
-    )
-
-    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
-
-    # Acquire OAuth2 token if configured
-    oauth2_token = None
-    oauth2_provider = None
-    if args.client_id:
-        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
-            IMAP_SERVER, args.client_id, USERNAME, args.client_secret
-        )
-
-    print("\n--- Configuration Summary ---")
-    print(f"Host            : {IMAP_SERVER}")
-    print(f"User            : {USERNAME}")
-    print(f"Auth Method     : {imap_oauth2.auth_description(oauth2_provider)}")
-    print("-----------------------------\n")
-
-    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
 
+from imap_count import main
 
 if __name__ == "__main__":
-    main()
+    print(
+        "WARNING: 'count_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_count.py'.",
+        file=sys.stderr,
+    )
+    print("Redirecting to 'imap_count.py'...", file=sys.stderr)
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_backup.py

@@ -0,0 +1,899 @@
+"""
+IMAP Email Backup Script
+
+Backs up emails from an IMAP account to a local directory.
+Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
+most email clients (Thunderbird, Apple Mail, Outlook, etc.).
+
+Features:
+- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
+- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
+- Folder Replication: Recreates the IMAP folder structure locally.
+- Parallel Processing: Uses multithreading for fast downloads.
+- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
+
+Configuration (Environment Variables):
+    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).
+  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
+  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
+  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
+  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
+  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
+              Default is "false".
+
+Usage:
+    python3 imap_backup.py \
+        --src-host "imap.example.com" \
+        --src-user "you@example.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup"
+
+Gmail Labels:
+    python3 imap_backup.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.
+"""
+
+import argparse
+import concurrent.futures
+import json
+import os
+import sys
+import threading
+
+from auth import imap_oauth2
+from core import imap_session
+from providers import provider_exchange, provider_gmail
+from utils import imap_common
+
+# Defaults
+MAX_WORKERS = 10
+BATCH_SIZE = 10
+MANIFEST_FILENAME = "labels_manifest.json"
+
+# Thread-local storage
+thread_local = threading.local()
+safe_print = imap_common.safe_print
+
+
+def process_single_uid(src, uid, folder_name, local_folder_path):
+    """
+    Fetch and save a single email by UID.
+
+    Args:
+        src: IMAP connection
+        uid: Message UID to fetch
+        folder_name: Current folder name (for logging)
+        local_folder_path: Local directory to save email
+
+    Returns:
+        Tuple of (success: bool, connection):
+        - (True, src): UID processed successfully (or skipped)
+        - (False, src): Auth error occurred, caller should retry after reconnect
+    """
+    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
+
+    try:
+        resp, data = src.uid("fetch", uid, "(RFC822)")
+        if resp != "OK" or not data or data[0] is None:
+            safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
+            return (True, src)  # Don't retry, move on
+
+        raw_email = None
+        for item in data:
+            if isinstance(item, tuple):
+                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)
+            clean_subject = clean_subject[:100]
+
+        filename = f"{uid_str}_{clean_subject}.eml"
+        full_path = os.path.join(local_folder_path, filename)
+
+        if os.path.exists(full_path):
+            return (True, src)  # Already exists
+
+        if raw_email:
+            try:
+                with open(full_path, "wb") as f:
+                    f.write(raw_email)
+                safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
+            except OSError as e:
+                safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
+        else:
+            safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
+
+        return (True, src)
+
+    except Exception as e:
+        if imap_oauth2.is_auth_error(e):
+            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
+            return (False, src)  # Signal retry needed
+        else:
+            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
+            return (True, src)  # Don't retry other errors
+
+
+def process_batch(uids, folder_name, src_conf, local_folder_path):
+    src = imap_session.get_thread_connection(thread_local, "src", src_conf)
+    if not src:
+        safe_print("Error: Could not establish connection for batch.")
+        return
+
+    try:
+        src.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
+        return
+
+    for uid in uids:
+        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
+        max_retries = 2
+
+        for attempt in range(max_retries):
+            src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
+            thread_local.src = src
+            if not ok:
+                safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
+                return
+
+            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
+            thread_local.src = src
+
+            if success:
+                break
+            if attempt < max_retries - 1:
+                src = None
+                thread_local.src = None
+
+
+def get_existing_uids(local_path):
+    """
+    Scans the local directory for files matching pattern matches {UID}_*.eml
+    Returns a set of UIDs (as strings).
+    """
+    existing = set()
+    if not os.path.exists(local_path):
+        return existing
+
+    try:
+        for filename in os.listdir(local_path):
+            if filename.endswith(".eml") and "_" in filename:
+                # Expecting UID_Subject.eml
+                parts = filename.split("_", 1)
+                if parts[0].isdigit():
+                    existing.add(parts[0])
+    except Exception:
+        pass
+    return existing
+
+
+# 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 = imap_common.PRESERVABLE_FLAGS
+
+
+def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
+    """
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
+    with OAuth2 session management.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_info, imap_conn) where message_info is:
+        { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+        The returned imap_conn may be different if reconnection occurred.
+    """
+    message_info = {}
+
+    # Ensure connection is healthy (refresh OAuth2 token if needed) before initial select
+    if src_conf:
+        imap_conn = imap_session.ensure_connection(imap_conn, src_conf)
+        if not imap_conn:
+            safe_print(f"Could not establish connection for folder {folder_name}")
+            return (message_info, None)
+
+    try:
+        imap_conn.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Could not select folder {folder_name}: {e}")
+        return (message_info, imap_conn)
+
+    try:
+        resp, data = imap_conn.uid("search", None, "ALL")
+        if resp != "OK" or not data or not data[0]:
+            return (message_info, imap_conn)
+
+        uids = data[0].split()
+        if not uids:
+            return (message_info, imap_conn)
+
+        total_uids = len(uids)
+
+        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
+        batch_size = 200
+        for i in range(0, len(uids), batch_size):
+            batch = uids[i : i + batch_size]
+            uid_range = b",".join(batch)
+
+            # Report progress
+            if progress_callback:
+                progress_callback(min(i + batch_size, total_uids), total_uids)
+
+            # Proactively refresh token and ensure folder is selected
+            if src_conf:
+                imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
+                if not ok:
+                    safe_print(f"ERROR: Connection/folder lost in {folder_name}")
+                    break
+
+            try:
+                # Fetch both Message-ID header and FLAGS
+                resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
+                if resp != "OK":
+                    continue
+
+                # Parse response - items come in pairs for each message
+                for item in items:
+                    if isinstance(item, tuple) and len(item) >= 2:
+                        # First element contains UID and FLAGS info
+                        meta_str = (
+                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
+                        )
+
+                        # Extract all preservable flags from the metadata
+                        flags = []
+                        for flag in imap_common.PRESERVABLE_FLAGS:
+                            if flag in meta_str:
+                                flags.append(flag)
+
+                        # Second element contains the header
+                        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}")
+                continue
+
+    except Exception as e:
+        safe_print(f"Error searching folder {folder_name}: {e}")
+
+    return (message_info, imap_conn)
+
+
+def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
+    """
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
+
+    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
+    return message_info
+
+
+def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
+    """
+    Returns a set of Message-IDs for all emails in a given folder,
+    with OAuth2 session management.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
+        The returned imap_conn may be different if reconnection occurred.
+    """
+    info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
+    return (set(info.keys()), imap_conn)
+
+
+def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
+    """
+    Returns a set of Message-IDs for all emails in a given folder.
+    This is a convenience wrapper around get_message_info_in_folder.
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
+    return set(info.keys())
+
+
+def build_labels_manifest(imap_conn, local_path, src_conf=None):
+    """
+    Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
+    Scans all folders (labels) in the account and records which Message-IDs
+    appear in each label, plus their flags from [Gmail]/All Mail.
+
+    Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
+    Saves the manifest to labels_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
+    """
+    import time
+
+    manifest = {}
+    start_time = time.time()
+    total_emails_scanned = 0
+
+    safe_print("--- Building Gmail Labels Manifest ---")
+
+    # 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
+    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
+    all_mail_start = time.time()
+
+    def all_mail_progress_cb(current, total):
+        elapsed = time.time() - start_time
+        print(
+            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+            end="",
+            flush=True,
+        )
+
+    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
+        imap_conn, provider_gmail.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
+    )
+    print()  # New line after progress
+
+    # Initialize manifest with flags from All Mail
+    for msg_id, info in all_mail_info.items():
+        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
+
+    all_mail_elapsed = time.time() - all_mail_start
+    # Count flag statistics
+    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")
+
+    # Keep connection alive
+    try:
+        imap_conn.noop()
+    except Exception:
+        pass
+
+    # Parse folder names and filter to label folders
+    label_folders = [f for f in all_folders if provider_gmail.is_label_folder(f)]
+
+    total_folders = len(label_folders)
+    safe_print(f"Found {total_folders} label folders to scan.\n")
+
+    # Scan each label folder
+    for folder_idx, folder_name in enumerate(label_folders, 1):
+        folder_start = time.time()
+
+        # Progress callback for this folder
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
+        print()  # New line after progress
+
+        # Keep connection alive between folders
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+        # Determine the label name to store
+        # For [Gmail]/Sent Mail -> "Sent Mail"
+        # For [Gmail]/Starred -> "Starred"
+        # For INBOX -> "INBOX"
+        # For user folders -> folder name as-is
+        if folder_name.startswith("[Gmail]/"):
+            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
+        else:
+            label_name = folder_name
+
+        for msg_id in message_ids:
+            if msg_id not in manifest:
+                # Email not in All Mail (rare, but handle it)
+                manifest[msg_id] = {"labels": [], "flags": []}
+            if label_name not in manifest[msg_id]["labels"]:
+                manifest[msg_id]["labels"].append(label_name)
+
+        folder_elapsed = time.time() - folder_start
+        total_emails_scanned += len(message_ids)
+        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    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}")
+    safe_print(f"  - Unique emails: {len(manifest)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    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)
+        safe_print(f"\nLabels manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
+    """
+    Builds a manifest mapping Message-IDs to their IMAP flags.
+    For non-Gmail servers, scans specified folders (or all folders if not specified).
+
+    Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
+    Saves the manifest to flags_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
+    """
+    import time
+
+    manifest = {}
+    start_time = time.time()
+
+    safe_print("--- Building Flags Manifest ---")
+
+    # Get folders to scan
+    if folders_to_scan:
+        all_folders = folders_to_scan
+    else:
+        try:
+            typ, folders = imap_conn.list()
+            if typ != "OK":
+                safe_print("Error: Could not list folders.")
+                return manifest
+            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
+        except Exception as e:
+            safe_print(f"Error listing folders: {e}")
+            return manifest
+
+    total_folders = len(all_folders)
+    safe_print(f"Found {total_folders} folders to scan.\n")
+
+    # Scan each folder
+    for folder_idx, folder_name in enumerate(all_folders, 1):
+        folder_start = time.time()
+
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
+        print()  # New line after progress
+
+        # Merge into manifest (keep first occurrence of flags)
+        for msg_id, info in folder_info.items():
+            if msg_id not in manifest:
+                manifest[msg_id] = {"flags": info.get("flags", [])}
+
+        folder_elapsed = time.time() - folder_start
+        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
+
+        # Keep connection alive
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    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)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    manifest_path = os.path.join(local_path, "flags_manifest.json")
+    try:
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2, ensure_ascii=False)
+        safe_print(f"\nFlags manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+def delete_orphan_local_files(local_folder_path, server_uids):
+    """
+    Delete local .eml files that no longer exist on the server.
+    Args:
+        local_folder_path: Path to local folder containing .eml files
+        server_uids: Set of UID strings currently on server
+    Returns:
+        Count of deleted files
+    """
+    deleted_count = 0
+    if not os.path.exists(local_folder_path):
+        return deleted_count
+
+    try:
+        for filename in os.listdir(local_folder_path):
+            if not filename.endswith(".eml") or "_" not in filename:
+                continue
+
+            # Extract UID from filename (format: {UID}_{Subject}.eml)
+            parts = filename.split("_", 1)
+            if not parts[0].isdigit():
+                continue
+
+            local_uid = parts[0]
+            if local_uid not in server_uids:
+                file_path = os.path.join(local_folder_path, filename)
+                try:
+                    os.remove(file_path)
+                    safe_print(f"  -> Deleted orphan: {filename}")
+                    deleted_count += 1
+                except Exception as e:
+                    safe_print(f"  -> Error deleting {filename}: {e}")
+    except Exception as e:
+        safe_print(f"Error scanning for orphan files: {e}")
+
+    return deleted_count
+
+
+def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
+    safe_print(f"--- Processing Folder: {folder_name} ---")
+
+    # create local path
+    # Handle folder separators. IMAP output might be "Parent/Child"
+    # We rely on OS to handle "Parent/Child" as subdirectories using join
+    # But clean the segments
+    cleaned_name = folder_name.replace("/", os.sep)
+    local_folder_path = os.path.join(local_base_path, cleaned_name)
+
+    try:
+        os.makedirs(local_folder_path, exist_ok=True)
+    except Exception as e:
+        safe_print(f"Error creating directory {local_folder_path}: {e}")
+        return
+
+    # Select IMAP folder
+    try:
+        src_main.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Skipping {folder_name}: {e}")
+        return
+
+    # Search all
+    resp, data = src_main.uid("search", None, "ALL")
+    if resp != "OK":
+        return
+
+    uids = data[0].split()
+    total_on_server = len(uids)
+
+    # Build set of server UIDs for comparison
+    server_uid_set = set()
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        server_uid_set.add(u_str)
+
+    if total_on_server == 0:
+        safe_print(f"Folder {folder_name} is empty.")
+        # If dest_delete enabled, delete all local files
+        if dest_delete:
+            deleted = delete_orphan_local_files(local_folder_path, set())
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+        return
+
+    # Incremental Optimization
+    # Read local directory to find UIDs we already have
+    existing_uids = get_existing_uids(local_folder_path)
+
+    # Delete orphan local files if dest_delete is enabled
+    if dest_delete:
+        orphan_uids = existing_uids - server_uid_set
+        if orphan_uids:
+            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
+            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+
+    # Filter UIDs
+    # decode uid first if bytes
+    uids_to_download = []
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        if u_str not in existing_uids:
+            uids_to_download.append(u)
+
+    skipped = total_on_server - len(uids_to_download)
+    if skipped > 0:
+        safe_print(f"Skipping {skipped} emails (already exist locally).")
+
+    if not uids_to_download:
+        safe_print(f"Folder {folder_name} is up to date.")
+        return
+
+    safe_print(f"Downloading {len(uids_to_download)} new emails...")
+
+    # Create batches
+    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
+
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
+    try:
+        futures = []
+        for batch in uid_batches:
+            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
+
+        for future in concurrent.futures.as_completed(futures):
+            future.result()
+
+    except KeyboardInterrupt:
+        executor.shutdown(wait=False, cancel_futures=True)
+        raise
+    finally:
+        executor.shutdown(wait=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    # Source
+    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,
+        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")
+    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
+
+    # Gmail Labels
+    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-labels",
+        action="store_true",
+        default=env_preserve_labels,
+        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
+    )
+    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-flags",
+        action="store_true",
+        default=env_preserve_flags,
+        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
+    )
+    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
+    parser.add_argument(
+        "--manifest-only",
+        action="store_true",
+        default=env_manifest_only,
+        help="Gmail only: Build the labels manifest and exit without downloading emails",
+    )
+    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
+    parser.add_argument(
+        "--gmail-mode",
+        action="store_true",
+        default=env_gmail_mode,
+        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
+    )
+
+    # Sync mode: delete local files not on server
+    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
+    parser.add_argument(
+        "--dest-delete",
+        action="store_true",
+        default=env_dest_delete,
+        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
+    )
+
+    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
+
+    args = parser.parse_args()
+
+    global MAX_WORKERS, BATCH_SIZE
+    MAX_WORKERS = args.workers
+    BATCH_SIZE = args.batch
+
+    # Build connection config (acquires OAuth2 token if configured)
+    src_conf = imap_session.build_imap_conf(
+        args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
+    )
+
+    # Expand path (~/...)
+    local_path = os.path.expanduser(args.dest_path)
+    if not os.path.exists(local_path):
+        try:
+            os.makedirs(local_path)
+            print(f"Created backup directory: {local_path}")
+        except Exception as e:
+            print(f"Error creating backup directory: {e}")
+            sys.exit(1)
+
+    print("\n--- Configuration Summary ---")
+    print(f"Source Host     : {args.src_host}")
+    print(f"Source User     : {args.src_user}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
+    print(f"Destination Path: {local_path}")
+    if args.gmail_mode:
+        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
+    elif args.manifest_only:
+        print("Mode            : Manifest Only (no email download)")
+    elif args.folder:
+        print(f"Target Folder   : {args.folder}")
+    if args.preserve_labels or args.manifest_only or args.gmail_mode:
+        print("Preserve Labels : Yes (Gmail)")
+    if args.preserve_flags or args.gmail_mode:
+        print("Preserve Flags  : Yes (read/starred/answered/draft)")
+    if args.dest_delete:
+        print("Dest Delete     : Yes (remove local orphans)")
+    print("-----------------------------\n")
+
+    try:
+        src = imap_common.get_imap_connection_from_conf(src_conf)
+        if not src:
+            sys.exit(1)
+
+        # Build labels manifest BEFORE backing up emails (Gmail mode)
+        # This way we capture the label state at backup time
+        if args.preserve_labels or args.manifest_only or args.gmail_mode:
+            print("Building Gmail labels manifest...")
+            print("This scans all folders to map Message-IDs to labels and flags.\n")
+            build_labels_manifest(src, local_path, src_conf)
+            print("")  # Blank line after manifest building
+        # Build flags-only manifest for non-Gmail servers
+        elif args.preserve_flags:
+            print("Building flags manifest...")
+            print("This scans folders to capture read/starred/etc status.\n")
+            # Get folders to scan
+            folders_to_scan = [args.folder] if args.folder else None
+            build_flags_manifest(src, local_path, folders_to_scan, src_conf)
+            print("")  # Blank line after manifest building
+
+        # If manifest-only mode, we're done
+        if args.manifest_only:
+            try:
+                src.logout()
+            except Exception:
+                pass  # Connection may already be closed
+            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:")
+            print(f'  python3 imap_backup.py --dest-path "{local_path}" "[Gmail]/All Mail"')
+            sys.exit(0)
+
+        # Gmail mode: backup only [Gmail]/All Mail
+        if args.gmail_mode:
+            backup_folder(src, provider_gmail.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:
+            # Reconnect after potentially long manifest building
+            src = imap_session.ensure_connection(src, src_conf)
+            if not src:
+                print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
+                sys.exit(0)
+            folders = imap_common.list_selectable_folders(src)
+            for name in folders:
+                if provider_exchange.is_special_folder(name):
+                    print(f"Skipping Exchange system folder: {name}")
+                    continue
+                src = imap_session.ensure_connection(src, src_conf)
+                if not src:
+                    print("Fatal: Could not reconnect to IMAP server. Aborting.")
+                    sys.exit(1)
+                backup_folder(src, name, local_path, src_conf, args.dest_delete)
+
+        try:
+            src.logout()
+        except Exception:
+            pass  # Connection may already be closed
+        print("\nBackup completed successfully.")
+
+        if args.preserve_labels or args.gmail_mode:
+            manifest_path = os.path.join(local_path, "labels_manifest.json")
+            print(f"\nGmail labels manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply labels and flags to emails.")
+        elif args.preserve_flags:
+            manifest_path = os.path.join(local_path, "flags_manifest.json")
+            print(f"\nFlags manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply read/starred status to emails.")
+
+    except KeyboardInterrupt:
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_compare.py

@@ -0,0 +1,337 @@
+"""
+IMAP Folder Comparison Script
+
+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.
+
+Configuration (Environment Variables):
+  Source Account:
+    SRC_IMAP_HOST       : Source IMAP Host
+    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 imap_compare.py
+
+Examples:
+        # IMAP -> IMAP
+        python3 imap_compare.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 imap_compare.py \
+            --src-path "./my_backup" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # IMAP -> Local
+        python3 imap_compare.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 auth import imap_oauth2
+from utils import imap_common
+
+
+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
+        typ, data = conn.select(f'"{folder_name}"', readonly=True)
+        if typ != "OK":
+            return None
+
+        # SELECT command returns the number of messages in data[0]
+        # data[0] is bytes, e.g. b'123'
+        if data and data[0]:
+            return int(data[0])
+        return 0
+
+    except Exception:
+        # print(f"Error checking {folder_name}: {e}")
+        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("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+    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
+    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
+    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()
+
+    src_is_local = bool(args.src_path)
+    dest_is_local = bool(args.dest_path)
+
+    SRC_HOST = args.src_host
+    SRC_USER = args.src_user
+    DEST_HOST = args.dest_host
+    DEST_USER = args.dest_user
+
+    # Acquire OAuth2 tokens if configured
+    src_oauth2_token = None
+    src_oauth2_provider = None
+    if not src_is_local and args.src_client_id:
+        src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
+            SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
+        )
+
+    dest_oauth2_token = None
+    dest_oauth2_provider = None
+    if not dest_is_local and args.dest_client_id:
+        dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
+            DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
+        )
+
+    print("\n--- Configuration Summary ---")
+    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     : {imap_oauth2.auth_description(src_oauth2_provider)}")
+
+    if dest_is_local:
+        print(f"Destination (Local): {args.dest_path}")
+    else:
+        print(f"Destination Host: {args.dest_host}")
+        print(f"Destination User: {args.dest_user}")
+        print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
+    print("-----------------------------\n")
+
+    src = None
+    dest = None
+
+    try:
+        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
+
+        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...")
+        if src_is_local:
+            folders = imap_common.list_local_folders(args.src_path)
+        else:
+            folders = imap_common.list_selectable_folders(src)
+
+        if not folders:
+            print("Failed to list source folders.")
+            return
+
+        # Prepare Table Header
+        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
+        print("-" * len(header))
+        print(header)
+        print("-" * len(header))
+
+        total_src = 0
+        total_dest = 0
+
+        # Iterate through Source folders
+        for folder_name in folders:
+            # Get Counts
+            if src_is_local:
+                src_count = imap_common.get_local_email_count(args.src_path, folder_name)
+            else:
+                src_count = get_email_count(src, folder_name)
+
+            if dest_is_local:
+                dest_count = imap_common.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"
+            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
+
+            diff_str = ""
+            if src_count is not None and dest_count is not None:
+                diff = src_count - dest_count
+                diff_str = str(diff)
+                total_src += src_count
+                total_dest += dest_count
+            elif src_count is not None:
+                total_src += src_count
+
+            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
+
+        print("-" * len(header))
+        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
+
+    except KeyboardInterrupt:
+        # Re-raise to be handled by the outer block, but ensure finally runs
+        raise
+
+    finally:
+        # Check source connection state and logout if possible
+        if src:
+            try:
+                src.logout()
+            except BaseException:
+                pass
+
+        # Check dest connection state and logout if possible
+        if dest:
+            try:
+                dest.logout()
+            except BaseException:
+                pass
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_count.py

@@ -0,0 +1,247 @@
+"""IMAP Email Counting Script.
+
+Counts emails per folder from either:
+- An IMAP account, or
+- A local backup folder created by ``imap_backup.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)
+
+    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
+
+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 imap_count.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 imap_count.py
+
+    # Count a local backup
+    python3 imap_count.py --path "./my_backup"
+
+    # Or set a default local backup path via env var
+    export BACKUP_LOCAL_PATH="./my_backup"
+    python3 imap_count.py
+"""
+
+import argparse
+import imaplib
+import os
+import sys
+from typing import Optional
+
+from auth import imap_oauth2
+from utils import imap_common
+
+
+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, oauth2_token)
+        if not mail:
+            return
+
+        # List all mailboxes
+        print("Listing mailboxes...")
+        folders = imap_common.list_selectable_folders(mail)
+
+        if not folders:
+            print("Failed to list mailboxes.")
+            return
+
+        total_all_folders = 0
+        print(f"{'Folder Name':<40} {'Count':>10}")
+        print("-" * 52)
+
+        for folder_name in folders:
+            display_name = folder_name
+
+            try:
+                # Select the mailbox (read-only is sufficient for counting)
+                # folder_name extracted from list usually handles quotes correctly for select
+                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
+                if rv != "OK":
+                    print(f"{display_name:<40} {'Skipped':>10}")
+                    continue
+
+                # Search for all emails
+                status, data = mail.search(None, "ALL")
+
+                if status == "OK":
+                    # data[0] is space separated IDs
+                    email_ids = data[0].split()
+                    count = len(email_ids)
+                    print(f"{display_name:<40} {count:>10}")
+                    total_all_folders += count
+                else:
+                    print(f"{display_name:<40} {'Error':>10}")
+
+            except imaplib.IMAP4.error:
+                print(f"{display_name:<40} {'Error':>10}")
+
+        print("-" * 52)
+        print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+        # Logout
+        mail.logout()
+
+    except imaplib.IMAP4.error as e:
+        print(f"IMAP Error: {e}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def count_local_emails(local_path: str) -> None:
+    print(f"Scanning local backup: {local_path}")
+
+    folders = imap_common.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 = imap_common.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("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    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,
+        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)",
+    )
+
+    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
+
+    # Acquire OAuth2 token if configured
+    oauth2_token = None
+    oauth2_provider = None
+    if args.client_id:
+        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
+            IMAP_SERVER, args.client_id, USERNAME, args.client_secret
+        )
+
+    print("\n--- Configuration Summary ---")
+    print(f"Host            : {IMAP_SERVER}")
+    print(f"User            : {USERNAME}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(oauth2_provider)}")
+    print("-----------------------------\n")
+
+    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_restore.py

@@ -0,0 +1,976 @@
+"""
+IMAP Email Restore Script
+
+Restores emails from a local backup to an IMAP account.
+Reads .eml files from a local directory and uploads them to the destination server.
+
+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).
+- Parallel Processing: Uses multithreading for fast uploads.
+- Date Preservation: Restores emails with their original dates.
+
+Configuration (Environment Variables):
+    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).
+  APPLY_LABELS: Set to "true" to apply Gmail labels from manifest. Default is "false".
+  APPLY_FLAGS: Set to "true" to apply IMAP flags from manifest. Default is "false".
+  GMAIL_MODE: Set to "true" for Gmail restore mode. Default is "false".
+  DEST_DELETE: Set to "true" to delete emails from destination not found in local backup.
+              Default is "false".
+
+Usage:
+    python3 imap_restore.py \
+        --src-path "./my_backup" \
+        --dest-host "imap.gmail.com" \
+        --dest-user "you@gmail.com" \
+        --dest-pass "your-app-password"
+
+Gmail Labels Restoration:
+    python3 imap_restore.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.
+"""
+
+import argparse
+import concurrent.futures
+import os
+import sys
+import threading
+import time
+from email import policy
+from email.parser import BytesParser
+from email.utils import parsedate_to_datetime
+from enum import Enum
+from typing import Optional
+
+from auth import imap_oauth2
+from core import imap_session
+from providers import provider_gmail
+from utils import imap_common, restore_cache
+
+
+class UploadResult(Enum):
+    """Result of an email upload operation."""
+
+    SUCCESS = "success"
+    ALREADY_EXISTS = "already_exists"
+    FAILURE = "failure"
+
+
+# Defaults
+MAX_WORKERS = 4  # Lower default for restore to avoid rate limits
+BATCH_SIZE = 10
+
+# Thread-local storage
+thread_local = threading.local()
+safe_print = imap_common.safe_print
+
+
+def get_flags_from_manifest(manifest, message_id):
+    """
+    Extract IMAP flags string from manifest entry.
+    Returns flags string like "\\Seen \\Flagged" or None.
+    """
+    if not message_id or message_id not in manifest:
+        return None
+
+    entry = manifest[message_id]
+
+    if isinstance(entry, dict) and "flags" in entry:
+        flags = entry.get("flags", [])
+        if flags:
+            return " ".join(flags)
+
+    return None
+
+
+def parse_eml_file(file_path):
+    """
+    Parse an .eml file and extract metadata.
+    Returns (message_id, date_str, raw_content, subject) or (None, None, None, None) on error.
+    """
+    try:
+        with open(file_path, "rb") as f:
+            raw_content = f.read()
+
+        # Use compat32 to preserve raw headers with continuation lines
+        parser = BytesParser(policy=policy.compat32)
+        msg = parser.parsebytes(raw_content, headersonly=True)
+
+        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
+        date_str = None
+        if date_header:
+            try:
+                dt = parsedate_to_datetime(date_header)
+                # Format: "DD-Mon-YYYY HH:MM:SS +ZZZZ"
+                date_str = dt.strftime('"%d-%b-%Y %H:%M:%S %z"')
+            except Exception:
+                pass
+
+        return message_id, date_str, raw_content, subject
+
+    except Exception as e:
+        safe_print(f"Error parsing {file_path}: {e}")
+        return None, None, None, None
+
+
+def get_eml_files(folder_path):
+    """
+    Get all .eml files in a folder.
+    Returns list of (file_path, filename) tuples.
+    """
+    eml_files = []
+    if not os.path.exists(folder_path):
+        return eml_files
+
+    try:
+        for filename in os.listdir(folder_path):
+            if filename.endswith(".eml"):
+                file_path = os.path.join(folder_path, filename)
+                eml_files.append((file_path, filename))
+    except Exception as e:
+        safe_print(f"Error listing {folder_path}: {e}")
+
+    return eml_files
+
+
+def upload_email(dest, folder_name, raw_content, date_str, flags=None):
+    """
+    Upload a single email to the destination folder.
+    Returns UploadResult enum: SUCCESS or FAILURE.
+
+    Callers are expected to check for duplicates via load_folder_msg_ids()
+    before calling this function, and to ensure the folder exists.
+
+    Args:
+        flags: Optional string of IMAP flags like "\\Seen" for read emails.
+    """
+    try:
+        success = imap_common.append_email(
+            dest,
+            folder_name,
+            raw_content,
+            date_str,
+            flags,
+            ensure_folder=False,
+        )
+        return UploadResult.SUCCESS if success else UploadResult.FAILURE
+
+    except Exception as e:
+        safe_print(f"Error uploading to {folder_name}: {e}")
+        return UploadResult.FAILURE
+
+
+def get_labels_from_manifest(manifest, message_id):
+    """
+    Get the list of labels for a message from the manifest.
+    Returns list of label strings or empty list.
+    """
+    if not message_id or message_id not in manifest:
+        return []
+
+    entry = manifest[message_id]
+    if isinstance(entry, dict):
+        return entry.get("labels", [])
+    elif isinstance(entry, list):
+        return entry  # Old format: list of labels
+    return []
+
+
+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.
+
+    Args:
+        folder_name: Target folder, or "__GMAIL_MODE__" for per-email folder selection
+        manifest: Combined manifest with labels and/or flags
+        apply_labels: Whether to apply Gmail labels from manifest
+        apply_flags: Whether to apply IMAP flags from manifest
+    """
+    dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
+    if not dest:
+        safe_print("Error: Could not establish connection for batch.")
+        return
+
+    gmail_mode = folder_name == "__GMAIL_MODE__"
+
+    for file_path, filename in eml_files:
+        # Proactively refresh token if needed
+        dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
+        if not dest:
+            safe_print(f"ERROR: Connection lost for {filename}")
+            return
+
+        try:
+            message_id, date_str, raw_content, subject = parse_eml_file(file_path)
+            if raw_content is None:
+                continue  # No content, skip to next file
+
+            size = len(raw_content)
+            size_str = f"{size / 1024:.1f}KB"
+
+            # Truncate subject for display
+            display_subject = (subject[:40] + "...") if len(subject) > 40 else subject
+
+            # Get flags from manifest if apply_flags is enabled
+            flags = None
+            if apply_flags:
+                flags = get_flags_from_manifest(manifest, message_id)
+
+            # Get labels for this message
+            labels = get_labels_from_manifest(manifest, message_id) if apply_labels else []
+
+            # Determine target folder and remaining labels
+            if gmail_mode:
+                target_folder, remaining_labels = provider_gmail.resolve_target(labels)
+            else:
+                target_folder = folder_name
+                remaining_labels = labels
+
+            existing_dest_msg_ids = imap_common.load_folder_msg_ids(
+                dest,
+                target_folder,
+                existing_dest_msg_ids_by_folder,
+                existing_dest_msg_ids_lock,
+                progress_cache_data,
+                progress_cache_lock,
+                dest_host,
+                dest_user,
+            )
+
+            # Check if email already exists on destination using pre-fetched set
+            email_already_on_dest = (
+                message_id and existing_dest_msg_ids is not None and message_id in existing_dest_msg_ids
+            )
+
+            # Incremental default: skip entirely if already present
+            if email_already_on_dest and not full_restore:
+                safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
+                continue  # Skip to next file
+
+            if email_already_on_dest:
+                # Full restore: treat as existing for flag sync
+                upload_result = UploadResult.ALREADY_EXISTS
+            else:
+                # Upload — skip per-message SEARCH since we have a pre-fetched set
+                upload_result = upload_email(
+                    dest,
+                    target_folder,
+                    raw_content,
+                    date_str,
+                    flags,
+                )
+
+            # Only record progress when upload succeeds or email already exists.
+            # Failed uploads should not be marked as processed to allow retry on next run.
+            if upload_result in (UploadResult.SUCCESS, UploadResult.ALREADY_EXISTS):
+                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 upload_result == UploadResult.ALREADY_EXISTS:
+                safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}")
+                # Full restore preserves legacy behavior: sync flags on existing email if requested
+                if full_restore and apply_flags and flags and message_id:
+                    imap_common.sync_flags_on_existing(dest, target_folder, message_id, flags, size)
+            elif upload_result == UploadResult.SUCCESS:
+                safe_print(f"[{target_folder}] UPLOADED      | {size_str:<8} | {display_subject}")
+                # Show applied flags in same style as labels
+                if flags:
+                    for flag in flags.split():
+                        safe_print(f"  -> Applied flag: {flag}")
+            else:  # FAILURE
+                safe_print(f"[{target_folder}] FAILED        | {size_str:<8} | {display_subject}")
+
+            # 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 (upload_result == UploadResult.SUCCESS or full_restore):
+                for label in remaining_labels:
+                    label_folder = provider_gmail.label_to_folder(label)
+
+                    # Skip if this is the same as the target folder
+                    if label_folder == target_folder:
+                        continue
+
+                    # Skip system folders we can't upload to
+                    if label_folder in (
+                        provider_gmail.GMAIL_ALL_MAIL,
+                        provider_gmail.GMAIL_SPAM,
+                        provider_gmail.GMAIL_TRASH,
+                    ):
+                        continue
+
+                    try:
+                        # Get or fetch Message-IDs for label folder (one-time server fetch per folder)
+                        label_folder_msg_ids = imap_common.load_folder_msg_ids(
+                            dest,
+                            label_folder,
+                            existing_dest_msg_ids_by_folder,
+                            existing_dest_msg_ids_lock,
+                            progress_cache_data,
+                            progress_cache_lock,
+                            dest_host,
+                            dest_user,
+                        )
+
+                        # Check duplicate using pre-fetched set instead of per-message SEARCH
+                        label_already_exists = (
+                            label_folder_msg_ids is not None and message_id and message_id in label_folder_msg_ids
+                        )
+
+                        if not label_already_exists:
+                            append_success = imap_common.append_email(
+                                dest,
+                                label_folder,
+                                raw_content,
+                                date_str,
+                                flags,
+                                ensure_folder=False,
+                            )
+                            if append_success:
+                                restore_cache.record_progress(
+                                    message_id=message_id,
+                                    folder_name=label_folder,
+                                    existing_dest_msg_ids=label_folder_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,
+                                )
+                                safe_print(f"  -> Applied label: {label}")
+                            else:
+                                safe_print(f"  -> Failed to apply label {label} (will retry on next restore)")
+                        # If email exists in this label folder, sync flags (full restore only)
+                        elif full_restore and apply_flags and flags:
+                            imap_common.sync_flags_on_existing(dest, label_folder, message_id, flags, size)
+                    except Exception as e:
+                        if imap_oauth2.is_auth_error(e):
+                            raise
+                        safe_print(f"  -> Error applying label {label}: {e}")
+
+        except Exception as e:
+            safe_print(f"Error processing {filename}: {e}")
+
+
+def get_local_message_ids(local_folder_path):
+    """
+    Get a set of Message-IDs from all .eml files in a local folder.
+    Used for destination deletion sync.
+    """
+    message_ids = set()
+    eml_files = get_eml_files(local_folder_path)
+    for file_path, _filename in eml_files:
+        message_id, _, _, _ = parse_eml_file(file_path)
+        if message_id:
+            message_ids.add(message_id)
+    return message_ids
+
+
+def pre_filter_eml_files(eml_files, dest_msg_ids):
+    """Filter out .eml files whose Message-IDs already exist in the destination."""
+    safe_print("Pre-filtering duplicates...")
+    files_to_restore = []
+    skipped = 0
+    for file_path, filename in eml_files:
+        msg_id = imap_common.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.")
+    return files_to_restore
+
+
+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,
+    progress_cache_file: Optional[str] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+):
+    """
+    Restore all emails from a local folder to the destination IMAP server.
+    """
+    safe_print(f"--- Restoring Folder: {folder_name} ---")
+
+    eml_files = get_eml_files(local_folder_path)
+    if not eml_files:
+        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_from_conf(dest_conf)
+            if dest:
+                imap_common.delete_orphan_emails(dest, folder_name, set())
+                dest.logout()
+        return
+
+    safe_print(f"Found {len(eml_files)} emails to restore.")
+
+    cache_root = cache_root or local_folder_path
+    cache_path = progress_cache_file
+    cache_data = progress_cache_data
+    cache_lock = progress_cache_lock
+    if cache_path is None or cache_data is None or cache_lock is None:
+        cache_path, cache_data, cache_lock = imap_common.load_progress_cache(
+            cache_root,
+            dest_conf["host"],
+            dest_conf["user"],
+            log_fn=safe_print,
+        )
+
+    # 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:
+        safe_print("Building local Message-ID index for sync...")
+        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.")
+
+        # Update destination Message-IDs with what we processed
+        if not full_restore:
+            with existing_dest_msg_ids_lock:
+                existing_dest_msg_ids_by_folder.setdefault(folder_name, set()).update(dest_msg_ids)
+        else:
+            existing_dest_msg_ids_by_folder[folder_name] = dest_msg_ids
+        dest_msg_ids = existing_dest_msg_ids_by_folder[folder_name]
+
+        # Pre-filter files to skip duplicates (skip in full_restore: need to process all for flag/label sync)
+        if dest_msg_ids and not full_restore:
+            files_to_restore = pre_filter_eml_files(eml_files, dest_msg_ids)
+
+    if not files_to_restore:
+        safe_print("No new emails to restore.")
+        if dest_delete and local_msg_ids is not None:
+            safe_print("Syncing destination: removing emails not in local backup...")
+            dest = imap_session.ensure_connection(None, dest_conf)
+            if dest:
+                imap_common.delete_orphan_emails(dest, folder_name, local_msg_ids)
+                dest.logout()
+
+        restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
+        return
+
+    safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...")
+
+    # Create batches
+    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 = []
+        for batch in batches:
+            futures.append(
+                executor.submit(
+                    process_restore_batch,
+                    batch,
+                    folder_name,
+                    dest_conf,
+                    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"),
+                )
+            )
+
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                future.result()
+            except Exception as e:
+                safe_print(f"Batch error: {e}")
+
+    # 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_session.ensure_connection(None, dest_conf)
+        if dest:
+            imap_common.delete_orphan_emails(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,
+    full_restore: bool = False,
+    progress_cache_file: Optional[str] = None,
+    progress_cache_data: Optional[dict] = None,
+    progress_cache_lock: Optional[threading.Lock] = None,
+):
+    """
+    Special restoration mode for Gmail: Upload emails to their first label folder
+    and then apply additional labels from the manifest.
+
+    This avoids putting all emails in INBOX - emails only appear in INBOX
+    if they originally had the INBOX label.
+    """
+    # Find the All Mail folder in the backup
+    all_mail_path = os.path.join(local_path, "[Gmail]", "All Mail")
+    if not os.path.exists(all_mail_path):
+        # Try without subfolder structure
+        all_mail_path = local_path
+
+    safe_print("--- Gmail Restore with Labels ---")
+    safe_print(f"Source path: {all_mail_path}")
+    safe_print(f"Entries in manifest: {len(manifest)}")
+
+    eml_files = get_eml_files(all_mail_path)
+    if not eml_files:
+        safe_print("No .eml files found for restoration.")
+        return
+
+    safe_print(f"Found {len(eml_files)} emails to restore.\n")
+
+    # Process in batches
+    total = len(eml_files)
+    start_time = time.time()
+
+    batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)]
+
+    cache_path = progress_cache_file
+    if cache_path is None or progress_cache_data is None or progress_cache_lock is None:
+        cache_path, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
+            local_path,
+            dest_conf["host"],
+            dest_conf["user"],
+            log_fn=safe_print,
+        )
+    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:
+            # For Gmail, we use a special folder marker to indicate gmail-mode
+            # The process_restore_batch will determine the target folder per-email
+            # based on the manifest labels
+            futures.append(
+                executor.submit(
+                    process_restore_batch,
+                    batch,
+                    "__GMAIL_MODE__",  # Special marker - target determined per-email from manifest
+                    dest_conf,
+                    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"),
+                )
+            )
+
+        completed = 0
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                future.result()
+                completed += 1
+                elapsed = time.time() - start_time
+                progress = (completed / len(batches)) * 100
+                safe_print(f"Progress: {progress:.1f}% ({elapsed:.0f}s elapsed)")
+            except Exception as e:
+                safe_print(f"Batch error: {e}")
+
+    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 main():
+    parser = argparse.ArgumentParser(description="Restore IMAP emails from local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
+
+    # Source (Local Path)
+    env_path = os.getenv("BACKUP_LOCAL_PATH")
+    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=default_dest_host,
+        required=not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--dest-user",
+        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-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
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=int(os.getenv("MAX_WORKERS", 4)),
+        help="Thread count (default: 4)",
+    )
+    parser.add_argument(
+        "--batch",
+        type=int,
+        default=int(os.getenv("BATCH_SIZE", 10)),
+        help="Emails per batch",
+    )
+
+    # Gmail Labels
+    env_apply_labels = os.getenv("APPLY_LABELS", "false").lower() == "true"
+    parser.add_argument(
+        "--apply-labels",
+        action="store_true",
+        default=env_apply_labels,
+        help="Apply Gmail labels from labels_manifest.json",
+    )
+    env_apply_flags = os.getenv("APPLY_FLAGS", "false").lower() == "true"
+    parser.add_argument(
+        "--apply-flags",
+        action="store_true",
+        default=env_apply_flags,
+        help="Apply IMAP flags (read/starred/answered/draft) from manifest",
+    )
+    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
+    parser.add_argument(
+        "--gmail-mode",
+        action="store_true",
+        default=env_gmail_mode,
+        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(
+        "--dest-delete",
+        action="store_true",
+        default=env_dest_delete,
+        help="Delete emails from destination that don't exist in local backup (sync mode)",
+    )
+
+    # Optional folder filter
+    parser.add_argument("folder", nargs="?", help="Specific folder to restore")
+
+    args = parser.parse_args()
+
+    global MAX_WORKERS, BATCH_SIZE
+    MAX_WORKERS = args.workers
+    BATCH_SIZE = args.batch
+
+    # Build connection config (acquires OAuth2 token if configured)
+    dest_conf = imap_session.build_imap_conf(
+        args.dest_host, args.dest_user, args.dest_pass, args.dest_client_id, args.dest_client_secret, "destination"
+    )
+
+    # Expand path
+    local_path = os.path.expanduser(args.src_path)
+    if not os.path.exists(local_path):
+        print(f"Error: Source path does not exist: {local_path}")
+        sys.exit(1)
+
+    # Load manifest(s) if needed
+    manifest = {}
+    apply_labels = args.apply_labels or args.gmail_mode
+    apply_flags = args.apply_flags or args.gmail_mode
+
+    if apply_labels or apply_flags:
+        # Try to load labels manifest first (contains both labels and flags for Gmail backups)
+        manifest = imap_common.load_manifest(local_path, "labels_manifest.json")
+
+        # If no labels manifest, try flags-only manifest
+        if not manifest and apply_flags:
+            manifest = imap_common.load_manifest(local_path, "flags_manifest.json")
+
+        if not manifest:
+            print("Warning: No manifest found. Labels/flags will not be applied.")
+
+    progress_cache_file = None
+    progress_cache_data = None
+    progress_cache_lock = None
+    try:
+        progress_cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
+            local_path,
+            args.dest_host,
+            args.dest_user,
+            log_fn=safe_print,
+        )
+    except Exception as e:
+        safe_print(f"Warning: Failed to load progress cache: {e}")
+
+    print("\n--- Configuration Summary ---")
+    print(f"Source Path     : {local_path}")
+    print(f"Destination Host: {args.dest_host}")
+    print(f"Destination User: {args.dest_user}")
+    print(f"Destination Auth: {imap_oauth2.auth_description(dest_conf['oauth2'] and dest_conf['oauth2']['provider'])}")
+    print(f"Workers         : {args.workers}")
+    if args.gmail_mode:
+        print("Mode            : Gmail Restore with Labels + Flags")
+    elif args.folder:
+        print(f"Target Folder   : {args.folder}")
+    if apply_labels:
+        print(f"Apply Labels    : Yes ({len(manifest)} mappings)")
+    if apply_flags:
+        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_from_conf(dest_conf)
+        if not dest:
+            print("Error: Could not connect to destination server.")
+            sys.exit(1)
+
+        if args.gmail_mode:
+            dest.logout()
+            # Special Gmail mode
+            restore_gmail_with_labels(
+                local_path,
+                dest_conf,
+                manifest,
+                apply_flags,
+                full_restore=args.full_restore,
+                progress_cache_file=progress_cache_file,
+                progress_cache_data=progress_cache_data,
+                progress_cache_lock=progress_cache_lock,
+            )
+            dest = None  # Connection handled by restore_gmail_with_labels
+        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,
+                full_restore=args.full_restore,
+                cache_root=local_path,
+                progress_cache_file=progress_cache_file,
+                progress_cache_data=progress_cache_data,
+                progress_cache_lock=progress_cache_lock,
+            )
+            dest.logout()
+        else:
+            # Restore all folders
+            folders = imap_common.get_backup_folders(local_path)
+            if not folders:
+                print("No backup folders found.")
+                sys.exit(1)
+
+            print(f"Found {len(folders)} folders to restore.\n")
+            for folder_name, folder_path in folders:
+                # Skip manifest files
+                if folder_name in ("labels_manifest.json", "flags_manifest.json"):
+                    continue
+
+                # Proactively refresh OAuth2 token and ensure connection is healthy between folders
+                dest = imap_session.ensure_connection(dest, dest_conf)
+                if not dest:
+                    print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
+                    sys.exit(1)
+
+                restore_folder(
+                    folder_name,
+                    folder_path,
+                    dest_conf,
+                    manifest,
+                    apply_labels,
+                    apply_flags,
+                    args.dest_delete,
+                    full_restore=args.full_restore,
+                    cache_root=local_path,
+                    progress_cache_file=progress_cache_file,
+                    progress_cache_data=progress_cache_data,
+                    progress_cache_lock=progress_cache_lock,
+                )
+
+            dest.logout()
+
+        print("\nRestore completed successfully.")
+
+    except KeyboardInterrupt:
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)

src/providers/provider_exchange.py

@@ -7,9 +7,10 @@ Constants and functions specific to Microsoft Exchange/Outlook IMAP implementati
 # Exchange/Outlook folders that typically can't be backed up via IMAP
 # These often contain proprietary data that Exchange returns as error messages
 EXCHANGE_SKIP_FOLDERS = {
-    "Suggested Contacts",
-    "Conversation History",
     "Calendar",
+    "Contacts",
+    "Conversation History",
+    "Suggested Contacts",
 }
 
 

src/providers/provider_gmail.py

@@ -4,8 +4,8 @@ Gmail-Specific IMAP Utilities
 Constants and functions specific to Gmail/Google Workspace IMAP implementation.
 """
 
-import imap_common
-import imap_oauth2
+from auth import imap_oauth2
+from utils import imap_common
 
 # Gmail system folders
 GMAIL_ALL_MAIL = "[Gmail]/All Mail"

src/utils/imap_common.py

@@ -16,9 +16,39 @@ import urllib.parse
 from email import policy
 from email.header import decode_header
 from email.parser import BytesParser
+from importlib.metadata import PackageNotFoundError, version
+
+from auth import imap_oauth2
+from utils import imap_compress, imap_retry, restore_cache
+
+
+def get_version() -> str:
+    """
+    Return the package version string.
+
+    Tries to get the version from the installed package metadata.
+    If not installed, falls back to reading pyproject.toml from the project root.
+    """
+    try:
+        return version("imap-migration-tools")
+    except PackageNotFoundError:
+        # Fallback: Try to read from pyproject.toml for local development
+        try:
+            # Assuming src/imap_common.py structure
+            root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            pyproject_path = os.path.join(root_dir, "pyproject.toml")
+
+            if os.path.exists(pyproject_path):
+                with open(pyproject_path, encoding="utf-8") as f:
+                    for line in f:
+                        # Looking for: version = "1.2.3"
+                        match = re.search(r'^version\s*=\s*"([^"]+)"', line.strip())
+                        if match:
+                            return match.group(1)
+        except Exception:
+            pass
+        return "0.0.0-dev"
 
-import imap_oauth2
-import restore_cache
 
 # Standard IMAP flags
 FLAG_SEEN = "\\Seen"
@@ -213,7 +243,9 @@ def append_email(
             else:
                 normalized_flags = f"({stripped})"
 
-    resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+    resp, data = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+    if resp != "OK":
+        safe_print(f"APPEND failed for {folder_name}: {resp} {data}")
     return resp == "OK"
 
 
@@ -287,7 +319,8 @@ def get_imap_connection(host, user, password=None, oauth2_token=None):
             conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
         else:
             conn.login(user, password)
-        return conn
+        imap_compress.enable_compression(conn, log_fn=safe_print)
+        return imap_retry.ConnectionProxy(conn, log_fn=safe_print)
     except Exception as e:
         print(f"Connection error to {host}: {e}")
         return None
@@ -460,28 +493,6 @@ def parse_message_id_and_subject_from_bytes(raw_message):
         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.
-    Only considers non-deleted messages (UNDELETED) so that messages pending expunge
-    don't block uploads of fresh copies.
-    Returns True if found, False otherwise.
-    """
-    if not msg_id:
-        return False
-
-    clean_id = msg_id.replace('"', '\\"')
-    try:
-        typ, data = dest_conn.search(None, f'UNDELETED (HEADER Message-ID "{clean_id}")')
-        if typ != "OK":
-            return False
-
-        dest_ids = data[0].split()
-        return len(dest_ids) > 0
-    except Exception:
-        return False
-
-
 def get_message_ids_in_folder(imap_conn):
     """
     Fetches all Message-IDs from the currently selected folder.
@@ -650,10 +661,12 @@ def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size):
 
         if flags_to_add:
             flags_str = " ".join(flags_to_add)
-            typ, _ = imap_conn.store(msg_num, "+FLAGS", f"({flags_str})")
+            typ, data = imap_conn.store(msg_num, "+FLAGS", f"({flags_str})")
             if typ == "OK":
                 for flag in flags_to_add:
                     safe_print(f"  -> Synced flag: {flag}")
+            else:
+                safe_print(f"STORE +FLAGS failed for {message_id} in {folder_name}: {typ} {data}")
 
     except Exception as e:
         safe_print(f"Error syncing flags for {message_id} in {folder_name}: {e}")
@@ -701,6 +714,55 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg
     return deleted_count
 
 
+def load_folder_msg_ids(
+    imap_conn,
+    folder_name,
+    msg_ids_by_folder,
+    msg_ids_lock,
+    progress_cache_data=None,
+    progress_cache_lock=None,
+    dest_host=None,
+    dest_user=None,
+):
+    """Load Message-IDs for a folder, fetching from server if not yet cached.
+
+    On first call for a folder, SELECTs the folder and fetches all Message-IDs
+    from the server (merged with progress cache). Subsequent calls return the
+    cached set without any IMAP operations.
+
+    Returns the set of Message-IDs, or None if tracking is disabled.
+    """
+    if msg_ids_by_folder is None or msg_ids_lock is None:
+        return None
+
+    with msg_ids_lock:
+        existing = msg_ids_by_folder.get(folder_name)
+
+    if existing is not None:
+        return existing
+
+    # Build from progress cache
+    built: set[str] = set()
+    if is_progress_cache_ready(progress_cache_data, progress_cache_lock) and dest_host and dest_user:
+        built = restore_cache.get_cached_message_ids(
+            progress_cache_data,
+            progress_cache_lock,
+            dest_host,
+            dest_user,
+            folder_name,
+        )
+
+    # Fetch from server for a comprehensive set (one-time cost per folder)
+    ensure_folder_exists(imap_conn, folder_name)
+    imap_conn.select(f'"{folder_name}"')
+    server_ids = set(get_message_ids_in_folder(imap_conn).values())
+    built.update(server_ids)
+
+    with msg_ids_lock:
+        msg_ids_by_folder.setdefault(folder_name, built)
+        return msg_ids_by_folder[folder_name]
+
+
 def load_manifest(local_path, filename):
     """Load a manifest JSON file from a backup directory.
 

src/utils/imap_compress.py

@@ -0,0 +1,180 @@
+"""IMAP COMPRESS=DEFLATE support (RFC 4978).
+
+Provides a transparent compression wrapper for stdlib imaplib connections.
+After authentication, call ``enable_compression(conn)`` to negotiate
+COMPRESS=DEFLATE with the server.  All subsequent IMAP traffic is then
+compressed/decompressed automatically via zlib.
+
+The implementation follows imaplib's own STARTTLS pattern: send the
+command, and on OK replace ``self.sock`` with a wrapper.
+"""
+
+from __future__ import annotations
+
+import io
+import zlib
+from collections.abc import Callable
+
+
+class _CompressedRawIO(io.RawIOBase):
+    """Minimal RawIOBase adapter so ``makefile()`` can return a BufferedReader
+    that reads through the decompression layer."""
+
+    def __init__(self, compressed_sock: _CompressedSocket):
+        self._sock = compressed_sock
+
+    def readable(self) -> bool:
+        return True
+
+    def readinto(self, b: bytearray | memoryview) -> int:
+        data = self._sock.recv(len(b))
+        if not data:
+            return 0
+        n = len(data)
+        b[:n] = data
+        return n
+
+
+class _CompressedSocket:
+    """Socket wrapper that adds zlib DEFLATE compression/decompression.
+
+    Wraps a real socket (typically SSL) so that:
+    - ``sendall()`` compresses data before sending
+    - ``recv()`` decompresses data after receiving
+
+    Python 3.13's imaplib reads via ``self.sock.recv()`` directly,
+    so intercepting at the socket level is sufficient.
+    """
+
+    def __init__(self, sock, initial_compressed: bytes = b""):
+        self._sock = sock
+        self._compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        self._decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
+        # Buffer for decompressed data not yet consumed by recv()
+        self._recv_buf = b""
+        # If there was read-ahead data in imaplib's buffer after the
+        # COMPRESS OK response, it is already compressed and must be
+        # fed to the decompressor.
+        if initial_compressed:
+            self._recv_buf = self._decompressor.decompress(initial_compressed)
+
+    # -- outgoing (compress) ---------------------------------------------------
+
+    def sendall(self, data: bytes) -> None:
+        compressed = self._compressor.compress(data)
+        compressed += self._compressor.flush(zlib.Z_SYNC_FLUSH)
+        self._sock.sendall(compressed)
+
+    def send(self, data: bytes) -> int:
+        self.sendall(data)
+        return len(data)
+
+    # -- incoming (decompress) -------------------------------------------------
+
+    def recv(self, bufsize: int) -> bytes:
+        # Return buffered decompressed data first
+        if self._recv_buf:
+            chunk = self._recv_buf[:bufsize]
+            self._recv_buf = self._recv_buf[bufsize:]
+            return chunk
+
+        # Read and decompress until we have at least 1 byte or true EOF.
+        # zlib can produce b"" for partial DEFLATE blocks; returning that
+        # would signal EOF to imaplib and break the connection.
+        decompressed = b""
+        while not decompressed:
+            raw = self._sock.recv(max(bufsize, 16384))
+            if not raw:
+                return b""  # real EOF
+            decompressed = self._decompressor.decompress(raw)
+
+        if len(decompressed) > bufsize:
+            self._recv_buf = decompressed[bufsize:]
+            return decompressed[:bufsize]
+        return decompressed
+
+    # -- lifecycle / proxy -----------------------------------------------------
+
+    def shutdown(self, how: int) -> None:
+        self._sock.shutdown(how)
+
+    def close(self) -> None:
+        self._sock.close()
+
+    def makefile(self, *args, **kwargs):
+        # Older Python imaplib reads from conn._file (a file object)
+        # rather than conn.sock.recv() directly.  The file must read
+        # through the decompression layer, not from the raw socket.
+        return io.BufferedReader(_CompressedRawIO(self))
+
+    def __getattr__(self, name: str):
+        return getattr(self._sock, name)
+
+
+def enable_compression(
+    conn,
+    *,
+    log_fn: Callable[[str], None] | None = None,
+) -> bool:
+    """Negotiate COMPRESS=DEFLATE on a raw imaplib IMAP4/IMAP4_SSL connection.
+
+    Must be called **after** authentication (login/XOAUTH2).
+
+    Returns True if compression was successfully enabled, False otherwise
+    (e.g. server does not advertise the capability, or the command failed).
+    """
+    raw_caps = getattr(conn, "capabilities", ())
+    caps = {c.decode() if isinstance(c, bytes) else c for c in raw_caps}
+    if "COMPRESS=DEFLATE" not in caps:
+        if log_fn is not None:
+            log_fn("Deflate compression not supported by server")
+        return False
+
+    try:
+        typ, _data = conn._simple_command("COMPRESS", "DEFLATE")
+    except Exception:
+        if log_fn is not None:
+            log_fn("Deflate compression negotiation failed")
+        return False
+
+    if typ != "OK":
+        if log_fn is not None:
+            log_fn("Deflate compression rejected by server")
+        return False
+
+    # Drain any read-ahead data from imaplib's internal buffers.
+    # After the server sends OK, all subsequent data is compressed.
+    # If recv() read past the OK response, the leftover bytes are
+    # already compressed and must be fed to the decompressor.
+    leftover = b""
+
+    # conn._file (BufferedReader) may have read-ahead bytes beyond
+    # the OK response that are already compressed.  Drain them.
+    old_file = getattr(conn, "_file", None)
+    if old_file is not None:
+        try:
+            buffered = old_file.peek(65536)
+            if isinstance(buffered, (bytes, bytearray)) and buffered:
+                leftover += bytes(buffered)
+        except Exception:
+            pass
+
+    readbuf = getattr(conn, "_readbuf", None)
+    if readbuf:
+        if isinstance(readbuf, (bytes, bytearray)):
+            leftover += bytes(readbuf)
+        else:
+            leftover += b"".join(readbuf)
+        if isinstance(readbuf, bytes):
+            conn._readbuf = b""
+        else:
+            readbuf.clear()
+
+    # Replace the socket (follows the STARTTLS pattern in imaplib).
+    conn.sock = _CompressedSocket(conn.sock, initial_compressed=leftover)
+    conn._file = conn.sock.makefile("rb")
+
+    if log_fn is not None:
+        log_fn("Deflate compression enabled")
+
+    return True

src/utils/imap_retry.py

@@ -0,0 +1,80 @@
+"""
+IMAP Retry Logic
+
+Transparent retry wrapper for IMAP connections that handles transient
+server errors (e.g. Microsoft 365 "Server Busy") with exponential backoff.
+"""
+
+from __future__ import annotations
+
+import time
+
+
+class ConnectionProxy:
+    """Transparent proxy that retries IMAP commands on transient server errors.
+
+    Wraps an imaplib.IMAP4 or IMAP4_SSL connection. For methods in
+    RETRYABLE_METHODS that return (typ, data) tuples, retries on transient
+    errors with exponential backoff.
+    """
+
+    TRANSIENT_PATTERNS = [b"UNAVAILABLE", b"Server Busy", b"try again", b"THROTTLED"]
+
+    # Methods that are safe to retry and return (typ, data)
+    RETRYABLE_METHODS = frozenset(
+        {
+            "uid",
+            "select",
+            "search",
+            "fetch",
+            "append",
+            "store",
+            "list",
+            "create",
+            "expunge",
+            "noop",
+        }
+    )
+
+    def __init__(self, conn, max_retries=3, initial_wait=5, log_fn=print):
+        if max_retries < 1:
+            raise ValueError(f"max_retries must be >= 1, got {max_retries}")
+        if initial_wait < 0:
+            raise ValueError(f"initial_wait must be >= 0, got {initial_wait}")
+        self._conn = conn
+        self._max_retries = max_retries
+        self._initial_wait = initial_wait
+        self._log_fn = log_fn
+
+    @classmethod
+    def _is_transient_error(cls, data):
+        """Check if IMAP response data contains transient error patterns."""
+        for item in data:
+            if isinstance(item, bytes):
+                for pattern in cls.TRANSIENT_PATTERNS:
+                    if pattern in item:
+                        return True
+        return False
+
+    def __getattr__(self, name):
+        attr = getattr(self._conn, name)
+        if name not in self.RETRYABLE_METHODS or not callable(attr):
+            return attr
+
+        def wrapper(*args, **kwargs):
+            last_result = None
+            for attempt in range(self._max_retries):
+                result = attr(*args, **kwargs)
+                if not isinstance(result, tuple) or len(result) < 2:
+                    return result
+                typ, data = result[0], result[1]
+                if typ == "OK" or not self._is_transient_error(data):
+                    return result
+                last_result = result
+                if attempt + 1 < self._max_retries:
+                    wait = self._initial_wait * (2**attempt)  # 5s, 10s, 20s
+                    self._log_fn(f"Server busy, retrying in {wait}s... (attempt {attempt + 1}/{self._max_retries})")
+                    time.sleep(wait)
+            return last_result
+
+        return wrapper

test/auth/test_imap_oauth2.py

@@ -14,11 +14,9 @@ from unittest.mock import patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import imap_oauth2
-import oauth2_google
-import oauth2_microsoft
+from auth import imap_oauth2, oauth2_google, oauth2_microsoft
 
 
 @pytest.fixture(autouse=True)

test/auth/test_oauth2_google.py

@@ -12,9 +12,9 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import oauth2_google
+from auth import oauth2_google
 
 
 @pytest.fixture(autouse=True)

test/auth/test_oauth2_microsoft.py

@@ -14,9 +14,9 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import oauth2_microsoft
+from auth import oauth2_microsoft
 from conftest import temp_env
 from mock_oauth_server import MOCK_TENANT_ID
 

test/conftest.py

@@ -33,11 +33,19 @@ def create_server_pair(src_data=None, dest_data=None):
     while p2 == p1:
         p2 = get_free_port()
 
-    src_t, src_s = start_server_thread(p1, src_data)
-    dest_t, dest_s = start_server_thread(p2, dest_data)
+    src_s, p1_actual = start_server_thread(p1, src_data)
+    dest_s, p2_actual = start_server_thread(p2, dest_data)
     time.sleep(0.3)
 
-    return (src_t, src_s, p1), (dest_t, dest_s, p2)
+    # We don't have direct access to thread object anymore as it is managed internally by TCPServer/ThreadingMixin or start_server_thread wrapper?
+    # Actually start_server_thread in tools/mock_imap_server.py now returns (server, port).
+    # The server object is a ThreadingMixIn so it handles per-request threads, but the main serve_forever loop is in a thread we created.
+    # But wait, my modified start_server_thread does: t.start(), return server, actual_port.
+    # It dropped returning 't'. I need to fix that or update consumers to not need 't'.
+    # Shutdown() stops serve_forever loop. join() is nice but maybe not strictly required if we trust shutdown.
+    # However, to avoid ResourceWarnings, we might want 't'.
+
+    return (None, src_s, p1_actual), (None, dest_s, p2_actual)
 
 
 def shutdown_server_pair(src_tuple, dest_tuple):
@@ -45,9 +53,13 @@ def shutdown_server_pair(src_tuple, dest_tuple):
     src_t, src_s, _ = src_tuple
     dest_t, dest_s, _ = dest_tuple
     src_s.shutdown()
+    src_s.server_close()  # Explicitly close socket
     dest_s.shutdown()
-    src_t.join(timeout=2)
-    dest_t.join(timeout=2)
+    dest_s.server_close()
+    if src_t:
+        src_t.join(timeout=2)
+    if dest_t:
+        dest_t.join(timeout=2)
 
 
 @pytest.fixture
@@ -82,16 +94,16 @@ def single_mock_server():
 
     def _create(initial_data=None):
         port = get_free_port()
-        thread, server = start_server_thread(port, initial_data)
+        server, actual_port = start_server_thread(port, initial_data)
         time.sleep(0.3)
-        servers.append((thread, server))
-        return server, port
+        servers.append(server)
+        return server, actual_port
 
     yield _create
 
-    for thread, server in servers:
+    for server in servers:
         server.shutdown()
-        thread.join(timeout=2)
+        server.server_close()
 
 
 @contextmanager

test/core/test_imap_session.py

@@ -12,9 +12,9 @@ import os
 import sys
 from unittest.mock import MagicMock, patch
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import imap_session
+from core import imap_session
 
 
 class TestBuildImapConf:

test/providers/test_provider_exchange.py

@@ -9,9 +9,9 @@ Tests cover:
 import os
 import sys
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import provider_exchange
+from providers import provider_exchange
 
 
 class TestIsSpecialFolder:
@@ -29,6 +29,10 @@ class TestIsSpecialFolder:
         """Test that Calendar is detected as special folder."""
         assert provider_exchange.is_special_folder("Calendar") is True
 
+    def test_contacts_is_special(self):
+        """Test that Contacts is detected as special folder."""
+        assert provider_exchange.is_special_folder("Contacts") is True
+
     def test_inbox_is_not_special(self):
         """Test that INBOX is not a special folder."""
         assert provider_exchange.is_special_folder("INBOX") is False
@@ -76,11 +80,12 @@ class TestExchangeConstants:
         assert "Suggested Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
         assert "Conversation History" in provider_exchange.EXCHANGE_SKIP_FOLDERS
         assert "Calendar" in provider_exchange.EXCHANGE_SKIP_FOLDERS
+        assert "Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
 
     def test_exchange_skip_folders_count(self):
         """Test that EXCHANGE_SKIP_FOLDERS has expected number of entries."""
         # This test will catch if folders are accidentally added/removed
-        assert len(provider_exchange.EXCHANGE_SKIP_FOLDERS) == 3
+        assert len(provider_exchange.EXCHANGE_SKIP_FOLDERS) == 4
 
     def test_inbox_not_in_skip_folders(self):
         """Test that INBOX is not in EXCHANGE_SKIP_FOLDERS."""

test/providers/test_provider_gmail.py

@@ -12,9 +12,9 @@ import os
 import sys
 from unittest.mock import MagicMock, patch
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import provider_gmail
+from providers import provider_gmail
 
 
 class TestIsLabelFolder:
@@ -181,7 +181,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        import imap_common
+        from utils import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),
@@ -221,7 +221,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        import imap_common
+        from utils import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),
@@ -251,7 +251,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        import imap_common
+        from utils import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),

test/test_imap_backup.py

@@ -17,10 +17,10 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import backup_imap_emails
-import imap_common
-import provider_gmail
+import imap_backup as backup_imap_emails
 from conftest import temp_argv, temp_env
+from providers import provider_gmail
+from utils import imap_common
 
 
 def _mock_imap_env(port):

test/test_imap_compare.py

@@ -18,7 +18,7 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import compare_imap_folders
+import imap_compare as compare_imap_folders
 from conftest import temp_argv, temp_env
 
 

test/test_imap_count.py

@@ -16,9 +16,9 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import count_imap_emails
-import imap_common
+import imap_count as count_imap_emails
 from conftest import temp_argv, temp_env
+from utils import imap_common
 
 
 def _mock_imap_env(port):
@@ -196,7 +196,6 @@ class TestImapCommonHelpers:
         with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
             conn = imap_common.get_imap_connection("host", "user", oauth2_token="token")
 
-        assert isinstance(conn, FakeIMAP)
         assert conn.auth_called is True
         assert conn.login_called is False
 
@@ -217,7 +216,6 @@ class TestImapCommonHelpers:
         with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
             conn = imap_common.get_imap_connection("host", "user", password="pass")
 
-        assert isinstance(conn, FakeIMAP)
         assert conn.login_called is True
         assert conn.auth_called is False
 

test/test_imap_migrate.py

@@ -19,9 +19,9 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_common
-import migrate_imap_emails
+import imap_migrate as migrate_imap_emails
 from conftest import temp_argv, temp_env
+from utils import imap_common
 
 
 def _mock_migrate_env(src_port, dest_port):
@@ -285,7 +285,7 @@ class TestCacheHitWithoutLock:
 
         src.select('"INBOX"', readonly=False)
 
-        existing_dest_msg_ids = {msg_id}
+        existing_dest_msg_ids_by_folder = {"INBOX": {msg_id}}
         success, _src, _dest, deleted = migrate_imap_emails.process_single_uid(
             src,
             dest,
@@ -298,7 +298,7 @@ class TestCacheHitWithoutLock:
             None,
             True,
             False,
-            existing_dest_msg_ids=existing_dest_msg_ids,
+            existing_dest_msg_ids_by_folder=existing_dest_msg_ids_by_folder,
             existing_dest_msg_ids_lock=None,
             progress_cache_path=None,
             progress_cache_data=None,
@@ -357,16 +357,18 @@ class TestMigrateErrorHandling:
 
         from unittest.mock import patch
 
+        original_get_imap_connection = imap_common.get_imap_connection
+
         def side_effect(host, user, pwd, oauth2_token=None):
             if threading.current_thread() is threading.main_thread():
-                return imap_common.get_imap_connection(host, user, pwd, oauth2_token)
+                return original_get_imap_connection(host, user, pwd, oauth2_token)
             return None  # Simulate connection failure in worker
 
         env = _mock_migrate_env(p1, p2)
         with (
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
-            patch("imap_common.get_imap_connection", side_effect=side_effect),
+            patch("utils.imap_common.get_imap_connection", side_effect=side_effect),
         ):
             # Should not crash, but log error
             migrate_imap_emails.main()
@@ -442,7 +444,7 @@ class TestMigrateErrorHandling:
 
         env = _mock_migrate_env(p1, p2)
         with (
-            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid),
+            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid, autospec=True),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
         ):
@@ -463,7 +465,7 @@ class TestMigrateErrorHandling:
             "DEST_IMAP_USERNAME": "dest_user",
             "DEST_IMAP_PASSWORD": "p",
         }
-        with patch("imap_common.get_imap_connection", return_value=None), temp_env(env):
+        with patch("utils.imap_common.get_imap_connection", return_value=None), temp_env(env):
             with pytest.raises(SystemExit) as exc:
                 migrate_imap_emails.main()
         assert exc.value.code == 1
@@ -480,7 +482,7 @@ class TestMigrateErrorHandling:
         env = _mock_migrate_env(p1, p2)
         with (
             patch(
-                "imap_common.load_progress_cache",
+                "utils.imap_common.load_progress_cache",
                 lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
             ),
             temp_env(env),
@@ -507,7 +509,7 @@ class TestTrashHandling:
         env = _mock_migrate_env(p1, p2)
         env["DELETE_FROM_SOURCE"] = "true"
         with (
-            patch("imap_common.detect_trash_folder", return_value="Trash"),
+            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py"]),
         ):
@@ -529,7 +531,7 @@ class TestTrashHandling:
         env = _mock_migrate_env(p1, p2)
         env["DELETE_FROM_SOURCE"] = "true"
         with (
-            patch("imap_common.detect_trash_folder", return_value="Trash"),
+            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
         ):

test/test_imap_restore.py

@@ -18,10 +18,9 @@ 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
+import imap_restore as restore_imap_emails
 from conftest import temp_argv, temp_env
+from utils import imap_common, restore_cache
 
 
 def _mock_restore_env(port):
@@ -509,45 +508,12 @@ class TestRestoreE2EHelpers:
             "INBOX",
             b"Subject: Upload\r\nMessage-ID: <up@test>\r\n\r\nBody",
             '"15-Jan-2024 10:30:00 +0000"',
-            "<up@test>",
         )
 
         assert result == restore_imap_emails.UploadResult.SUCCESS
         assert len(server.folders["INBOX"]) == 1
         conn.logout()
 
-    def test_upload_email_duplicate(self, single_mock_server):
-        dest_data = {"INBOX": [b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody"]}
-        server, port = single_mock_server(dest_data)
-
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        result = restore_imap_emails.upload_email(
-            conn,
-            "INBOX",
-            b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody",
-            '"15-Jan-2024 10:30:00 +0000"',
-            "<dup@test>",
-            check_duplicate=True,
-        )
-
-        assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS
-        assert len(server.folders["INBOX"]) == 1
-        conn.logout()
-
-    def test_email_exists_in_folder(self, single_mock_server):
-        dest_data = {"INBOX": [b"Subject: Exists\r\nMessage-ID: <exists@test>\r\n\r\nBody"]}
-        _server, port = single_mock_server(dest_data)
-
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-        conn.select('"INBOX"')
-
-        assert restore_imap_emails.email_exists_in_folder(conn, "<exists@test>") is True
-        assert restore_imap_emails.email_exists_in_folder(conn, "<missing@test>") is False
-        conn.logout()
-
     def test_sync_flags_on_existing(self, single_mock_server):
         dest_data = {
             "INBOX": [{"uid": 1, "flags": set(), "content": b"Subject: Flag\r\nMessage-ID: <flag@test>\r\n\r\nBody"}]
@@ -787,8 +753,8 @@ Body content.
 
         # Track calls to record_progress
         with (
-            patch("restore_cache.record_progress") as mock_record_progress,
-            patch("imap_common.get_imap_connection", return_value=mock_conn),
+            patch("utils.restore_cache.record_progress") as mock_record_progress,
+            patch("utils.imap_common.get_imap_connection", return_value=mock_conn),
         ):
             env = {
                 "DEST_IMAP_HOST": "localhost",

test/test_imap_retry.py

@@ -0,0 +1,188 @@
+"""
+Tests for imap_retry.py
+
+Tests cover:
+- Transient error detection
+- ConnectionProxy transparent proxying
+- Retry with exponential backoff on transient errors
+- Pass-through for non-retryable methods
+- Pass-through for non-transient errors
+"""
+
+import imaplib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools")))
+
+from mock_imap_server import start_server_thread as start_mock_server
+from utils import imap_retry
+
+
+class TestConnectionProxy:
+    @pytest.fixture(scope="class")
+    def imap_server_info(self):
+        folders = {"INBOX": []}
+        server, port = start_mock_server(folders)
+        yield server, port
+        server.shutdown()
+        server.server_close()
+
+    @pytest.fixture
+    def imap_conn(self, imap_server_info):
+        server, port = imap_server_info
+        client = imaplib.IMAP4("127.0.0.1", port)
+        client.login("user", "pass")
+        yield client
+        try:
+            client.logout()
+        except:
+            pass
+
+    @pytest.fixture
+    def captured_logs(self):
+        logs = []
+
+        def log_fn(msg):
+            logs.append(msg)
+
+        return logs, log_fn
+
+    @pytest.fixture
+    def proxy(self, imap_conn, captured_logs):
+        # Short wait to speed up tests
+        logs, log_fn = captured_logs
+        return imap_retry.ConnectionProxy(imap_conn, max_retries=3, initial_wait=0.01, log_fn=log_fn)
+
+    def test_transient_retry_logic_unavailable(self, proxy, captured_logs):
+        """Test detection of 'NO [UNAVAILABLE]' error pattern."""
+        try:
+            logs, _ = captured_logs
+            # Server will echo "NO [UNAVAILABLE] Server Busy" and fail continuously
+            typ, data = proxy.select('"NO [UNAVAILABLE] Server Busy"')
+            assert typ == "NO"
+            # Should have retried max_retries-1 times before returning last error
+            assert len(logs) == 2  # Retries on attempt 1 and 2, then fail on 3
+            assert "attempt 1/3" in logs[0]
+            assert "attempt 2/3" in logs[1]
+        except Exception:
+            raise
+
+    def test_transient_retry_logic_server_busy(self, proxy, captured_logs):
+        """Test detection of 'NO Server Busy' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO Server Busy"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_try_again(self, proxy, captured_logs):
+        """Test detection of 'try again' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO try again later"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_throttled(self, proxy, captured_logs):
+        """Test detection of 'THROTTLED' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [THROTTLED]"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_non_transient_no_retry_empty(self, proxy, captured_logs):
+        """Test mismatch on empty/irrelevant data."""
+        logs, _ = captured_logs
+        # EMPTY arg causes mock to return OK with 0 items but generally here we want an error condition
+        # Mock server "EMPTY" folder returns OK.
+        # We need an error that is NOT transient.
+        # "NO [UNKNOWN]"
+        typ, data = proxy.select('"NO [UNKNOWN]"')
+        assert typ == "NO"
+        assert len(logs) == 0
+
+    def test_ok_response_returns_immediately(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"INBOX"')
+        assert typ == "OK"
+        assert len(logs) == 0
+
+    def test_non_transient_error_not_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [AUTHENTICATIONFAILED]"')
+        assert typ == "NO"
+        assert b"AUTHENTICATIONFAILED" in data[0]
+        assert len(logs) == 0
+
+    def test_transient_error_retried_then_succeeds(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        # Retry 2 times (fail count 2, so 3rd succeeds)
+        # RETRY_2 triggers 2 failures then OK.
+        typ, data = proxy.select('"RETRY_2"')
+        assert typ == "OK"
+        assert len(logs) == 2
+        assert "attempt 1/3" in logs[0]
+        assert "attempt 2/3" in logs[1]
+
+    def test_exponential_backoff_logic(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.fetch("1", "(BODY RETRY_2)")
+        assert typ == "OK"
+        assert len(logs) == 2
+        # Verify message content for backoff?
+        # approximate wait time is hard to verify without mocking sleep, but we verify calls were made.
+
+    def test_max_retries_exhausted_returns_last_error(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        # Fail 5 times. Max retries is 3. Result should be error.
+        typ, data = proxy.store("1", "+FLAGS", "(\\Seen RETRY_5)")
+        assert typ == "NO"
+        assert b"UNAVAILABLE" in data[0]
+        assert len(logs) == 2
+
+    def test_non_retryable_method_passes_through(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        res = proxy.capability()
+        assert res[0] == "OK"
+        assert len(logs) == 0
+
+    def test_non_callable_attribute_passes_through(self, proxy):
+        assert proxy.state == "AUTH"
+
+    def test_non_tuple_return_passes_through(self):
+        class DummyConn:
+            def noop(self):
+                return "unexpected"
+
+        d = DummyConn()
+        p = imap_retry.ConnectionProxy(d)
+        assert p.noop() == "unexpected"
+
+    def test_uid_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.uid("SEARCH", "RETRY_1")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_append_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.append('"RETRY_1"', None, None, b"data")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_max_retries_zero_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=0)
+
+    def test_max_retries_negative_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=-1)
+
+    def test_initial_wait_negative_raises(self):
+        with pytest.raises(ValueError, match="initial_wait must be >= 0"):
+            imap_retry.ConnectionProxy(None, initial_wait=-1)

test/test_migrate_resume.py

@@ -6,9 +6,9 @@ import sys
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import migrate_imap_emails
-import restore_cache
+import imap_migrate as migrate_imap_emails
 from conftest import temp_argv, temp_env
+from utils import restore_cache
 
 
 def _run_migrate(cache_dir, src_port, dest_port):

test/test_migrate_with_cache.py

@@ -10,11 +10,10 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_common
-import imap_session
-import migrate_imap_emails
-import restore_cache
+import imap_migrate as migrate_imap_emails
 from conftest import make_mock_connection, temp_argv, temp_env
+from core import imap_session
+from utils import imap_common, restore_cache
 
 
 def _run_migrate(cache_dir, src_port, dest_port, full_migrate=False, extra_env=None):
@@ -168,7 +167,7 @@ class TestMigrationCache:
         dest.login("dest", "p")
 
         with (
-            patch("imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
             patch.object(
                 imap_common,
                 "load_progress_cache",
@@ -215,7 +214,7 @@ class TestMigrationCache:
         dest.login("dest", "p")
 
         with (
-            patch("imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
             patch.object(
                 restore_cache,
                 "get_cached_message_ids",

test/utils/test_imap_common.py

@@ -18,10 +18,10 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
 
-import imap_common
 from conftest import temp_env
+from utils import imap_common
 
 
 class TestVerifyEnvVars:
@@ -320,48 +320,6 @@ class TestDetectTrashFolder:
         assert result is None
 
 
-class TestMessageExistsInFolder:
-    """Tests for message_exists_in_folder function."""
-
-    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)
-        assert result is False
-
-    def test_search_fails(self):
-        """Test returns False when search fails."""
-        mock_conn = Mock()
-        mock_conn.search.return_value = ("NO", [])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-    def test_no_matches(self):
-        """Test returns False when no matches found."""
-        mock_conn = Mock()
-        mock_conn.search.return_value = ("OK", [b""])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-    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"])
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is True
-
-    def test_search_exception(self):
-        """Test returns False when search raises an exception."""
-        mock_conn = Mock()
-        mock_conn.search.side_effect = Exception("Connection error")
-
-        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
-        assert result is False
-
-
 class TestExtractMessageId:
     """Tests for extract_message_id function."""
 
@@ -807,3 +765,77 @@ class TestGetMessageIdsInFolder:
 
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert set(result.values()) == {"<valid@example.com>"}
+
+
+class TestLoadFolderMsgIds:
+    """Tests for load_folder_msg_ids() using the mock IMAP server."""
+
+    def test_returns_none_when_disabled(self):
+        """Returns None if dict or lock is None (tracking disabled)."""
+        mock_conn = Mock()
+        import threading
+
+        lock = threading.Lock()
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, lock) is None
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", {}, None) is None
+        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, None) is None
+
+    def test_fetches_message_ids_from_server(self, single_mock_server):
+        """Fetches Message-IDs from server and caches them."""
+        import imaplib
+        import threading
+
+        server_data = {
+            "TestFolder": [
+                b"Subject: A\r\nMessage-ID: <a@test.com>\r\n\r\nBody A",
+                b"Subject: B\r\nMessage-ID: <b@test.com>\r\n\r\nBody B",
+            ]
+        }
+        _server, port = single_mock_server(server_data)
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
+        assert result == {"<a@test.com>", "<b@test.com>"}
+        assert "TestFolder" in by_folder
+
+        # Second call returns the same cached set object
+        result2 = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
+        assert result2 is result
+        conn.logout()
+
+    def test_empty_folder(self, single_mock_server):
+        """Returns empty set for a folder with no messages."""
+        import imaplib
+        import threading
+
+        _server, port = single_mock_server({"INBOX": []})
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "INBOX", by_folder, lock)
+        assert result == set()
+        conn.logout()
+
+    def test_creates_folder_if_missing(self, single_mock_server):
+        """Creates folder on first access if it doesn't exist."""
+        import imaplib
+        import threading
+
+        server, port = single_mock_server({"INBOX": []})
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        lock = threading.Lock()
+        by_folder: dict[str, set[str]] = {}
+
+        result = imap_common.load_folder_msg_ids(conn, "NewFolder", by_folder, lock)
+        assert result == set()
+        assert "NewFolder" in server.folders
+        conn.logout()

test/utils/test_imap_compress.py

@@ -0,0 +1,394 @@
+"""Tests for IMAP COMPRESS=DEFLATE support."""
+
+import zlib
+from unittest.mock import MagicMock, patch
+
+from utils import imap_compress
+
+
+class TestCompressedSocket:
+    """Tests for the _CompressedSocket wrapper."""
+
+    def _make_compressed(self, plaintext: bytes) -> bytes:
+        """Compress data using raw DEFLATE with Z_SYNC_FLUSH (server-side)."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        return c.compress(plaintext) + c.flush(zlib.Z_SYNC_FLUSH)
+
+    def test_sendall_compresses_data(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        ws.sendall(b"AAAA0 SELECT INBOX\r\n")
+
+        sock.sendall.assert_called_once()
+        compressed = sock.sendall.call_args[0][0]
+        # Verify round-trip: decompress the sent data
+        d = zlib.decompressobj(-zlib.MAX_WBITS)
+        assert d.decompress(compressed) == b"AAAA0 SELECT INBOX\r\n"
+
+    def test_send_compresses_and_returns_length(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        result = ws.send(b"hello")
+        assert result == 5
+        sock.sendall.assert_called_once()
+
+    def test_recv_decompresses_data(self):
+        plaintext = b"* OK IMAP ready\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        result = ws.recv(4096)
+
+        assert result == plaintext
+
+    def test_recv_buffers_excess(self):
+        plaintext = b"A" * 100
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        # Request only 30 bytes — rest should be buffered
+        chunk1 = ws.recv(30)
+        assert chunk1 == b"A" * 30
+        assert len(ws._recv_buf) == 70
+
+        # Next recv should return from buffer without hitting socket
+        chunk2 = ws.recv(70)
+        assert chunk2 == b"A" * 70
+        sock.recv.assert_called_once()  # Only one real socket read
+
+    def test_recv_empty_returns_empty(self):
+        sock = MagicMock()
+        sock.recv.return_value = b""
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == b""
+
+    def test_initial_compressed_data_decompressed(self):
+        plaintext = b"* 5 EXISTS\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock, initial_compressed=compressed)
+
+        # Should return the pre-decompressed data without hitting socket
+        result = ws.recv(4096)
+        assert result == plaintext
+        sock.recv.assert_not_called()
+
+    def test_shutdown_proxied(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+        ws.shutdown(2)
+        sock.shutdown.assert_called_once_with(2)
+
+    def test_close_proxied(self):
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+        ws.close()
+        sock.close.assert_called_once()
+
+    def test_makefile_returns_buffered_reader(self):
+        import io
+
+        plaintext = b"* OK ready\r\n"
+        compressed = self._make_compressed(plaintext)
+
+        sock = MagicMock()
+        sock.recv.return_value = compressed
+
+        ws = imap_compress._CompressedSocket(sock)
+        f = ws.makefile("rb")
+        assert isinstance(f, io.BufferedReader)
+        assert f.readline() == plaintext
+
+    def test_getattr_proxied(self):
+        sock = MagicMock()
+        sock.gettimeout.return_value = 30
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.gettimeout() == 30
+
+    def test_multiple_sendall_share_compressor_state(self):
+        """Verify the compressor maintains state across calls (streaming)."""
+        sock = MagicMock()
+        ws = imap_compress._CompressedSocket(sock)
+
+        ws.sendall(b"first ")
+        ws.sendall(b"second")
+
+        assert sock.sendall.call_count == 2
+        # Both calls should produce decompressible data
+        d = zlib.decompressobj(-zlib.MAX_WBITS)
+        result = b""
+        for call in sock.sendall.call_args_list:
+            result += d.decompress(call[0][0])
+        assert result == b"first second"
+
+    def test_recv_loops_on_partial_deflate_block(self):
+        """Decompress can return b'' for partial blocks; recv must not return that as EOF."""
+        sock = MagicMock()
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+
+        # Build a compressed stream, then split it mid-block so the first
+        # chunk decompresses to b"" and the second completes the block.
+        full_compressed = c.compress(b"hello\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+        # Split at byte 1 — the first fragment is too short to produce output
+        part1 = full_compressed[:1]
+        part2 = full_compressed[1:]
+        sock.recv.side_effect = [part1, part2]
+
+        ws = imap_compress._CompressedSocket(sock)
+        result = ws.recv(4096)
+
+        assert result == b"hello\r\n"
+        assert sock.recv.call_count == 2
+
+    def test_recv_returns_empty_on_real_eof(self):
+        """True socket EOF (b'') should still return b''."""
+        sock = MagicMock()
+        sock.recv.return_value = b""
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == b""
+
+    def test_recv_multiple_chunks(self):
+        """Simulate receiving data across multiple socket reads."""
+        chunk1 = b"line one\r\n"
+        chunk2 = b"line two\r\n"
+
+        sock = MagicMock()
+        # Server compressor maintains state across writes
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed1 = c.compress(chunk1) + c.flush(zlib.Z_SYNC_FLUSH)
+        compressed2 = c.compress(chunk2) + c.flush(zlib.Z_SYNC_FLUSH)
+        sock.recv.side_effect = [compressed1, compressed2]
+
+        ws = imap_compress._CompressedSocket(sock)
+        assert ws.recv(4096) == chunk1
+        assert ws.recv(4096) == chunk2
+
+
+class TestEnableCompression:
+    """Tests for the enable_compression() function."""
+
+    def test_returns_false_when_capability_missing(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "IDLE")
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_capability_missing_bytes(self):
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"IDLE")
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_no_capabilities_attr(self):
+        conn = MagicMock(spec=[])
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_command_fails(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("NO", [b"compression not available"])
+
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_returns_false_when_command_raises(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.side_effect = Exception("network error")
+
+        assert imap_compress.enable_compression(conn) is False
+
+    def test_enables_compression_with_bytes_capabilities(self):
+        """Real imaplib stores capabilities as bytes; verify they are handled."""
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        result = imap_compress.enable_compression(conn)
+
+        assert result is True
+        assert isinstance(conn.sock, imap_compress._CompressedSocket)
+
+    def test_enables_compression_on_success(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        result = imap_compress.enable_compression(conn)
+
+        assert result is True
+        assert isinstance(conn.sock, imap_compress._CompressedSocket)
+
+    def test_drains_readbuf_on_success(self):
+        # Simulate leftover compressed data in the read buffer
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 1 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = [compressed]
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        # readbuf should have been cleared
+        assert len(conn._readbuf) == 0
+        # The compressed leftover should be pre-decompressed in the wrapper
+        assert conn.sock._recv_buf == b"* 1 EXISTS\r\n"
+
+    def test_drains_bytearray_readbuf(self):
+        """Real imaplib uses bytearray for _readbuf; verify drain handles it."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 2 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = bytearray(compressed)
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        assert len(conn._readbuf) == 0
+        assert conn.sock._recv_buf == b"* 2 EXISTS\r\n"
+
+    def test_drains_bytes_readbuf(self):
+        """Some imaplib versions use bytes for _readbuf; verify drain handles it."""
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 3 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = bytes(compressed)
+        conn.sock = MagicMock()
+
+        imap_compress.enable_compression(conn)
+
+        assert conn._readbuf == b""
+        assert conn.sock._recv_buf == b"* 3 EXISTS\r\n"
+
+    def test_drains_file_buffer(self):
+        """BufferedReader may have read-ahead bytes beyond the OK response."""
+        import io
+
+        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
+        compressed = c.compress(b"* 4 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = b""
+        conn.sock = MagicMock()
+
+        # Simulate a BufferedReader with read-ahead compressed bytes
+        old_file = io.BufferedReader(io.BytesIO(compressed))
+        old_file.peek(65536)  # fill the buffer
+        conn._file = old_file
+
+        imap_compress.enable_compression(conn)
+
+        assert conn.sock._recv_buf == b"* 4 EXISTS\r\n"
+
+    def test_calls_log_fn(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = []
+        conn.sock = MagicMock()
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "compression enabled" in messages[0]
+
+    def test_logs_when_not_supported(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1",)  # No COMPRESS
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "not supported" in messages[0]
+
+    def test_logs_when_rejected(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("NO", [b"not available"])
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "rejected" in messages[0]
+
+    def test_logs_when_negotiation_fails(self):
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.side_effect = Exception("network error")
+
+        messages = []
+        imap_compress.enable_compression(conn, log_fn=messages.append)
+
+        assert len(messages) == 1
+        assert "failed" in messages[0]
+
+    def test_recreates_file(self):
+        import io
+
+        conn = MagicMock()
+        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
+        conn._simple_command.return_value = ("OK", [b"ok"])
+        conn._readbuf = []
+        real_sock = MagicMock()
+        conn.sock = real_sock
+
+        imap_compress.enable_compression(conn)
+
+        # _file should be a BufferedReader over the compressed socket
+        assert isinstance(conn._file, io.BufferedReader)
+
+
+class TestIntegrationWithGetImapConnection:
+    """Test that enable_compression is called during connection setup."""
+
+    @patch("utils.imap_common.imaplib")
+    @patch("utils.imap_compress.enable_compression")
+    def test_compression_attempted_after_login(self, mock_compress, mock_imaplib):
+        from utils import imap_common
+
+        mock_conn = MagicMock()
+        mock_imaplib.IMAP4_SSL.return_value = mock_conn
+
+        imap_common.get_imap_connection("imap.example.com", "user", password="pass")
+
+        mock_conn.login.assert_called_once()
+        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)
+
+    @patch("utils.imap_common.imaplib")
+    @patch("utils.imap_compress.enable_compression")
+    def test_compression_attempted_after_oauth2(self, mock_compress, mock_imaplib):
+        from utils import imap_common
+
+        mock_conn = MagicMock()
+        mock_imaplib.IMAP4_SSL.return_value = mock_conn
+
+        imap_common.get_imap_connection("imap.example.com", "user", oauth2_token="token123")
+
+        mock_conn.authenticate.assert_called_once()
+        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)

tools/mock_imap_server.py

@@ -1,4 +1,3 @@
-import re
 import socketserver
 import threading
 
@@ -15,7 +14,15 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
     def handle(self):
         self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n")
         self.selected_folder = None
-        self.current_folders = self.server.folders
+
+        # Check if server has folders, if not, initialize
+        if hasattr(self.server, "folders"):
+            self.current_folders = self.server.folders
+        else:
+            # Just in case it's used standalone without the threaded wrapper
+            self.current_folders = {"INBOX": []}
+
+        self.retry_state = {}  # key -> count
 
         while True:
             try:
@@ -31,9 +38,73 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 cmd = parts[1].upper()
                 args = parts[2] if len(parts) > 2 else ""
 
+                # Check for RETRY_N injection in args (for testing imap_retry)
+                # Pattern: RETRY_N where N is integer
+                import re
+
+                retry_match = re.search(r"RETRY_(\d+)", args)
+                if retry_match:
+                    count = int(retry_match.group(1))
+                    # Create a unique key for this command invocation context
+                    # Just using command + args is roughly sufficient
+                    key = f"{cmd}_{args}"
+                    current_fails = self.retry_state.get(key, 0)
+
+                    if current_fails < count:
+                        # Logic Fix: Only increment if we are going to fail
+                        self.retry_state[key] = current_fails + 1
+                        self.send_response(tag, "NO [UNAVAILABLE] Server Busy (Simulated)")
+                        continue
+                    # Else fall through to normal command processing
+
                 if cmd == "LOGIN":
                     self.send_response(tag, "OK LOGIN completed")
 
+                elif cmd == "APPEND":
+                    # Check for literal size
+                    import re
+
+                    size_match = re.search(r"\{(\d+)\}$", args)
+                    data = b""
+                    if size_match:
+                        size = int(size_match.group(1))
+                        # Send continuation
+                        self.wfile.write(b"+\r\n")
+                        self.wfile.flush()
+                        # Read data
+                        data = self.rfile.read(size)
+
+                    # Extract mailbox name to store the message
+                    # Simplistic parsing: check for quoted or unquoted first argument
+                    mailbox = "INBOX"  # Default fallback
+                    clean_args = args.lstrip()
+                    if clean_args.startswith('"'):
+                        end_quote = clean_args.find('"', 1)
+                        if end_quote != -1:
+                            mailbox = clean_args[1:end_quote]
+                    else:
+                        mailbox = clean_args.split(" ")[0]
+
+                    if mailbox not in self.current_folders:
+                        self.current_folders[mailbox] = []
+
+                    # Parse flags if present: look for (...)
+                    msg_flags = set()
+                    flag_match = re.search(r"\(([^)]+)\)", args)
+                    if flag_match:
+                        # e.g. (\Seen \Deleted)
+                        flag_content = flag_match.group(1)
+                        msg_flags = set(flag_content.split())
+
+                    # Store message
+                    new_uid = len(self.current_folders[mailbox]) + 1
+                    msg_obj = {"uid": new_uid, "flags": msg_flags, "content": data}
+                    self.current_folders[mailbox].append(msg_obj)
+
+                    # Always succeed for test if we passed retry logic
+                    self.wfile.write(tag.encode() + b" OK [APPENDUID 1 100] APPEND completed\r\n")
+                    self.wfile.flush()
+
                 elif cmd == "LOGOUT":
                     self.send_response(tag, "OK LOGOUT completed")
                     break
@@ -46,8 +117,12 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                     # Minimal XOAUTH2 support for tests.
                     self.wfile.write(b"+ \r\n")
                     _ = self.rfile.readline()
+                    # Also consume potential retry line? No.
                     self.send_response(tag, "OK AUTHENTICATE completed")
 
+                elif cmd == "NOOP":
+                    self.send_response(tag, "OK NOOP completed")
+
                 elif cmd == "LIST":
                     for folder in self.current_folders:
                         self.wfile.write(f'* LIST (\\HasNoChildren) "/" "{folder}"\r\n'.encode())
@@ -55,23 +130,40 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                 elif cmd == "SELECT":
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    # For retry testing: if folder starts with NO, fail immediately
+                    if folder.startswith("NO"):
+                        # e.g. "NO [UNAVAILABLE]"
+                        self.send_response(tag, folder)
+                        continue
+
+                    # Allow RETRY_N folders to succeed if they passed the retry intercept logic
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))  # Default empty for retry folders
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
                         self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n")
                         self.send_response(tag, "OK [READ-WRITE] SELECT completed")
+                    elif folder == "EMPTY":
+                        # Special case for "empty_data" test in retry
+                        self.wfile.write(b"* 0 EXISTS\r\n")
+                        self.send_response(tag, "OK SELECT completed")
                     else:
                         self.send_response(tag, "NO [NONEXISTENT] Folder not found")
 
                 elif cmd == "EXAMINE":
                     # EXAMINE is like SELECT but read-only
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
@@ -412,7 +504,11 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 else:
                     self.send_response(tag, "BAD Command not recognized")
 
-            except Exception:
+            except Exception as e:
+                print(f"MockServer EXCEPTION: {e}")
+                import traceback
+
+                traceback.print_exc()
                 break
 
     def send_response(self, tag, message):
@@ -438,9 +534,19 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
             self.folders = {"INBOX": []}
 
 
-def start_server_thread(port=10143, initial_folders=None):
+def start_server_thread(port=0, initial_folders=None):
+    # Use port 0 to let OS select free port
+    if isinstance(port, dict):
+        # Handle legacy call where port was inadvertently passed as dict
+        initial_folders = port
+        port = 0
+
     server = MockIMAPServer(("localhost", port), MockIMAPHandler, initial_folders)
+
+    # Get the actual port if 0 was used
+    actual_port = server.server_address[1]
+
     t = threading.Thread(target=server.serve_forever)
     t.daemon = True
     t.start()
-    return t, server
+    return server, actual_port