From fbbcb4e1828495c0ca193da66101a869c49de246 Mon Sep 17 00:00:00 2001
From: Philippe Loctaux <p@philippeloctaux.com>
Date: Mon, 1 May 2023 11:58:29 +0200
Subject: [PATCH] totp: added page to verify totp token

---
 crates/ezidam/src/page.rs                     |   5 +
 crates/ezidam/src/routes/oauth.rs             |  12 ++
 crates/ezidam/src/routes/oauth/totp.rs        | 143 ++++++++++++++++++
 .../templates/pages/oauth/totp.html.tera      |  74 +++++++++
 4 files changed, 234 insertions(+)
 create mode 100644 crates/ezidam/src/routes/oauth/totp.rs
 create mode 100644 crates/ezidam/templates/pages/oauth/totp.html.tera

diff --git a/crates/ezidam/src/page.rs b/crates/ezidam/src/page.rs
index 0551d24..1fec61b 100644
--- a/crates/ezidam/src/page.rs
+++ b/crates/ezidam/src/page.rs
@@ -28,6 +28,7 @@ pub enum Page {
     ForgotPassword,
     ResetPassword(ResetPassword),
     UserSecurityTotp(UserSecurityTotp),
+    AuthorizeTotp(AuthorizeTotp),
 }
 
 impl Page {
@@ -52,6 +53,7 @@ impl Page {
             Page::ForgotPassword => "pages/forgot-password",
             Page::ResetPassword(_) => "pages/reset-password",
             Page::UserSecurityTotp(_) => "pages/settings/totp",
+            Page::AuthorizeTotp(_) => "pages/oauth/totp",
         }
     }
 
@@ -76,6 +78,7 @@ impl Page {
             Page::ForgotPassword => "Forgot password",
             Page::ResetPassword(_) => "Reset password",
             Page::UserSecurityTotp(_) => "Enable One-time password",
+            Page::AuthorizeTotp(_) => "Verifying your account",
         }
     }
 
@@ -102,6 +105,7 @@ impl Page {
             Page::ForgotPassword => None,
             Page::ResetPassword(_) => None,
             Page::UserSecurityTotp(_) => Some(UserMenu::Settings.into()),
+            Page::AuthorizeTotp(_) => None,
         }
     }
 
@@ -126,6 +130,7 @@ impl Page {
             Page::ForgotPassword => Box::new(()),
             Page::ResetPassword(reset) => Box::new(reset),
             Page::UserSecurityTotp(totp) => Box::new(totp),
+            Page::AuthorizeTotp(totp) => Box::new(totp),
         }
     }
 }
diff --git a/crates/ezidam/src/routes/oauth.rs b/crates/ezidam/src/routes/oauth.rs
index b6a2622..2553d15 100644
--- a/crates/ezidam/src/routes/oauth.rs
+++ b/crates/ezidam/src/routes/oauth.rs
@@ -1,3 +1,4 @@
+use self::totp::*;
 use authorize::*;
 use redirect::*;
 use rocket::{routes, Route};
@@ -7,6 +8,7 @@ use userinfo::*;
 pub mod authorize;
 pub mod redirect;
 pub mod token;
+pub mod totp;
 pub mod userinfo;
 
 pub fn routes() -> Vec<Route> {
@@ -17,6 +19,8 @@ pub fn routes() -> Vec<Route> {
         redirect_page,
         request_token,
         get_userinfo,
+        totp_page,
+        totp_verify,
     ]
 }
 
@@ -42,4 +46,12 @@ pub mod content {
         pub username: String,
         pub home_page: String,
     }
+
+    #[derive(Serialize)]
+    #[serde(crate = "rocket::serde")]
+    #[derive(Clone)]
+    pub struct AuthorizeTotp {
+        pub name: Option<String>,
+        pub username: String,
+    }
 }
diff --git a/crates/ezidam/src/routes/oauth/totp.rs b/crates/ezidam/src/routes/oauth/totp.rs
new file mode 100644
index 0000000..26655b9
--- /dev/null
+++ b/crates/ezidam/src/routes/oauth/totp.rs
@@ -0,0 +1,143 @@
+use crate::routes::oauth::{redirect_uri, AUTHORIZATION_CODE_LEN};
+use crate::routes::prelude::*;
+use apps::App;
+use authorization_codes::AuthorizationCode;
+use hash::SecretString;
+use rocket::http::{Cookie, CookieJar};
+use rocket::{get, post};
+use users::totp_login_request::TOTP_REQUEST_COOKIE_NAME;
+use users::User;
+
+#[get("/oauth/totp?<auth_request..>")]
+pub async fn totp_page(
+    totp_request: TotpRequest,
+    mut db: Connection<Database>,
+    flash: Option<FlashMessage<'_>>,
+    auth_request: AuthenticationRequest<'_>,
+) -> Result<Template> {
+    let mut transaction = db.begin().await?;
+
+    // Get app info
+    let _app = App::get_valid_app(
+        &mut transaction,
+        auth_request.response_type,
+        auth_request.scope,
+        auth_request.client_id,
+        auth_request.redirect_uri,
+    )
+    .await?;
+
+    // Get totp request
+    let totp_request =
+        users::totp_login_request::TotpLoginRequest::get_one(&mut transaction, &totp_request.0)
+            .await?
+            .ok_or_else(|| Error::not_found("Failed to find totp request"))?;
+
+    if totp_request.has_expired() {
+        return Err(Error::bad_request("Totp request has expired"));
+    }
+
+    if totp_request.used_at().is_some() {
+        return Err(Error::bad_request("Totp request has been used"));
+    }
+
+    // Get user
+    let user = User::get_by_id(&mut transaction, totp_request.user())
+        .await?
+        .ok_or_else(|| Error::not_found("Failed to find user"))?;
+
+    transaction.commit().await?;
+
+    let page = Page::AuthorizeTotp(super::content::AuthorizeTotp {
+        name: user.name().map(|name| name.to_string()),
+        username: user.username().into(),
+    });
+
+    Ok(flash
+        .map(|flash| Page::with_flash(page.clone(), flash))
+        .unwrap_or_else(|| page.into()))
+}
+
+#[derive(Debug, FromForm)]
+pub struct TotpVerifyForm<'r> {
+    pub code: &'r str,
+}
+
+#[post("/oauth/totp?<auth_request..>", data = "<form>")]
+pub async fn totp_verify(
+    totp_request: TotpRequest,
+    form: Form<TotpVerifyForm<'_>>,
+    mut db: Connection<Database>,
+    auth_request: AuthenticationRequest<'_>,
+    cookie_jar: &CookieJar<'_>,
+) -> Result<Either<Redirect, Flash<Redirect>>> {
+    let mut transaction = db.begin().await?;
+
+    // Get app info
+    let app = App::get_valid_app(
+        &mut transaction,
+        auth_request.response_type,
+        auth_request.scope,
+        auth_request.client_id,
+        auth_request.redirect_uri,
+    )
+    .await?;
+
+    // Get totp request
+    let totp_request =
+        users::totp_login_request::TotpLoginRequest::get_one(&mut transaction, &totp_request.0)
+            .await?
+            .ok_or_else(|| Error::not_found("Failed to find totp request"))?;
+
+    if totp_request.has_expired() {
+        return Err(Error::bad_request("Totp request has expired"));
+    }
+
+    if totp_request.used_at().is_some() {
+        return Err(Error::bad_request("Totp request has been used"));
+    }
+
+    // Get user
+    let user = User::get_by_id(&mut transaction, totp_request.user())
+        .await?
+        .ok_or_else(|| Error::not_found("Failed to find user"))?;
+
+    transaction.commit().await?;
+
+    let totp_secret = user
+        .totp_secret()
+        .ok_or_else(|| Error::bad_request("TOTP is not enabled for user"))?;
+
+    // Create totp
+    let totp = totp::new(totp_secret, None, user.username().to_string())?;
+
+    // Verify totp code
+    if !totp.check_current(form.code)? {
+        return Ok(Either::Right(Flash::new(
+            Redirect::to(uri!(totp_page(auth_request))),
+            FlashKind::Danger,
+            "Wrong code. Please try again.",
+        )));
+    }
+
+    // Generate authorization code
+    let code = task::spawn_blocking(|| SecretString::new(AUTHORIZATION_CODE_LEN)).await?;
+
+    let mut transaction = db.begin().await?;
+
+    // Save authorization code
+    AuthorizationCode::insert(&mut transaction, code.as_ref(), app.id(), user.id()).await?;
+
+    // Mark totp token as used
+    totp_request.use_code(&mut transaction).await?;
+
+    transaction.commit().await?;
+
+    // Delete cookie
+    cookie_jar.remove(Cookie::named(TOTP_REQUEST_COOKIE_NAME));
+
+    // Construct uri to redirect to
+    let uri = redirect_uri(auth_request, &app, &code);
+
+    Ok(Either::Left(Redirect::found(uri)))
+}
diff --git a/crates/ezidam/templates/pages/oauth/totp.html.tera b/crates/ezidam/templates/pages/oauth/totp.html.tera
new file mode 100644
index 0000000..650d31a
--- /dev/null
+++ b/crates/ezidam/templates/pages/oauth/totp.html.tera
@@ -0,0 +1,74 @@
+{% extends "base" %}
+
+{% import "utils/user" as user %}
+{% import "utils/form" as form %}
+
+{% block page %}
+<body class=" d-flex flex-column">
+<script src="/js/demo-theme.min.js"></script>
+<div>
+    <div class="min-vh-100 d-flex flex-column justify-content-between">
+        <div class="container container-tight py-4">
+            <div class="text-center mb-4">
+                {% include "utils/logo" %}
+            </div>
+
+            {% if flash %}
+            <div class="alert alert-{{flash.0}}" role="alert">
+                <h4 class="alert-title">{{ flash.1 }}</h4>
+            </div>
+            {% endif %}
+
+            <div class="card card-md">
+                <div class="card-body">
+                    <div class="text-center mb-4">
+                        <div class="mb-4">
+                            <h2 class="card-title">Verify your account</h2>
+                        </div>
+                        <div class="mb-4">
+                            {{ user::avatar(username=username, name=name, size="lg", css="mb-3") }}
+                            <h3>
+                                {% if name %}
+                                {{ name }}
+                                {% else %}
+                                {{ username }}
+                                {% endif %}
+                            </h3>
+                        </div>
+                    </div>
+                    <form id="totp_form" action="" method="post" autocomplete="off" novalidate class="mt-4">
+                        <div class="mb-3">
+                            <label class="form-label required" for="code">
+                                Enter the code displayed on your device
+                            </label>
+                            <input
+                                    class="form-control"
+                                    type="text"
+                                    name="code"
+                                    id="code"
+                                    inputmode="numeric"
+                                    pattern="[0-9]*"
+                                    autocomplete="one-time-code"
+                                    required
+                            >
+                        </div>
+
+                        <div class="form-footer">
+                            <button type="submit" class="btn btn-primary w-100">Verify code</button>
+                        </div>
+                    </form>
+                </div>
+            </div>
+        </div>
+        {% include "shell/footer" %}
+    </div>
+</div>
+<!-- Libs JS -->
+<!-- Tabler Core -->
+<script src="/js/tabler.min.js" defer></script>
+<script src="/js/demo.min.js" defer></script>
+
+{{ form::disable_button_delay_submit(form_id="totp_form") }}
+
+</body>
+{% endblock page %}