From 650d6e35f048de9ba9d93e0d12dd0111e91798b8 Mon Sep 17 00:00:00 2001
From: devStorm <59678453+developStorm@users.noreply.github.com>
Date: Tue, 5 May 2020 05:03:29 -0700
Subject: [PATCH] FIDO login middleware

---
 app/auth/__init__.py                          |   1 +
 app/auth/templates/auth/fido.html             |  60 ++++++++++
 app/auth/views/fido.py                        | 105 ++++++++++++++++++
 .../templates/dashboard/fido_setup.html       |   2 +-
 app/dashboard/views/fido_setup.py             |   4 +-
 5 files changed, 169 insertions(+), 3 deletions(-)
 create mode 100644 app/auth/templates/auth/fido.html
 create mode 100644 app/auth/views/fido.py

diff --git a/app/auth/__init__.py b/app/auth/__init__.py
index 322cbd81..6a51b706 100644
--- a/app/auth/__init__.py
+++ b/app/auth/__init__.py
@@ -11,5 +11,6 @@ from .views import (
     facebook,
     change_email,
     mfa,
+    fido,
     social,
 )
diff --git a/app/auth/templates/auth/fido.html b/app/auth/templates/auth/fido.html
new file mode 100644
index 00000000..abde9d4f
--- /dev/null
+++ b/app/auth/templates/auth/fido.html
@@ -0,0 +1,60 @@
+{% extends "single.html" %}
+
+
+{% block title %}
+  Verify Your Security Key
+{% endblock %}
+
+{% block head %}
+  <script src="{{ url_for('static', filename='assets/js/vendors/base64.js') }}"></script>
+  <script src="{{ url_for('static', filename='assets/js/vendors/webauthn.js') }}"></script>
+{% endblock %}
+
+{% block single_content %}
+  <div class="bg-white p-6" style="margin: auto">
+
+    <div class="mb-2">
+      Your account is protected with your security key (WebAuthn). <br><br>
+      Follow your browser's steps to continue the sign-in process.
+    </div>
+
+    <form id="formRegisterKey" method="post">
+      {{ fido_token_form.csrf_token }}
+      {{ fido_token_form.sk_assertion(class="form-control", placeholder="") }}
+    </form>
+    <div class="text-center">
+      <button id="btnVerifyKey" class="btn btn-success mt-2">Use your security key</button>
+    </div>
+
+    <script>
+      async function verifyKey () {
+        $("#btnVerifyKey").prop('disabled', true);
+        $("#btnVerifyKey").text('Waiting for Security Key...');
+
+        const credentialRequestOptions = transformCredentialRequestOptions(
+          JSON.parse('{{webauthn_assertion_options|tojson|safe}}')
+        )
+
+        let assertion;
+        try {
+          assertion = await navigator.credentials.get({
+            publicKey: credentialRequestOptions
+          });
+        } catch (err) {
+          toastr.error("An error occurred when we trying to verify your key.");
+          $("#btnVerifyKey").prop('disabled', false);
+          $("#btnVerifyKey").text('Use your security key');
+          return console.error("Error when trying to get credential:", err);
+        }
+
+        const skAssertion = transformAssertionForServer(assertion);
+        $('#sk_assertion').val(JSON.stringify(skAssertion));
+        $('#formRegisterKey').submit();
+      }
+
+      $("#btnVerifyKey").click(verifyKey);
+    </script>
+
+  </div>
+
+{% endblock %}
\ No newline at end of file
diff --git a/app/auth/views/fido.py b/app/auth/views/fido.py
new file mode 100644
index 00000000..e4ef5ddd
--- /dev/null
+++ b/app/auth/views/fido.py
@@ -0,0 +1,105 @@
+import json
+import secrets
+import webauthn
+from app.config import URL as SITE_URL
+from urllib.parse import urlparse
+
+from flask import request, render_template, redirect, url_for, flash, session
+from flask_login import login_user
+from flask_wtf import FlaskForm
+from wtforms import HiddenField, validators
+
+from app.auth.base import auth_bp
+from app.config import MFA_USER_ID
+from app.log import LOG
+from app.models import User
+from app.extensions import db
+
+
+class FidoTokenForm(FlaskForm):
+    sk_assertion = HiddenField("sk_assertion", validators=[validators.DataRequired()])
+
+
+@auth_bp.route("/fido", methods=["GET", "POST"])
+def fido():
+    # passed from login page
+    user_id = session.get(MFA_USER_ID)
+
+    # user access this page directly without passing by login page
+    if not user_id:
+        flash("Unknown error, redirect back to main page", "warning")
+        return redirect(url_for("auth.login"))
+
+    user = User.get(user_id)
+
+    if not (user and (user.fido_uuid is not None)):
+        flash("Only user with security key linked should go to this page", "warning")
+        return redirect(url_for("auth.login"))
+
+    fido_token_form = FidoTokenForm()
+
+    next_url = request.args.get("next")
+
+    rp_id = urlparse(SITE_URL).hostname
+
+    webauthn_user = webauthn.WebAuthnUser(
+            user.fido_uuid, user.email, user.name, False,
+            user.fido_credential_id, user.fido_pk, user.fido_sign_count, rp_id)
+
+    # Handling POST requests
+    if fido_token_form.validate_on_submit():
+        try:
+            sk_assertion = json.loads(fido_token_form.sk_assertion.data)
+        except Exception as e:
+            flash('Key registration failed. Error: Invalid Payload', "warning")
+            return redirect(url_for("dashboard.index"))
+                    
+        challenge = session['fido_challenge']
+        credential_id = sk_assertion['id']
+
+        webauthn_assertion_response = webauthn.WebAuthnAssertionResponse(
+            webauthn_user,
+            sk_assertion,
+            challenge,
+            SITE_URL,
+            uv_required=False
+        )
+
+        new_sign_count = False
+        new_sign_count = webauthn_assertion_response.verify()
+        try:
+            pass
+        except Exception as e:
+            flash('Key verification failed. Error: {}'.format(e), "warning")
+
+        if new_sign_count != False:
+            user.fido_sign_count = new_sign_count
+            db.session.commit()
+            del session[MFA_USER_ID]
+
+            login_user(user)
+            flash(f"Welcome back {user.name}!", "success")
+
+            # User comes to login page from another page
+            if next_url:
+                LOG.debug("redirect user to %s", next_url)
+                return redirect(next_url)
+            else:
+                LOG.debug("redirect user to dashboard")
+                return redirect(url_for("dashboard.index"))
+        else:
+            # Verification failed, put else here to make structure clear
+            pass
+    
+    # Prepare infomation for key registration process
+    session.pop('challenge', None)
+    challenge = secrets.token_urlsafe(32)
+    
+    session['fido_challenge'] = challenge.rstrip('=')
+
+    webauthn_assertion_options = webauthn.WebAuthnAssertionOptions(
+        webauthn_user, challenge)
+    webauthn_assertion_options = webauthn_assertion_options.assertion_dict
+
+    return render_template("auth/fido.html", fido_token_form=fido_token_form, 
+                        webauthn_assertion_options=webauthn_assertion_options)
\ No newline at end of file
diff --git a/app/dashboard/templates/dashboard/fido_setup.html b/app/dashboard/templates/dashboard/fido_setup.html
index 713c29b2..3c8b66e5 100644
--- a/app/dashboard/templates/dashboard/fido_setup.html
+++ b/app/dashboard/templates/dashboard/fido_setup.html
@@ -20,7 +20,7 @@
       {{ fido_token_form.sk_assertion(class="form-control", placeholder="") }}
     </form>
     <div class="text-center">
-      <button id="btnRegisterKey" onclick="registerKey($, toastr)" class="btn btn-lg btn-primary mt-2">Register Key</button>
+      <button id="btnRegisterKey" class="btn btn-lg btn-primary mt-2">Register Key</button>
     </div>
 
     <script>
diff --git a/app/dashboard/views/fido_setup.py b/app/dashboard/views/fido_setup.py
index 399b26e4..b614ef64 100644
--- a/app/dashboard/views/fido_setup.py
+++ b/app/dashboard/views/fido_setup.py
@@ -55,8 +55,8 @@ def fido_setup():
             flash('Key registration failed. Error: {}'.format(e), "warning")
             return redirect(url_for("dashboard.index"))
 
-        current_user.fido_pk = fido_uuid
-        current_user.fido_uuid = str(fido_credential.public_key, "utf-8")
+        current_user.fido_pk = str(fido_credential.public_key, "utf-8")
+        current_user.fido_uuid = fido_uuid
         current_user.fido_sign_count = fido_credential.sign_count
         current_user.fido_credential_id = str(fido_credential.credential_id, "utf-8")
         db.session.commit()