crabmail/src/smtp_client.rs

use std::sync::Arc;
use rustls_pki_types::ServerName;
use tokio::io::{AsyncWriteExt, AsyncReadExt};
use tokio::net::TcpStream;
use std::convert::TryFrom;
use std::error::Error;
use anyhow::anyhow;
use base64::encode;
use lettre::{Address, Message};
use crate::config::Config;

#[cfg(not(target_os = "wasi"))]
use native_tls::{TlsConnector, TlsStream};
#[cfg(not(target_os = "wasi"))]
use std::io::{Read, Write};
#[cfg(not(target_os = "wasi"))]
use std::net::TcpStream as StdTcpStream;

#[cfg(target_os = "wasi")]
use tokio_rustls::rustls::{ClientConfig, RootCertStore};
#[cfg(target_os = "wasi")]
use tokio_rustls::client::TlsStream;

#[cfg(target_os = "wasi")]
async fn send_command_tcp(stream: &mut TcpStream, command: &str) -> anyhow::Result<()> {
    stream.write_all(command.as_bytes()).await;
    let mut buffer = [0; 1024];
    stream.read(&mut buffer).await;
    println!("Server: {}", String::from_utf8_lossy(&buffer));

    Ok(())
}

#[cfg(not(target_os = "wasi"))]
async fn send_command_tcp(stream: &mut StdTcpStream, command: &str) -> anyhow::Result<()> {
    stream.write_all(command.as_bytes())?;
    let mut buffer = [0; 1024];
    stream.read(&mut buffer)?;
    println!("Server: {}", String::from_utf8_lossy(&buffer));
    
    Ok(())
}

#[cfg(target_os = "wasi")]
async fn send_command_bytes(stream: &mut tokio_rustls::client::TlsStream<TcpStream>, command: Vec<u8>) -> anyhow::Result<()> {
    stream.write_all(&*command).await?;
    println!("SMTP server: {:?}", read_response(stream).await?);
    Ok(())
}

#[cfg(not(target_os = "wasi"))]
async fn send_command_bytes(stream: &mut TlsStream<StdTcpStream>, command: Vec<u8>) -> anyhow::Result<()> {
    stream.write_all(&*command)?;
    println!("SMTP server: {:?}", read_response(stream).await?);
    
    Ok(())
}

#[cfg(target_os = "wasi")]
async fn send_command(stream: &mut tokio_rustls::client::TlsStream<TcpStream>, command: &str) -> anyhow::Result<()> {
    stream.write_all(command.as_bytes()).await?;
    println!("SMTP server: {:?}", read_response(stream).await?);
    Ok(())
}

#[cfg(not(target_os = "wasi"))]
async fn send_command(stream: &mut TlsStream<StdTcpStream>, command: &str) -> anyhow::Result<()> {
    stream.write_all(command.as_bytes())?;
    println!("SMTP server: {:?}", read_response(stream).await?);
    
    Ok(())
}

#[cfg(target_os = "wasi")]
async fn read_response(stream: &mut tokio_rustls::client::TlsStream<TcpStream>) -> anyhow::Result<String>{
    let mut buffer = [0; 1024];
    match stream.read(&mut buffer).await{
        Ok(_) => {
            Ok(String::from_utf8_lossy(&buffer).to_string())
        }
        Err(_) => {
            Ok("".to_string())
        }
    }
}

#[cfg(not(target_os = "wasi"))]
async fn read_response(stream: &mut TlsStream<StdTcpStream>) -> anyhow::Result<String>{
    let mut buffer = [0; 1024];
    match stream.read(&mut buffer){
        Ok(_) => {
            Ok(String::from_utf8_lossy(&buffer).to_string())
        }
        Err(_) => {
            Ok("".to_string())
        }
    }
}

#[cfg(target_os = "wasi")]
async fn login() -> Result<tokio_rustls::client::TlsStream<TcpStream>, Box<dyn std::error::Error>> {
    let mut stream = tokio::net::TcpStream::connect(
        &(Config::global().smtp_domain.clone(), Config::global().smtp_port.clone() as u16)
    ).await?;
    let mut buffer = [0; 1024];
    stream.read(&mut buffer).await?;
    println!("Server: {:?}", String::from_utf8(Vec::from(buffer)));

    // Send EHLO command
    send_command_tcp(&mut stream, "EHLO localhost\r\n").await;

    // Start TLS
    send_command_tcp(&mut stream, "STARTTLS\r\n").await;

    // Wrap the stream with TLS
    let mut root_cert_store = RootCertStore::empty();
    root_cert_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    let config = ClientConfig::builder()
        .with_root_certificates(root_cert_store)
        .with_no_client_auth();
    let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
    let dnsname = ServerName::try_from(Config::global().smtp_domain.clone()).unwrap();
    let mut stream = connector.connect(dnsname, stream).await?;


    // Send EHLO command again after establishing TLS
    send_command(&mut stream, "EHLO localhost\r\n").await;

    // Authenticate using LOGIN method
    send_command(&mut stream, "AUTH LOGIN\r\n").await;
    send_command(&mut stream, &(encode(Config::global().username.lock().unwrap().clone()) + "\r\n")).await;
    send_command(&mut stream, &(encode(Config::global().password.lock().unwrap().clone()) + "\r\n")).await;
    
    Ok(stream)
}

#[cfg(not(target_os = "wasi"))]
async fn login() -> Result<TlsStream<StdTcpStream>, Box<dyn std::error::Error>> {
    let connector = TlsConnector::new().unwrap();

    let mut stream = StdTcpStream::connect(format!("{}:{}", Config::global().smtp_domain.clone(), Config::global().smtp_port.clone() as u16)).unwrap();
    let mut buffer = [0; 1024];
    stream.read(&mut buffer).unwrap();
    println!("Server: {:?}", String::from_utf8(Vec::from(buffer)));

    // Send EHLO command
    send_command_tcp(&mut stream, "EHLO localhost\r\n").await?;

    // Start TLS
    send_command_tcp(&mut stream, "STARTTLS\r\n").await?;

    // Wrap the stream with TLS
    let mut stream = connector.connect(&*Config::global().smtp_domain.clone(), stream).unwrap();

    // Send EHLO command again after establishing TLS
    send_command(&mut stream, "EHLO localhost\r\n").await?;

    // Authenticate using LOGIN method
    send_command(&mut stream, "AUTH LOGIN\r\n").await?;
    send_command(&mut stream, &(encode(Config::global().username.lock().unwrap().clone()) + "\r\n")).await?;
    send_command(&mut stream, &(encode(Config::global().password.lock().unwrap().clone()) + "\r\n")).await?; // TODO use Engine::encode()

    Ok(stream)
}

#[cfg(target_os = "wasi")]
async fn logout(mut stream: &mut tokio_rustls::client::TlsStream<TcpStream>) -> anyhow::Result<()>{
    // Send the QUIT command to end the session
    send_command(&mut stream, "QUIT\r\n").await; // TODO maybe '?'

    // Close the connection
    stream.shutdown().await?;
    
    Ok(())
}

#[cfg(not(target_os = "wasi"))]
async fn logout(mut stream: &mut TlsStream<StdTcpStream>) -> anyhow::Result<()>{
    // Send the QUIT command to end the session
    send_command(&mut stream, "QUIT\r\n").await?;

    // Close the connection
    stream.shutdown()?;
    
    Ok(())
}

pub async fn send_email(email: Message) -> anyhow::Result<()>{
    let mut stream= match login().await{
        Ok(stream) => {stream}
        Err(_) => {return Err(anyhow!("Unable to establish connection with SMTP server"))}
    };
    
    // Specify the sender
    let from = match email.envelope().from() {
        None => {return Err(anyhow!("Wrong sender address"))}
        Some(from) => {from}
    };
    send_command(&mut stream, format!("MAIL FROM:<{}@{}>\r\n", 
                                      from.user(),
                                      from.domain())
        .to_owned().as_str()).await?;
    
    // Specify the recipient
    let recipients = email.envelope().to();
    for recipient in recipients {
        send_command(&mut stream, format!("RCPT TO:<{}@{}>\r\n", 
                                          recipient.user(), 
                                          recipient.domain()
        ).to_owned().as_str()).await?;
    }
    
    // Send the DATA command to start the message content
    send_command(&mut stream, "DATA\r\n").await?;
    
    // Send the message content
    let mut email = email.formatted();
    email.append(&mut Vec::from("\r\n.\r\n".as_bytes()));
    send_command_bytes(&mut stream, email).await?;
    
    logout(&mut stream).await?;
    
    Ok(())
}