From ce7118ff811bd905b8afd2a029f72fa0da8476fa Mon Sep 17 00:00:00 2001
From: Pierre Hubert <pierre.git@communiquons.org>
Date: Tue, 19 Apr 2022 19:24:07 +0200
Subject: [PATCH] Display form to enter OTP code

---
 src/controllers/login_controller.rs | 45 ++++++++++++++++++++++++++++-
 src/data/user.rs                    |  6 +++-
 src/main.rs                         |  2 ++
 templates/login/opt_input.html      | 24 +++++++++++++++
 4 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 templates/login/opt_input.html

diff --git a/src/controllers/login_controller.rs b/src/controllers/login_controller.rs
index cb0cb4e..90cda06 100644
--- a/src/controllers/login_controller.rs
+++ b/src/controllers/login_controller.rs
@@ -11,7 +11,7 @@ use crate::controllers::base_controller::{FatalErrorPage, redirect_user, redirec
 use crate::data::login_redirect_query::LoginRedirectQuery;
 use crate::data::remote_ip::RemoteIP;
 use crate::data::session_identity::{SessionIdentity, SessionStatus};
-use crate::data::user::{TwoFactor, User};
+use crate::data::user::{FactorID, TwoFactor, User};
 
 struct BaseLoginPage {
     danger: Option<String>,
@@ -42,6 +42,13 @@ struct ChooseSecondFactorTemplate<'a> {
     factors: &'a [TwoFactor],
 }
 
+#[derive(Template)]
+#[template(path = "login/opt_input.html")]
+struct LoginWithOTPTemplate<'a> {
+    _p: BaseLoginPage,
+    factor: &'a TwoFactor,
+}
+
 
 #[derive(serde::Deserialize)]
 pub struct LoginRequestBody {
@@ -253,4 +260,40 @@ pub async fn choose_2fa_method(id: Identity, query: web::Query<ChooseSecondFacto
             .render()
             .unwrap(),
     )
+}
+
+#[derive(serde::Deserialize)]
+pub struct LoginWithOTPQuery {
+    #[serde(default)]
+    redirect: LoginRedirectQuery,
+    id: FactorID,
+}
+
+
+/// Login with OTP
+pub async fn login_with_otp(id: Identity, query: web::Query<LoginWithOTPQuery>,
+                            users: web::Data<Addr<UsersActor>>) -> impl Responder {
+    if !SessionIdentity(&id).need_2fa_auth() {
+        return redirect_user_for_login(query.redirect.get());
+    }
+
+    let user: User = users.send(users_actor::GetUserRequest(SessionIdentity(&id).user_id()))
+        .await.unwrap().0.expect("Could not find user!");
+
+    let factor = match user.find_factor(&query.id) {
+        Some(f) => f,
+        None => return HttpResponse::Ok()
+            .body(FatalErrorPage { message: "Factor not found!" }.render().unwrap())
+    };
+
+    HttpResponse::Ok().body(LoginWithOTPTemplate {
+        _p: BaseLoginPage {
+            danger: None,
+            success: None,
+            page_title: "Two-Factor Auth",
+            app_name: APP_NAME,
+            redirect_uri: query.redirect.get_encoded(),
+        },
+        factor,
+    }.render().unwrap())
 }
\ No newline at end of file
diff --git a/src/data/user.rs b/src/data/user.rs
index 35269ed..a1b88e8 100644
--- a/src/data/user.rs
+++ b/src/data/user.rs
@@ -29,7 +29,7 @@ impl TwoFactor {
 
     pub fn login_url(&self, redirect_uri: &str) -> String {
         match self.kind {
-            TwoFactorType::TOTP(_) => format!("/2fa_totp?id={}&redirect_uri={}",
+            TwoFactorType::TOTP(_) => format!("/2fa_otp?id={}&redirect_uri={}",
                                               self.id.0, redirect_uri)
         }
     }
@@ -83,6 +83,10 @@ impl User {
     pub fn remove_factor(&mut self, factor_id: FactorID) {
         self.two_factor.retain(|f| f.id != factor_id);
     }
+
+    pub fn find_factor(&self, factor_id: &FactorID) -> Option<&TwoFactor> {
+        self.two_factor.iter().find(|f| f.id.eq(&factor_id))
+    }
 }
 
 impl PartialEq for User {
diff --git a/src/main.rs b/src/main.rs
index 959a961..84ab1af 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -112,6 +112,8 @@ async fn main() -> std::io::Result<()> {
             .route("/reset_password", web::get().to(login_controller::reset_password_route))
             .route("/reset_password", web::post().to(login_controller::reset_password_route))
             .route("/2fa_auth", web::get().to(login_controller::choose_2fa_method))
+            .route("/2fa_otp", web::get().to(login_controller::login_with_otp))
+            .route("/2fa_otp", web::post().to(login_controller::login_with_otp))
 
             // Logout page
             .route("/logout", web::get().to(login_controller::logout_route))
diff --git a/templates/login/opt_input.html b/templates/login/opt_input.html
new file mode 100644
index 0000000..417e7a4
--- /dev/null
+++ b/templates/login/opt_input.html
@@ -0,0 +1,24 @@
+{% extends "base_login_page.html" %}
+{% block content %}
+
+<div>
+    <p>Please go to your authenticator app <i>{{ factor.name }}</i>, generate a new code and enter it here:</p>
+    <form method="post" action="{{ factor.login_url(_p.redirect_uri) }}">
+
+        <div class="form-group">
+            <label for="code" class="form-label mt-4">Generated code</label>
+            <input type="text" name="code" minlength="6" maxlength="6" class="form-control" id="code"
+                   placeholder="XXXXXX">
+        </div>
+
+        <button class="w-100 btn btn-lg btn-primary" type="submit" style="margin-top: 20px;">Login</button>
+    </form>
+</div>
+
+
+<div style="margin-top: 10px;">
+    <a href="/2fa_auth?redirect=_p.redirect_uri">Sign in using another factor</a><br />
+    <a href="/logout">Sign out</a>
+</div>
+
+{% endblock content %}
\ No newline at end of file