Generate recovery codes for email and duo

2 files changed, 10 insertions, 3 deletions

diff --git a/src/api/core/two_factor/duo.rs b/src/api/core/two_factor/duo.rs
index 5b7d7886..1d8074d5 100644
--- a/src/api/core/two_factor/duo.rs
+++ b/src/api/core/two_factor/duo.rs
@@ -4,6 +4,7 @@ use rocket::Route;
 use rocket_contrib::json::Json;
 use serde_json;
 
+use crate::api::core::two_factor::_generate_recover_code;
 use crate::api::{ApiResult, EmptyResult, JsonResult, JsonUpcase, PasswordData};
 use crate::auth::Headers;
 use crate::crypto;
@@ -152,8 +153,9 @@ fn check_duo_fields_custom(data: &EnableDuoData) -> bool {
 #[post("/two-factor/duo", data = "<data>")]
 fn activate_duo(data: JsonUpcase<EnableDuoData>, headers: Headers, conn: DbConn) -> JsonResult {
     let data: EnableDuoData = data.into_inner().data;
+    let mut user = headers.user;
 
-    if !headers.user.check_valid_password(&data.MasterPasswordHash) {
+    if !user.check_valid_password(&data.MasterPasswordHash) {
         err!("Invalid password");
     }
 
@@ -167,8 +169,10 @@ fn activate_duo(data: JsonUpcase<EnableDuoData>, headers: Headers, conn: DbConn)
     };
 
     let type_ = TwoFactorType::Duo;
-    let twofactor = TwoFactor::new(headers.user.uuid, type_, data_str);
+    let twofactor = TwoFactor::new(user.uuid.clone(), type_, data_str);
     twofactor.save(&conn)?;
+    
+    _generate_recover_code(&mut user, &conn);
 
     Ok(Json(json!({
         "Enabled": true,
diff --git a/src/api/core/two_factor/email.rs b/src/api/core/two_factor/email.rs
index 97e55dfc..654d239f 100644
--- a/src/api/core/two_factor/email.rs
+++ b/src/api/core/two_factor/email.rs
@@ -2,6 +2,7 @@ use rocket::Route;
 use rocket_contrib::json::Json;
 use serde_json;
 
+use crate::api::core::two_factor::_generate_recover_code;
 use crate::api::{EmptyResult, JsonResult, JsonUpcase, PasswordData};
 use crate::auth::Headers;
 use crate::crypto;
@@ -172,7 +173,7 @@ struct EmailData {
 #[put("/two-factor/email", data = "<data>")]
 fn email(data: JsonUpcase<EmailData>, headers: Headers, conn: DbConn) -> JsonResult {
     let data: EmailData = data.into_inner().data;
-    let user = headers.user;
+    let mut user = headers.user;
 
     if !user.check_valid_password(&data.MasterPasswordHash) {
         err!("Invalid password");
@@ -197,6 +198,8 @@ fn email(data: JsonUpcase<EmailData>, headers: Headers, conn: DbConn) -> JsonRes
     twofactor.data = email_data.to_json();
     twofactor.save(&conn)?;
 
+    _generate_recover_code(&mut user, &conn);
+
     Ok(Json(json!({
         "Email": email_data.email,
         "Enabled": "true",