fileserver/backend/src/routes/auth/tfa.rs

use lazy_static::lazy_static;
use lettre::Transport;
use ring::rand::SecureRandom;
use warp::Filter;
use crate::config::CONFIG;
use crate::db::{DBConnection, DBPool, with_db, TfaTypes};
use crate::dto;
use crate::routes::{AppError, get_reply};
use crate::routes::filters::{authenticated, UserInfo};

fn build_mail_sender() -> lettre::SmtpTransport {
    lettre::SmtpTransport::builder_dangerous(CONFIG.smtp_server.clone())
        .port(CONFIG.smtp_port)
        .tls(
            lettre::transport::smtp::client::Tls::Required(
                lettre::transport::smtp::client::TlsParameters::new(
                    CONFIG.smtp_server.clone()
                ).unwrap()
            )
        )
        .credentials(lettre::transport::smtp::authentication::Credentials::new(CONFIG.smtp_user.clone(), CONFIG.smtp_password.clone()))
        .build()
}

lazy_static! {
    static ref MAIL_SENDER: lettre::SmtpTransport = build_mail_sender();
}

fn get_totp(user: &crate::db::User) -> totp_rs::TOTP {
    totp_rs::TOTP::from_rfc6238(
        totp_rs::Rfc6238::new(
            6,
            user.tfa_secret.clone().unwrap(),
            Some("MFileserver".to_owned()),
            user.name.clone()
        ).unwrap()
    ).unwrap()
}

pub fn verify2fa(user: &crate::db::User, code: String) -> bool {
    let allowed_skew = if user.tfa_type == TfaTypes::Totp {0} else {10};
    let totp = get_totp(user);
    let time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
    let base_step = time / totp.step - allowed_skew;
    for i in 0..allowed_skew + 1 {
        let step = (base_step + i) * totp.step;
        if totp.generate(step).eq(&code) {
            return true;
        }
    }
    false
}

pub fn send_2fa_mail(user: &crate::db::User) {
    let totp = get_totp(user);
    let code = totp.generate_current().unwrap();
    let mail = lettre::Message::builder()
        .from("fileserver@mattv.de".parse().unwrap())
        .to(user.name.parse().unwrap())
        .subject("MFileserver - Email 2fa code")
        .body(format!("Your code is: {}\r\nIt is valid for 5 minutes", code))
        .unwrap();

    MAIL_SENDER.send(&mail).expect("Failed to send mail");
}

pub fn build_routes(db: DBPool) -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
    let tfa_setup = warp::path!("auth" / "2fa" / "setup")
        .and(warp::post())
        .and(warp::body::json())
        .and(authenticated(db.clone()))
        .and(with_db(db.clone()))
        .and_then(tfa_setup);
    let tfa_complete = warp::path!("auth" / "2fa" / "complete")
        .and(warp::post())
        .and(warp::body::json())
        .and(authenticated(db.clone()))
        .and(with_db(db.clone()))
        .and_then(tfa_complete);
    let tfa_disable = warp::path!("auth" / "2fa" / "disable")
        .and(warp::post())
        .and(authenticated(db.clone()))
        .and(with_db(db))
        .and_then(tfa_disable);

    tfa_setup.or(tfa_complete).or(tfa_disable)
}

async fn tfa_setup(data: dto::requests::TfaSetup, mut info: UserInfo, mut db: DBConnection)
    -> Result<impl warp::Reply, warp::Rejection> {
    let mut secret: [u8; 32] = [0; 32];
    super::SEC_RANDOM.fill(&mut secret).expect("Failed to generate secret");
    let secret = Vec::from(secret);
    info.0.tfa_secret = Some(secret);
    db.save_user(&info.0);

    if data.mail {
        send_2fa_mail(&info.0);
        get_reply(&dto::responses::Success {
            statusCode: 200
        })
    } else {
        let totp = get_totp(&info.0);
        get_reply(&dto::responses::TfaSetup {
            statusCode: 200,
            secret: totp.get_secret_base32(),
            qrCode: "data:image/png;base64,".to_owned() + &totp.get_qr().expect("Failed to generate qr code")
        })
    }
}

async fn tfa_complete(data: dto::requests::TfaComplete, mut info: UserInfo, mut db: DBConnection)
    -> Result<impl warp::Reply, warp::Rejection> {
    info.0.tfa_type = if data.mail { TfaTypes::Email } else { TfaTypes::Totp };

    if verify2fa(&info.0, data.code) {
        db.save_user(&info.0);
        db.delete_all_tokens(info.0.id);
        get_reply(&dto::responses::Success {
            statusCode: 200
        })
    } else {
        AppError::BadRequest("Incorrect 2fa code").err()
    }
}

async fn tfa_disable(mut info: UserInfo, mut db: DBConnection)
    -> Result<impl warp::Reply, warp::Rejection> {
    info.0.tfa_secret = None;
    info.0.tfa_type = TfaTypes::None;
    db.save_user(&info.0);
    db.delete_all_tokens(info.0.id);
    get_reply(&dto::responses::Success {
        statusCode: 200
    })
}