src/controllers/mod.rs

@@ -6,3 +6,4 @@ pub mod admin_controller;
 pub mod admin_api;
 pub mod openid_controller;
 pub mod two_factors_controller;
+pub mod two_factors_api;

src/controllers/two_factors_api.rs

@@ -0,0 +1,41 @@
+use actix::Addr;
+use actix_web::{HttpResponse, Responder, web};
+
+use crate::actors::users_actor;
+use crate::actors::users_actor::UsersActor;
+use crate::data::current_user::CurrentUser;
+use crate::data::totp_key::TotpKey;
+use crate::data::user::{SecondFactor, User};
+
+#[derive(serde::Deserialize)]
+pub struct Request {
+    factor_name: String,
+    secret: String,
+    first_code: String,
+}
+
+pub async fn save_totp_key(user: CurrentUser, form: web::Json<Request>,
+                           users: web::Data<Addr<UsersActor>>) -> impl Responder {
+    let key = TotpKey::from_encoded_secret(&form.secret);
+
+    if !key.check_code(&form.first_code).unwrap_or(false) {
+        return HttpResponse::BadRequest()
+            .body(format!("Given code is invalid (expected {} or {})!",
+                          key.current_code().unwrap_or_default(),
+                          key.previous_code().unwrap_or_default()));
+    }
+
+    if form.factor_name.is_empty() {
+        return HttpResponse::BadRequest().body("Please give a name to the factor!");
+    }
+
+    let mut user = User::from(user);
+    user.add_factor(SecondFactor::TOTP(key));
+    let res = users.send(users_actor::UpdateUserRequest(user)).await.unwrap().0;
+
+    if !res {
+        HttpResponse::InternalServerError().body("Failed to update user information!")
+    } else {
+        HttpResponse::Ok().body("Added new factor!")
+    }
+}

src/data/totp_key.rs

@@ -1,14 +1,19 @@
+use std::io::ErrorKind;
+
 use base32::Alphabet;
 use rand::Rng;
+use totp_rfc6238::{HashAlgorithm, TotpGenerator};
 
 use crate::data::app_config::AppConfig;
 use crate::data::user::User;
+use crate::utils::err::Res;
+use crate::utils::time::time;
 
 const BASE32_ALPHABET: Alphabet = Alphabet::RFC4648 { padding: true };
-const NUM_DIGITS: i32 = 6;
-const PERIOD: i32 = 30;
+const NUM_DIGITS: usize = 6;
+const PERIOD: u64 = 30;
 
-#[derive(serde::Serialize, serde::Deserialize, Debug)]
+#[derive(Clone, serde::Serialize, serde::Deserialize, Debug)]
 pub struct TotpKey {
     encoded: String,
 }
@@ -17,11 +22,16 @@ impl TotpKey {
     /// Generate a new TOTP key
     pub fn new_random() -> Self {
         let random_bytes = rand::thread_rng().gen::<[u8; 10]>();
-        TotpKey {
+        Self {
             encoded: base32::encode(BASE32_ALPHABET, &random_bytes)
         }
     }
 
+    /// Get a key from an encoded secret
+    pub fn from_encoded_secret(s: &str) -> Self {
+        Self { encoded: s.to_string() }
+    }
+
     /// Get QrCode URL for user
     ///
     /// Based on https://github.com/google/google-authenticator/wiki/Key-Uri-Format
@@ -50,4 +60,38 @@ impl TotpKey {
     pub fn get_secret(&self) -> String {
         self.encoded.to_string()
     }
+
+    /// Get current code
+    pub fn current_code(&self) -> Res<String> {
+        self.get_code_at(|| time())
+    }
+
+    /// Get previous code
+    pub fn previous_code(&self) -> Res<String> {
+        self.get_code_at(|| time() - PERIOD)
+    }
+
+    /// Get the code at a specific time
+    fn get_code_at<F: Fn() -> u64>(&self, get_time: F) -> Res<String> {
+        let gen = TotpGenerator::new()
+            .set_digit(NUM_DIGITS).unwrap()
+            .set_step(PERIOD).unwrap()
+            .set_hash_algorithm(HashAlgorithm::SHA1)
+            .build();
+
+        let key = match base32::decode(BASE32_ALPHABET, &self.encoded) {
+            None => {
+                return Err(Box::new(
+                    std::io::Error::new(ErrorKind::Other, "Failed to decode base32 secret!")));
+            }
+            Some(k) => k,
+        };
+
+        Ok(gen.get_code_with(&key, get_time))
+    }
+
+    /// Check a code's validity
+    pub fn check_code(&self, code: &str) -> Res<bool> {
+        Ok(self.current_code()?.eq(code) || self.previous_code()?.eq(code))
+    }
 }

src/data/user.rs

@@ -1,9 +1,15 @@
 use crate::data::client::ClientID;
 use crate::data::entity_manager::EntityManager;
+use crate::data::totp_key::TotpKey;
 use crate::utils::err::Res;
 
 pub type UserID = String;
 
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub enum SecondFactor {
+    TOTP(TotpKey)
+}
+
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 pub struct User {
     pub uid: UserID,
@@ -16,6 +22,9 @@ pub struct User {
     pub enabled: bool,
     pub admin: bool,
 
+    /// 2FA
+    pub second_factors: Option<Vec<SecondFactor>>,
+
     /// None = all services
     /// Some([]) = no service
     pub authorized_clients: Option<Vec<ClientID>>,
@@ -36,6 +45,14 @@ impl User {
     pub fn verify_password<P: AsRef<[u8]>>(&self, pass: P) -> bool {
         verify_password(pass, &self.password)
     }
+
+    pub fn add_factor(&mut self, factor: SecondFactor) {
+        if self.second_factors.is_none() {
+            self.second_factors = Some(vec![]);
+        }
+
+        self.second_factors.as_mut().unwrap().push(factor);
+    }
 }
 
 impl PartialEq for User {
@@ -58,6 +75,7 @@ impl Default for User {
             need_reset_password: false,
             enabled: true,
             admin: false,
+            second_factors: Some(vec![]),
             authorized_clients: Some(Vec::new()),
         }
     }

src/main.rs

@@ -119,7 +119,10 @@ async fn main() -> std::io::Result<()> {
             .route("/settings/change_password", web::get().to(settings_controller::change_password_route))
             .route("/settings/change_password", web::post().to(settings_controller::change_password_route))
             .route("/settings/two_factors", web::get().to(two_factors_controller::two_factors_route))
-            .route("settings/two_factors/add_totp", web::get().to(two_factors_controller::add_totp_factor_route))
+            .route("/settings/two_factors/add_totp", web::get().to(two_factors_controller::add_totp_factor_route))
+
+            // User API
+            .route("/settings/api/two_factors/save_totp_key", web::post().to(two_factors_api::save_totp_key))
 
             // Admin routes
             .route("/admin", web::get()

templates/settings/add_2fa_totp_page.html

@@ -21,23 +21,83 @@
 
     <p>Once you have scanned the QrCode, please generate a code and type it below:</p>
 
+    <form id="validateForm" method="post">
+        <input type="hidden" name="secret" id="secretInput" value="{{ secret_key }}"/>
+
         <div class="form-group">
             <label for="inputDevName" class="form-label mt-4">Device name</label>
-        <input type="text" class="form-control" id="inputDevName" aria-describedby="emailHelp" placeholder="Enter email"
-               value="Authenticator app">
-        <small class="form-text text-muted">Please give a name to your device to identity it more easily later.</small>
+            <input type="text" class="form-control" id="inputDevName"
+                   placeholder="Device / Authenticator app name"
+                   value="Authenticator app" minlength="1" required/>
+            <small class="form-text text-muted">Please give a name to your device to identity it more easily
+                later.</small>
+            <div class="invalid-feedback">Please give a name to this authenticator app</div>
         </div>
 
         <div class="form-group">
             <label for="inputFirstCode" class="form-label mt-4">First code</label>
-        <input type="text" class="form-control" id="inputFirstCode" aria-describedby="emailHelp" placeholder="XXXXXX"
-               maxlength="6"/>
+            <input type="text" class="form-control" id="inputFirstCode"
+                   placeholder="XXXXXX"
+                   maxlength="6" minlength="6" required/>
             <small class="form-text text-muted">Check that your authenticator app is working correctly by typing a first
                 code.</small>
+            <div class="invalid-feedback">Please enter a first code (must have 6 digits)</div>
         </div>
 
+        <input type="submit" value="Register app" class="btn btn-primary">
+    </form>
+
     <script src="/assets/js/clipboard_utils.js"></script>
 
+    <script>
+        const form = document.getElementById("validateForm");
+        form.addEventListener("submit", async (e) => {
+            e.preventDefault();
+            const secret = document.getElementById("secretInput").value;
+            const factorNameInput = document.getElementById("inputDevName");
+            const firstCodeInput = document.getElementById("inputFirstCode");
+
+            let fail = false;
+            factorNameInput.classList.remove("is-invalid");
+            if (factorNameInput.value.length === 0) {
+                fail = true;
+                factorNameInput.classList.add("is-invalid");
+            }
+
+            firstCodeInput.classList.remove("is-invalid");
+            if (firstCodeInput.value.length != 6) {
+                fail = true;
+                firstCodeInput.classList.add("is-invalid");
+            }
+
+
+            if (fail)
+                return;
+
+            try {
+                const res = await fetch("/settings/api/two_factors/save_totp_key", {
+                    method: "post",
+                    headers: {
+                      'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({
+                        factor_name: factorNameInput.value,
+                        secret: secret,
+                        first_code: firstCodeInput.value,
+                    })
+                });
+
+                let text = await res.text();
+                alert(text);
+
+                if (res.status == 200)
+                    location.href = "/settings/two_factors";
+            } catch(e) {
+                console.error(e);
+                alert("Failed to register authenticator app!");
+            }
+        });
+    </script>
 </div>
 
 {% endblock content %}