eric / rustwarden

//! IMAP Bounce Processing Service
//!
//! This service connects to an IMAP server to retrieve and process
//! Delivery Status Notifications (DSN) / bounce messages.
//!
//! It parses DSN messages according to RFC 3464 and updates the
//! EmailDeliveryStatus table accordingly.

use crate::config::email::EmailConfig;
use crate::error::AppError;
use crate::models::email_delivery_status;
use chrono::Utc;
use native_tls::{Certificate, TlsConnector};
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
use std::net::TcpStream;
use std::sync::Arc;
use tokio::time::Duration;

/// IMAP Bounce Processing Service
///
/// Connects to IMAP to retrieve and process bounce messages (DSN).
pub struct ImapBounceService {
    config: EmailConfig,
    db: Arc<DatabaseConnection>,
}

impl ImapBounceService {
    /// Create a new ImapBounceService
    pub fn new(config: EmailConfig, db: Arc<DatabaseConnection>) -> Self {
        Self { config, db }
    }

    /// Check if IMAP bounce processing is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.imap_host.is_some()
    }

    /// Start the IMAP bounce processor as a background task
    pub fn start(self) {
        if !self.is_enabled() {
            tracing::info!("IMAP bounce processing disabled (IMAP_HOST not set)");
            return;
        }

        tokio::spawn(async move {
            let interval_secs = self.config.imap_poll_interval_seconds;
            let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
            interval.tick().await; // Skip first immediate tick

            loop {
                interval.tick().await;

                match self.process_bounces().await {
                    Ok(processed) => {
                        if processed > 0 {
                            tracing::info!(
                                "IMAP bounce processor: processed {} bounce message(s)",
                                processed
                            );
                        }
                    }
                    Err(e) => {
                        tracing::error!("IMAP bounce processor error: {}", e);
                    }
                }
            }
        });
    }

    /// Process bounce messages from IMAP
    async fn process_bounces(&self) -> Result<usize, AppError> {
        let host = self
            .config
            .imap_host
            .as_ref()
            .ok_or_else(|| AppError::Internal("IMAP host not configured".to_string()))?;
        let user = self
            .config
            .imap_user
            .as_ref()
            .ok_or_else(|| AppError::Internal("IMAP user not configured".to_string()))?;
        let password = self
            .config
            .imap_password
            .as_ref()
            .ok_or_else(|| AppError::Internal("IMAP password not configured".to_string()))?;

        // Connect to IMAP in a blocking task (imap crate is sync)
        let host_clone = host.clone();
        let port = self.config.imap_port;
        let use_tls = self.config.imap_use_tls;
        let use_starttls = self.config.imap_use_starttls;
        let accept_invalid_certs = self.config.imap_accept_invalid_certs;
        let ca_cert_path = self.config.imap_ca_cert_path.clone();
        let user_clone = user.clone();
        let password_clone = password.clone();
        let mailbox = self.config.imap_mailbox.clone();

        let db = self.db.clone();

        // Run IMAP operations in a blocking task
        let result = tokio::task::spawn_blocking(move || {
            Self::fetch_and_process_bounces(
                &host_clone,
                port,
                use_tls,
                use_starttls,
                accept_invalid_certs,
                ca_cert_path.as_deref(),
                &user_clone,
                &password_clone,
                &mailbox,
            )
        })
        .await
        .map_err(|e| AppError::Internal(format!("IMAP task error: {}", e)))?;

        let bounces = result?;

        // Update database with bounce information
        let mut processed = 0;
        for bounce in bounces {
            if let Err(e) = self.update_delivery_status(&db, &bounce).await {
                tracing::warn!("Failed to update delivery status for bounce: {}", e);
            } else {
                processed += 1;
            }
        }

        Ok(processed)
    }

    /// Connect to IMAP and fetch bounce messages (sync, runs in blocking task)
    #[allow(clippy::too_many_arguments)]
    fn fetch_and_process_bounces(
        host: &str,
        port: u16,
        use_tls: bool,
        use_starttls: bool,
        accept_invalid_certs: bool,
        ca_cert_path: Option<&str>,
        user: &str,
        password: &str,
        mailbox: &str,
    ) -> Result<Vec<BounceInfo>, AppError> {
        if use_tls {
            // Direct TLS connection (typically port 993)
            Self::fetch_bounces_tls(
                host,
                port,
                accept_invalid_certs,
                ca_cert_path,
                user,
                password,
                mailbox,
            )
        } else if use_starttls {
            // Plain connection upgraded to TLS via STARTTLS (typically port 143)
            Self::fetch_bounces_starttls(
                host,
                port,
                accept_invalid_certs,
                ca_cert_path,
                user,
                password,
                mailbox,
            )
        } else {
            // Plain connection without encryption (not recommended)
            Self::fetch_bounces_plain(host, port, user, password, mailbox)
        }
    }

    /// Fetch bounces over TLS connection (direct TLS, typically port 993)
    fn fetch_bounces_tls(
        host: &str,
        port: u16,
        accept_invalid_certs: bool,
        ca_cert_path: Option<&str>,
        user: &str,
        password: &str,
        mailbox: &str,
    ) -> Result<Vec<BounceInfo>, AppError> {
        let mut tls_builder = TlsConnector::builder();

        // Load custom CA certificate if configured
        if let Some(path) = ca_cert_path {
            let cert_pem = std::fs::read(path).map_err(|e| {
                AppError::Internal(format!("Failed to read CA certificate {}: {}", path, e))
            })?;
            let cert = Certificate::from_pem(&cert_pem).map_err(|e| {
                AppError::Internal(format!("Failed to parse CA certificate: {}", e))
            })?;
            tls_builder.add_root_certificate(cert);
            tracing::debug!("Added custom CA certificate for IMAP: {}", path);
        }

        if accept_invalid_certs {
            tls_builder.danger_accept_invalid_certs(true);
        }
        let tls = tls_builder
            .build()
            .map_err(|e| AppError::Internal(format!("TLS error: {}", e)))?;
        let client = imap::connect((host, port), host, &tls)
            .map_err(|e| AppError::Internal(format!("IMAP connect error: {}", e)))?;
        let mut session = client
            .login(user, password)
            .map_err(|e| AppError::Internal(format!("IMAP login error: {:?}", e.0)))?;

        let bounces = Self::process_session(&mut session, mailbox)?;
        let _ = session.logout();
        Ok(bounces)
    }

    /// Fetch bounces using STARTTLS (plain connection upgraded to TLS, typically port 143)
    fn fetch_bounces_starttls(
        host: &str,
        port: u16,
        accept_invalid_certs: bool,
        ca_cert_path: Option<&str>,
        user: &str,
        password: &str,
        mailbox: &str,
    ) -> Result<Vec<BounceInfo>, AppError> {
        tracing::debug!("Connecting to IMAP server {}:{} with STARTTLS", host, port);

        // Connect plain first with timeout
        let stream = TcpStream::connect((host, port)).map_err(|e| {
            AppError::Internal(format!("TCP connect error to {}:{}: {}", host, port, e))
        })?;

        tracing::debug!("TCP connection established to {}:{}", host, port);

        let timeout_secs = crate::config::tasks::imap_socket_timeout_secs();
        stream
            .set_read_timeout(Some(std::time::Duration::from_secs(timeout_secs)))
            .map_err(|e| AppError::Internal(format!("Failed to set read timeout: {}", e)))?;
        stream
            .set_write_timeout(Some(std::time::Duration::from_secs(timeout_secs)))
            .map_err(|e| AppError::Internal(format!("Failed to set write timeout: {}", e)))?;

        // Build TLS connector for STARTTLS upgrade
        let mut tls_builder = TlsConnector::builder();

        // Load custom CA certificate if configured
        if let Some(path) = ca_cert_path {
            let cert_pem = std::fs::read(path).map_err(|e| {
                AppError::Internal(format!("Failed to read CA certificate {}: {}", path, e))
            })?;
            let cert = Certificate::from_pem(&cert_pem).map_err(|e| {
                AppError::Internal(format!("Failed to parse CA certificate: {}", e))
            })?;
            tls_builder.add_root_certificate(cert);
            tracing::debug!("Added custom CA certificate for IMAP STARTTLS: {}", path);
        }

        if accept_invalid_certs {
            tls_builder.danger_accept_invalid_certs(true);
            tracing::debug!("IMAP STARTTLS: Accepting invalid certificates (INSECURE)");
        }
        let tls = tls_builder
            .build()
            .map_err(|e| AppError::Internal(format!("TLS connector build error: {}", e)))?;

        // Create IMAP client - this reads the initial greeting from the server
        // Note: Client::new() automatically reads the server's initial greeting response
        // If the server response is malformed or non-standard, this will fail with
        // "Unable to parse status response"
        tracing::debug!("Reading initial IMAP greeting from server");
        let client = imap::Client::new(stream);

        tracing::debug!("Sending STARTTLS command to {}:{}", host, port);

        // Upgrade connection to TLS via STARTTLS
        // The secure() method sends STARTTLS and handles the server response
        let tls_client = client.secure(host, &tls).map_err(|e| {
            // Provide more detailed error information with context
            let error_msg = format!(
                "IMAP STARTTLS error on {}:{}: {}. \
                Possible causes: \
                1) Server doesn't support STARTTLS \
                2) Server response was malformed or non-standard \
                3) Network issue or timeout \
                4) Certificate validation failed (check IMAP_ACCEPT_INVALID_CERTS or IMAP_CA_CERT_PATH) \
                \
                Diagnostic steps: \
                1) Test server response manually: \
                   echo 'a1 CAPABILITY' | openssl s_client -connect {}:{} -starttls imap -quiet 2>&1 \
                2) Check if server supports STARTTLS: \
                   telnet {} {} then type: a1 CAPABILITY (look for STARTTLS in response) \
                3) Verify certificate: \
                   openssl s_client -connect {}:{} -starttls imap -showcerts \
                4) If server works with other clients, consider using IMAP_USE_TLS=true with port 993 \
                   for direct TLS connection (bypasses STARTTLS negotiation)",
                host, port, e, host, port, host, port, host, port
            );
            tracing::error!("{}", error_msg);
            AppError::Internal(error_msg)
        })?;

        tracing::debug!("STARTTLS upgrade successful, connection is now encrypted");

        // Login after STARTTLS
        tracing::debug!("Attempting IMAP login for user: {}", user);
        let mut session = tls_client.login(user, password).map_err(|e| {
            tracing::error!("IMAP login failed for user {}: {:?}", user, e.0);
            AppError::Internal(format!("IMAP login error: {:?}", e.0))
        })?;

        tracing::debug!("IMAP login successful, processing mailbox: {}", mailbox);

        let bounces = Self::process_session(&mut session, mailbox)?;
        let _ = session.logout();
        tracing::debug!("IMAP session closed successfully");
        Ok(bounces)
    }

    /// Fetch bounces over plain connection (not recommended)
    fn fetch_bounces_plain(
        host: &str,
        port: u16,
        user: &str,
        password: &str,
        mailbox: &str,
    ) -> Result<Vec<BounceInfo>, AppError> {
        let stream = TcpStream::connect((host, port))
            .map_err(|e| AppError::Internal(format!("TCP connect error: {}", e)))?;
        let client = imap::Client::new(stream);
        let mut session = client
            .login(user, password)
            .map_err(|e| AppError::Internal(format!("IMAP login error: {:?}", e.0)))?;

        let bounces = Self::process_session(&mut session, mailbox)?;
        let _ = session.logout();
        Ok(bounces)
    }

    /// Process an IMAP session to fetch bounce messages
    fn process_session<T: std::io::Read + std::io::Write>(
        session: &mut imap::Session<T>,
        mailbox: &str,
    ) -> Result<Vec<BounceInfo>, AppError> {
        // Select mailbox
        session
            .select(mailbox)
            .map_err(|e| AppError::Internal(format!("IMAP select error: {}", e)))?;

        // RFC 3464: a DSN is multipart/report with report-type=delivery-status.
        // HEADER matches the top-level Content-Type, so both tokens are needed.
        // IMAP OR is a prefix operator, not an infix one.
        let search_result = session
            .search(
                r#"OR (UNSEEN HEADER Content-Type "delivery-status") (UNSEEN HEADER Content-Type "multipart/report")"#,
            )
            .map_err(|e| AppError::Internal(format!("IMAP search error: {e}")))?;

        let mut bounces = Vec::new();

        for msg_seq in search_result.iter() {
            // Fetch the message
            let messages = session
                .fetch(msg_seq.to_string(), "RFC822")
                .map_err(|e| AppError::Internal(format!("IMAP fetch error: {}", e)))?;

            for message in messages.iter() {
                if let Some(body) = message.body() {
                    if let Some(bounce_info) = Self::parse_dsn_message(body) {
                        bounces.push(bounce_info);

                        // Mark message as seen
                        let _ = session.store(msg_seq.to_string(), "+FLAGS (\\Seen)");
                    }
                }
            }
        }

        Ok(bounces)
    }

    /// Parse a DSN message to extract bounce information
    fn parse_dsn_message(body: &[u8]) -> Option<BounceInfo> {
        let body_str = String::from_utf8_lossy(body);

        // Look for common DSN patterns
        // Original-Recipient or Final-Recipient
        let recipient = Self::extract_header_value(&body_str, "Final-Recipient:")
            .or_else(|| Self::extract_header_value(&body_str, "Original-Recipient:"))
            .map(|r| r.replace("rfc822;", "").trim().to_string())?;

        // Action (failed, delayed, delivered, relayed, expanded)
        let action = Self::extract_header_value(&body_str, "Action:")
            .map(|a| a.trim().to_lowercase())
            .unwrap_or_else(|| "failed".to_string());

        // Status code (e.g., 5.1.1, 4.2.2)
        let status = Self::extract_header_value(&body_str, "Status:").map(|s| s.trim().to_string());

        // Diagnostic-Code
        let diagnostic =
            Self::extract_header_value(&body_str, "Diagnostic-Code:").map(|d| d.trim().to_string());

        // Try to extract original message ID
        let original_message_id = Self::extract_header_value(&body_str, "Original-Message-ID:")
            .or_else(|| Self::extract_header_value(&body_str, "Message-ID:"))
            .map(|m| m.trim().trim_matches(['<', '>']).to_string());

        // Map action to our state
        let state = match action.as_str() {
            "delivered" => "delivered",
            "relayed" | "expanded" => "accepted",
            "delayed" => "sending",
            _ => "bounced",
        };

        Some(BounceInfo {
            recipient,
            action,
            status,
            diagnostic,
            original_message_id,
            state: state.to_string(),
        })
    }

    /// Extract a header value from message body
    fn extract_header_value(body: &str, header: &str) -> Option<String> {
        for line in body.lines() {
            if line.to_lowercase().starts_with(&header.to_lowercase()) {
                return Some(line[header.len()..].trim().to_string());
            }
        }
        None
    }

    /// Update EmailDeliveryStatus with bounce information
    async fn update_delivery_status(
        &self,
        db: &DatabaseConnection,
        bounce: &BounceInfo,
    ) -> Result<(), AppError> {
        // Find the delivery status by recipient email
        let status = email_delivery_status::Entity::find()
            .filter(email_delivery_status::Column::Recipient.eq(&bounce.recipient))
            .filter(email_delivery_status::Column::State.ne("delivered"))
            .filter(email_delivery_status::Column::State.ne("bounced"))
            .one(db)
            .await?;

        if let Some(existing) = status {
            let now = Utc::now();
            let mut active: email_delivery_status::ActiveModel = existing.into();

            // Update state based on bounce action
            let new_state = if bounce.state == "bounced" || bounce.state == "delivered" {
                &bounce.state
            } else {
                // Keep current state for delayed/relayed
                return Ok(());
            };

            active.state = Set(new_state.clone());
            active.smtp_code = Set(bounce.status.clone());
            active.smtp_response = Set(bounce.diagnostic.clone());
            active.last_state_at = Set(now);
            active.updated_at = Set(now);

            active.update(db).await?;

            tracing::info!(
                "Updated delivery status for {}: {} (action={}, status={}, msg_id={:?})",
                bounce.recipient,
                new_state,
                bounce.action,
                bounce.status.as_deref().unwrap_or("no status"),
                bounce.original_message_id
            );
        }

        Ok(())
    }
}

/// Information extracted from a DSN bounce message
#[derive(Debug)]
struct BounceInfo {
    /// Recipient email address
    recipient: String,
    /// DSN action (failed, delayed, delivered, etc.)
    action: String,
    /// Status code (e.g., 5.1.1)
    status: Option<String>,
    /// Diagnostic message
    diagnostic: Option<String>,
    /// Original message ID (if available)
    original_message_id: Option<String>,
    /// Mapped state for our system
    state: String,
}