tuimail/proton-bridge/src/main.rs

mod api;
mod auth;
mod config;
mod crypto;
mod imap_server;
mod srp;
mod store;

use std::sync::Arc;

use api::{ApiClient, LABEL_INBOX};
use store::MessageStore;
use tokio::sync::Mutex;

// ── Shared bridge state ───────────────────────────────────────────────────────

/// State shared between the IMAP and SMTP server tasks.
pub struct BridgeState {
    pub store: MessageStore,
    pub key_pool: Vec<crypto::PrivateKey>,
    /// The underlying reqwest client (stored so we can create ApiClient on demand).
    pub http_client: reqwest::Client,
    pub session: auth::Session,
    pub config: config::Config,
}

pub type SharedState = Arc<Mutex<BridgeState>>;

// ── Entry point ───────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() {
    env_logger::init();

    let config = match config::Config::load() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Failed to load bridge.toml: {e}");
            std::process::exit(1);
        }
    };

    let client = match auth::build_client() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Failed to build HTTP client: {e}");
            std::process::exit(1);
        }
    };

    let session = match auth::authenticate(&client, &config.proton).await {
        Ok(s) => s,
        Err(e) => {
            eprintln!("Authentication failed: {e}");
            std::process::exit(1);
        }
    };
    println!("Authenticated as {}", config.proton.username);

    let api = ApiClient::new(&client, &session);

    // Unlock user + address keys.
    let key_pool = match unlock_key_pool(&api, &config).await {
        Ok(k) => k,
        Err(e) => {
            eprintln!("{e}");
            std::process::exit(1);
        }
    };
    println!("{} key(s) in pool", key_pool.len());

    // Load the initial inbox into the message store.
    let (messages, total) = match api.list_messages(LABEL_INBOX, 0, 50).await {
        Ok(r) => r,
        Err(e) => {
            eprintln!("list_messages failed: {e}");
            std::process::exit(1);
        }
    };
    let mut store = MessageStore::new();
    store.load_all(messages);
    println!("Inbox: {} messages ({total} total)", store.count());

    let state: SharedState = Arc::new(Mutex::new(BridgeState {
        store,
        key_pool,
        http_client: client,
        session,
        config: config.clone(),
    }));

    // Spawn IMAP server.
    let imap_state = Arc::clone(&state);
    let imap_port = config.bridge.imap_port;
    tokio::spawn(async move {
        if let Err(e) = imap_server::run(imap_state, imap_port).await {
            eprintln!("IMAP server error: {e}");
        }
    });

    println!(
        "Bridge ready — IMAP :{}, SMTP :{}  (Ctrl-C to stop)",
        config.bridge.imap_port, config.bridge.smtp_port
    );

    tokio::signal::ctrl_c()
        .await
        .expect("failed to listen for ctrl-c");
    println!("Shutting down.");
}

// ── Key pool setup ────────────────────────────────────────────────────────────

async fn unlock_key_pool(
    api: &ApiClient<'_>,
    config: &config::Config,
) -> Result<Vec<crypto::PrivateKey>, String> {
    // Fetch user keys, address keys, and bcrypt salts in parallel.
    let (user_keys_res, addresses_res, key_salts_res) = tokio::join!(
        api.get_user_keys(),
        api.get_addresses(),
        api.get_key_salts(),
    );

    let user_keys = user_keys_res.map_err(|e| format!("get_user_keys: {e}"))?;
    let addresses = addresses_res.map_err(|e| format!("get_addresses: {e}"))?;

    let key_salts = match key_salts_res {
        Ok(salts) => salts,
        Err(e) => {
            if e.contains("403") {
                auth::Session::delete();
                return Err(format!(
                    "Session 'locked' scope expired — session.json deleted. \
                     Re-run to enter TOTP and regain 'locked' scope. ({e})"
                ));
            }
            eprintln!("get_key_salts: {e} — continuing without salt");
            vec![]
        }
    };

    let password = config
        .proton
        .mailbox_password
        .as_deref()
        .unwrap_or(&config.proton.password);

    // Derive the passphrase for a key (bcrypt-31 or raw password fallback).
    let passphrase_for = |key_id: &str| -> String {
        let salt = key_salts
            .iter()
            .find(|s| s.id == key_id)
            .and_then(|s| s.key_salt.as_deref())
            .unwrap_or("");
        crypto::derive_key_passphrase(password, salt).unwrap_or_else(|_| password.to_string())
    };

    // Phase 1: unlock user keys.
    let mut user_private_keys: Vec<crypto::PrivateKey> = Vec::new();
    for uk in &user_keys {
        if uk.active != 1 {
            continue;
        }
        let salt = key_salts
            .iter()
            .find(|s| s.id == uk.id)
            .and_then(|s| s.key_salt.as_deref())
            .unwrap_or("");
        let pp = crypto::derive_key_passphrase(password, salt)
            .unwrap_or_else(|_| password.to_string());
        match crypto::PrivateKey::unlock(&uk.private_key, &pp) {
            Ok(k) => {
                println!("user key {} unlocked", &uk.id[..8.min(uk.id.len())]);
                user_private_keys.push(k);
            }
            Err(e) => eprintln!("user key {} FAILED: {e}", &uk.id[..8.min(uk.id.len())]),
        }
    }

    // Phase 2: unlock address keys via token (newer accounts) or shared passphrase.
    let mut key_pool: Vec<crypto::PrivateKey> = Vec::new();
    for addr in &addresses {
        for ak in &addr.keys {
            if ak.active != 1 {
                continue;
            }
            let pp = if let Some(token) = &ak.token {
                let mut token_pp = None;
                for uk in &user_private_keys {
                    if let Ok(p) = crypto::decrypt_token(token, uk) {
                        token_pp = Some(p);
                        break;
                    }
                }
                token_pp.unwrap_or_else(|| passphrase_for(&ak.id))
            } else {
                passphrase_for(&ak.id)
            };
            match crypto::PrivateKey::unlock(&ak.private_key, &pp) {
                Ok(k) => {
                    println!(
                        "address key {} ({}) unlocked",
                        &ak.id[..8.min(ak.id.len())],
                        addr.email
                    );
                    key_pool.push(k);
                }
                Err(e) => eprintln!(
                    "address key {} ({}): {e}",
                    &ak.id[..8.min(ak.id.len())],
                    addr.email
                ),
            }
        }
    }

    // Include user keys as fallback.
    key_pool.extend(user_private_keys);
    Ok(key_pool)
}