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

use lettre::Transport;
use once_cell::sync::Lazy;
use ring::rand::SecureRandom;
use rustracing_jaeger::Span;
use tiny_http::{Request, ResponseBox};

use crate::{
    config::CONFIG,
    db,
    db::{DBConnection, TfaTypes},
    dto,
    metrics,
    routes::{filters::UserInfo, get_reply, AppError}
};

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()
}

pub static MAIL_SENDER: Lazy<lettre::SmtpTransport> = Lazy::new(build_mail_sender);

fn get_totp(user: &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: &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(span: &Span, user: &db::User) {
    let _span = metrics::span("send_2fa_mail", span);
    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 tfa_setup(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    mut info: UserInfo,
    data: dto::requests::TfaSetup
) -> Result<ResponseBox, AppError> {
    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(span, &info.0);

    if data.mail {
        send_2fa_mail(span, &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")
        })
    }
}

pub fn tfa_complete(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    mut info: UserInfo,
    data: dto::requests::TfaComplete
) -> Result<ResponseBox, AppError> {
    info.0.tfa_type = if data.mail {
        TfaTypes::Email
    } else {
        TfaTypes::Totp
    };

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

pub fn tfa_disable(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    mut info: UserInfo
) -> Result<ResponseBox, AppError> {
    info.0.tfa_secret = None;
    info.0.tfa_type = TfaTypes::None;
    db.save_user(span, &info.0);
    db.delete_all_tokens(span, info.0.id);
    get_reply(&dto::responses::Success { statusCode: 200 })
}