use rocket::Route; use rocket::response::status::BadRequest; use rocket_contrib::{Json, Value}; use data_encoding::BASE32; use db::DbConn; use db::models::*; use util; use crypto; use auth::Headers; #[get("/two-factor")] fn get_twofactor(headers: Headers) -> Result> { let data = if headers.user.totp_secret.is_none() { Value::Null } else { json!([{ "Enabled": true, "Type": 0, "Object": "twoFactorProvider" }]) }; Ok(Json(json!({ "Data": data, "Object": "list" }))) } #[post("/two-factor/get-recover", data = "")] fn get_recover(data: Json, headers: Headers) -> Result> { let password_hash = data["masterPasswordHash"].as_str().unwrap(); if !headers.user.check_valid_password(password_hash) { err!("Invalid password"); } Ok(Json(json!({ "Code": headers.user.totp_recover, "Object": "twoFactorRecover" }))) } #[post("/two-factor/get-authenticator", data = "")] fn generate_authenticator(data: Json, headers: Headers) -> Result> { let password_hash = data["masterPasswordHash"].as_str().unwrap(); if !headers.user.check_valid_password(password_hash) { err!("Invalid password"); } let (enabled, key) = match headers.user.totp_secret { Some(secret) => (true, secret), _ => (false, BASE32.encode(&crypto::get_random(vec![0u8; 20]))) }; Ok(Json(json!({ "Enabled": enabled, "Key": key, "Object": "twoFactorAuthenticator" }))) } #[post("/two-factor/authenticator", data = "")] fn activate_authenticator(data: Json, headers: Headers, conn: DbConn) -> Result> { let password_hash = data["masterPasswordHash"].as_str().unwrap(); if !headers.user.check_valid_password(password_hash) { err!("Invalid password"); } let token = data["token"].as_str(); // 123456 let key = data["key"].as_str().unwrap(); // YI4SKBIXG32LOA6VFKH2NI25VU3E4QML // Validate key as base32 and 20 bytes length let decoded_key: Vec = match BASE32.decode(key.as_bytes()) { Ok(decoded) => decoded, _ => err!("Invalid totp secret") }; if decoded_key.len() != 20 { err!("Invalid key length") } // Set key in user.totp_secret let mut user = headers.user; user.totp_secret = Some(key.to_uppercase()); // Validate the token provided with the key if !user.check_totp_code(util::parse_option_string(token)) { err!("Invalid totp code") } // Generate totp_recover let totp_recover = BASE32.encode(&crypto::get_random(vec![0u8; 20])); user.totp_recover = Some(totp_recover); user.save(&conn); Ok(Json(json!({ "Enabled": true, "Key": key, "Object": "twoFactorAuthenticator" }))) } #[post("/two-factor/disable", data = "")] fn disable_authenticator(data: Json, headers: Headers, conn: DbConn) -> Result> { let _type = &data["type"]; let password_hash = data["masterPasswordHash"].as_str().unwrap(); if !headers.user.check_valid_password(password_hash) { err!("Invalid password"); } let mut user = headers.user; user.totp_secret = None; user.totp_recover = None; user.save(&conn); Ok(Json(json!({ "Enabled": false, "Type": 0, "Object": "twoFactorProvider" }))) }