From 7916365816d982f3aaf0d6239ec01e0578a4236c Mon Sep 17 00:00:00 2001 From: Shautvast Date: Sun, 22 Feb 2026 18:55:32 +0100 Subject: [PATCH] Add message store and refactor main into bridge startup Step 5: in-memory MessageStore mapping IMAP seq numbers to ProtonMail IDs. Oldest-first ordering, mark/expunge, range queries, subject/sender search. Refactor main.rs: extract unlock_key_pool(), build SharedState (Arc), load initial inbox, then wait for Ctrl-C. Ready to plug in IMAP/SMTP servers. Co-Authored-By: Claude Sonnet 4.6 --- Cargo.lock | 1 + proton-bridge/Cargo.toml | 2 +- proton-bridge/src/main.rs | 279 +++++++++++++++++++------------------ proton-bridge/src/store.rs | 104 ++++++++++++++ 4 files changed, 247 insertions(+), 139 deletions(-) create mode 100644 proton-bridge/src/store.rs diff --git a/Cargo.lock b/Cargo.lock index 70e0bfb..ef575e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4015,6 +4015,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", diff --git a/proton-bridge/Cargo.toml b/proton-bridge/Cargo.toml index 4a10bba..0d5df4f 100644 --- a/proton-bridge/Cargo.toml +++ b/proton-bridge/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] -tokio = { version = "1", features = ["net", "io-util", "rt-multi-thread", "macros"] } +tokio = { version = "1", features = ["net", "io-util", "rt-multi-thread", "macros", "signal"] } reqwest = { version = "0.12", default-features = false, features = ["json", "native-tls"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" diff --git a/proton-bridge/src/main.rs b/proton-bridge/src/main.rs index 6f63241..5e864a6 100644 --- a/proton-bridge/src/main.rs +++ b/proton-bridge/src/main.rs @@ -3,51 +3,108 @@ mod auth; mod config; mod crypto; 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, + /// 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>; + +// ── 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); } + 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); } + 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); } + Err(e) => { + eprintln!("Authentication failed: {e}"); + std::process::exit(1); + } }; - println!("Session UID: {}", session.uid); + println!("Authenticated as {}", config.proton.username); let api = ApiClient::new(&client, &session); - // ── Step 3: list inbox ──────────────────────────────────────────────────── - let (messages, total) = match api.list_messages(LABEL_INBOX, 0, 10).await { + // 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); } + Err(e) => { + eprintln!("list_messages failed: {e}"); + std::process::exit(1); + } }; - println!("\nInbox ({total} total):"); - for m in &messages { - println!( - " [{}] {:50} — {}", - if m.unread == 1 { "UNREAD" } else { " read" }, - m.subject, - m.sender.address, - ); - } + let mut store = MessageStore::new(); + store.load_all(messages); + println!("Inbox: {} messages ({total} total)", store.count()); - // ── Step 4: decrypt first message body ─────────────────────────────────── - let first = match messages.first() { - Some(m) => m, - None => { println!("\nInbox is empty — nothing to decrypt."); return; } - }; + let _state: SharedState = Arc::new(Mutex::new(BridgeState { + store, + key_pool, + http_client: client, + session, + config: config.clone(), + })); + 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, 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(), @@ -55,154 +112,100 @@ async fn main() { api.get_key_salts(), ); - let user_keys = user_keys_res.unwrap_or_else(|e| { eprintln!("get_user_keys: {e}"); vec![] }); - let addresses = addresses_res.unwrap_or_else(|e| { eprintln!("get_addresses: {e}"); vec![] }); + 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}"))?; - // key_salts requires the 'locked' OAuth scope. The server may revoke it even - // when our cached session still claims to have it. If we get 403, purge the - // session so the next run does a full SRP+TOTP login that re-grants 'locked'. let key_salts = match key_salts_res { Ok(salts) => salts, Err(e) => { - eprintln!("get_key_salts: {e}"); if e.contains("403") { auth::Session::delete(); - eprintln!("Session 'locked' scope expired server-side — session.json deleted."); - eprintln!("Re-run; you will be prompted for your TOTP code to regain 'locked' scope."); - std::process::exit(1); + 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![] } }; - // In two-password mode the mailbox password unlocks the keys. - let password = config.proton.mailbox_password + let password = config + .proton + .mailbox_password .as_deref() .unwrap_or(&config.proton.password); - // Derive passphrase for a key ID via bcrypt(password, keySalt). - // Falls back to raw password when salt is absent. + // Derive the passphrase for a key (bcrypt-31 or raw password fallback). let passphrase_for = |key_id: &str| -> String { - let salt = key_salts.iter() + 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()) + crypto::derive_key_passphrase(password, salt).unwrap_or_else(|_| password.to_string()) }; - // Phase 1: unlock user keys — try multiple passphrase candidates so we can - // identify which derivation is correct. - let mut user_private_keys: Vec<(String, crypto::PrivateKey)> = Vec::new(); + // Phase 1: unlock user keys. + let mut user_private_keys: Vec = Vec::new(); for uk in &user_keys { - if uk.active == 1 { - let salt = key_salts.iter() - .find(|s| s.id == uk.id) - .and_then(|s| s.key_salt.as_deref()) - .unwrap_or(""); - - // Derive key passphrase per ProtonMail's go-proton-api SaltForKey: - // bcrypt($2y$, cost=10, raw_salt_16_bytes, password+\0)[last 31 chars] - eprintln!(" [key-salt] raw API value: {:?} (len={})", salt, salt.len()); - - let pp = crypto::derive_key_passphrase(password, salt) - .unwrap_or_else(|_| password.to_string()); - eprintln!(" [pp] len={} last8=...{}", pp.len(), - pp.chars().rev().take(8).collect::().chars().rev().collect::()); - - let mut candidates: Vec<(&str, String)> = vec![ - ("bcrypt-31", pp), - ("raw-password", password.to_string()), - ]; - - // Manual override — set user_key_passphrase in bridge.toml to bypass derivation. - if let Some(ref override_pp) = config.proton.user_key_passphrase { - eprintln!(" [override] using user_key_passphrase from config"); - candidates.push(("config-override", override_pp.clone())); - } - - // Use the first candidate to store the parsed key structure. - // (unlock() only parses the key; actual passphrase is tested during decrypt.) - let first_pp = candidates[0].1.clone(); - match crypto::PrivateKey::unlock(&uk.private_key, &first_pp) { - Ok(k) => { - println!("user key {} parsed OK", &uk.id[..8]); - // Run the independent S2K+CFB+SHA1 diagnostic for all candidates. - let diag_candidates: Vec<(&str, &str)> = candidates - .iter() - .map(|(l, p)| (*l, p.as_str())) - .collect(); - crypto::diagnose_subkey_passphrase(&k, &diag_candidates); - // Store all candidate passphrases so token decrypt can try each. - for (label, pp) in &candidates { - match crypto::PrivateKey::unlock(&uk.private_key, pp) { - Ok(k2) => user_private_keys.push((label.to_string(), k2)), - Err(e) => eprintln!("user key {} [{}] FAILED: {}", &uk.id[..8], label, e), - } - } - } - Err(e) => eprintln!("user key {} parse FAILED: {}", &uk.id[..8], e), + 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. - // Newer accounts: the address key's passphrase is stored as a PGP-encrypted - // Token (encrypted to the user key). Decrypt the token, then unlock. - // Older accounts: address keys share the same bcrypt passphrase as user keys. + // Phase 2: unlock address keys via token (newer accounts) or shared passphrase. let mut key_pool: Vec = Vec::new(); for addr in &addresses { for ak in &addr.keys { - if ak.active == 1 { - let pp = if let Some(token) = &ak.token { - let mut token_pp = None; - for (label, uk) in &user_private_keys { - match crypto::decrypt_token(token, uk) { - Ok(p) => { - println!("token decrypted using [{}] passphrase!", label); - token_pp = Some(p); - break; - } - Err(e) => eprintln!("decrypt_token [{}] ({}): {e}", label, addr.email), - } + 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], addr.email); - key_pool.push(k); - } - Err(e) => eprintln!("address key {} ({}): {e}", &ak.id[..8], addr.email), } + 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 (messages may be encrypted to the user key). - key_pool.extend(user_private_keys.into_iter().map(|(_, k)| k)); - println!("\n{} key(s) in pool ({} address(es))", key_pool.len(), addresses.len()); - // Fetch and decrypt the first message. - let full_msg = match api.get_message(&first.id).await { - Ok(m) => m, - Err(e) => { eprintln!("get_message failed: {}", e); std::process::exit(1); } - }; - - crypto::dump_key_pool(&key_pool); - crypto::dump_message_pkesk(&full_msg.body); - - let key_refs: Vec<&crypto::PrivateKey> = key_pool.iter().collect(); - match crypto::decrypt_body(&full_msg.body, &key_refs) { - Ok(plaintext) => { - println!("\n── Decrypted body ({}) ──────────────────", full_msg.mime_type); - let preview: String = plaintext.chars().take(2000).collect(); - println!("{preview}"); - if plaintext.len() > 2000 { - println!("… ({} chars total)", plaintext.len()); - } - } - Err(e) => eprintln!("decrypt_body failed: {}", e), - } -} \ No newline at end of file + // Include user keys as fallback. + key_pool.extend(user_private_keys); + Ok(key_pool) +} diff --git a/proton-bridge/src/store.rs b/proton-bridge/src/store.rs new file mode 100644 index 0000000..1741918 --- /dev/null +++ b/proton-bridge/src/store.rs @@ -0,0 +1,104 @@ +/// Step 5: In-memory message store. +/// +/// Maps IMAP sequence numbers (1-based, contiguous) ↔ ProtonMail message IDs. +/// Messages are stored oldest-first so that seq 1 = oldest, seq N = newest. +/// The API returns messages newest-first, so `load_all` reverses the slice. +/// +/// After EXPUNGE the Vec automatically renumbers: the element that was at +/// index i is now at some new index j, and its seq number is j + 1. +use std::collections::HashSet; + +use crate::api::MessageMeta; + +pub struct MessageStore { + /// Oldest-first. `messages[i]` has IMAP sequence number `i + 1`. + messages: Vec, + /// Proton IDs of messages marked `\Deleted`, waiting for EXPUNGE. + deleted_pending: HashSet, +} + +impl MessageStore { + pub fn new() -> Self { + Self { + messages: Vec::new(), + deleted_pending: HashSet::new(), + } + } + + /// Replace the store contents with a fresh API response (newest-first). + /// Reverses to oldest-first and clears any pending deletions. + pub fn load_all(&mut self, mut messages: Vec) { + messages.reverse(); // API is newest-first; store oldest-first + self.messages = messages; + self.deleted_pending.clear(); + } + + /// Number of messages currently in the store (the IMAP EXISTS count). + pub fn count(&self) -> u32 { + self.messages.len() as u32 + } + + /// Look up the ProtonMail message ID for a given IMAP sequence number. + pub fn proton_id(&self, seq: u32) -> Option<&str> { + self.messages + .get(seq.wrapping_sub(1) as usize) + .map(|m| m.id.as_str()) + } + + /// Look up message metadata for a given IMAP sequence number. + pub fn get_meta(&self, seq: u32) -> Option<&MessageMeta> { + self.messages.get(seq.wrapping_sub(1) as usize) + } + + /// Return all (seq, meta) pairs in the inclusive range [from, to]. + /// Clamps to the actual message count. + pub fn get_range(&self, from: u32, to: u32) -> Vec<(u32, &MessageMeta)> { + let max = self.count(); + let from = from.max(1); + let to = to.min(max); + if from > to { + return vec![]; + } + (from..=to) + .filter_map(|seq| self.get_meta(seq).map(|m| (seq, m))) + .collect() + } + + /// Mark a message for deletion (STORE … +FLAGS (\Deleted)). + pub fn mark_deleted(&mut self, seq: u32) { + if let Some(m) = self.get_meta(seq) { + self.deleted_pending.insert(m.id.clone()); + } + } + + /// Remove all pending-deleted messages and return their ProtonMail IDs. + /// After this call the Vec indices automatically provide new seq numbers. + pub fn expunge(&mut self) -> Vec { + let removed: Vec = self + .messages + .iter() + .filter(|m| self.deleted_pending.contains(&m.id)) + .map(|m| m.id.clone()) + .collect(); + self.messages + .retain(|m| !self.deleted_pending.contains(&m.id)); + self.deleted_pending.clear(); + removed + } + + /// Return the seq numbers of messages whose subject, sender name, or + /// sender address contains `query` (case-insensitive substring match). + pub fn search_subject_or_from(&self, query: &str) -> Vec { + let q = query.to_lowercase(); + self.messages + .iter() + .enumerate() + .filter(|(_, m)| { + m.subject.to_lowercase().contains(&q) + || m.sender.address.to_lowercase().contains(&q) + || m.sender.name.to_lowercase().contains(&q) + }) + .map(|(i, _)| i as u32 + 1) + .collect() + } +}