add /api/auth/mfa

README.md

@@ -566,6 +566,20 @@ Output:
 - mfa_key: only useful when user enables MFA. In this case, user needs to enter their OTP token in order to login.
 - api_key: if MFA is not enabled, the `api key` is returned right away.
 
+The `api_key` is used in all subsequent requests. It's empty if MFA is enabled.
+If user hasn't enabled MFA, `mfa_key` is empty.
+
+#### POST /api/auth/mfa
+
+Input:
+- mfa_token: OTP token that user enters
+- mfa_key: MFA key obtained in previous auth request, e.g. /api/auth/login
+- device: the device name, used to create an ApiKey associated with this device
+
+Output:
+- name: user name, could be an empty string
+- api_key: if MFA is not enabled, the `api key` is returned right away.
+
 The `api_key` is used in all subsequent requests. It's empty if MFA is enabled.
 If user hasn't enabled MFA, `mfa_key` is empty. 
 

app/api/__init__.py

@@ -1 +1,8 @@
-from .views import alias_options, new_custom_alias, new_random_alias, user_info, auth_login
+from .views import (
+    alias_options,
+    new_custom_alias,
+    new_random_alias,
+    user_info,
+    auth_login,
+    auth_mfa,
+)

app/api/views/auth_mfa.py

@@ -0,0 +1,65 @@
+import pyotp
+from flask import jsonify, request
+from flask_cors import cross_origin
+from itsdangerous import Signer, BadSignature
+
+from app.api.base import api_bp
+from app.config import FLASK_SECRET
+from app.extensions import db
+from app.models import User, ApiKey
+
+
+@api_bp.route("/auth/mfa", methods=["POST"])
+@cross_origin()
+def auth_mfa():
+    """
+    Validate the OTP Token
+    Input:
+        mfa_token: OTP token that user enters
+        mfa_key: MFA key obtained in previous auth request, e.g. /api/auth/login
+        device: the device name, used to create an ApiKey associated with this device
+    Output:
+        200 and user info containing:
+        {
+            name: "John Wick",
+            api_key: "a long string"
+        }
+
+    """
+    data = request.get_json()
+    if not data:
+        return jsonify(error="request body cannot be empty"), 400
+
+    mfa_token = data.get("mfa_token")
+    mfa_key = data.get("mfa_key")
+    device = data.get("device")
+
+    s = Signer(FLASK_SECRET)
+    try:
+        user_id = int(s.unsign(mfa_key))
+    except BadSignature:
+        return jsonify(error="Invalid mfa_key"), 400
+
+    user = User.get(user_id)
+
+    if not user:
+        return jsonify(error="Invalid mfa_key"), 400
+    elif not user.enable_otp:
+        return (
+            jsonify(error="This endpoint should only be used by user who enables MFA"),
+            400,
+        )
+
+    totp = pyotp.TOTP(user.otp_secret)
+    if not totp.verify(mfa_token):
+        return jsonify(error="Wrong TOTP Token"), 400
+
+    ret = {
+        "name": user.name,
+    }
+
+    api_key = ApiKey.create(user.id, device)
+    db.session.commit()
+    ret["api_key"] = api_key.code
+
+    return jsonify(**ret), 200

tests/api/test_auth_mfa.py

@@ -0,0 +1,58 @@
+import pyotp
+from flask import url_for
+from itsdangerous import Signer
+
+from app.config import FLASK_SECRET
+from app.extensions import db
+from app.models import User
+
+
+def test_auth_mfa_success(flask_client):
+    user = User.create(
+        email="a@b.c",
+        password="password",
+        name="Test User",
+        activated=True,
+        enable_otp=True,
+        otp_secret="base32secret3232",
+    )
+    db.session.commit()
+
+    totp = pyotp.TOTP(user.otp_secret)
+    s = Signer(FLASK_SECRET)
+    mfa_key = s.sign(str(user.id))
+
+    r = flask_client.post(
+        url_for("api.auth_mfa"),
+        json={"mfa_token": totp.now(), "mfa_key": mfa_key, "device": "Test Device"},
+    )
+
+    assert r.status_code == 200
+    assert r.json["api_key"]
+    assert r.json["name"] == "Test User"
+
+
+def test_auth_wrong_mfa_key(flask_client):
+    user = User.create(
+        email="a@b.c",
+        password="password",
+        name="Test User",
+        activated=True,
+        enable_otp=True,
+        otp_secret="base32secret3232",
+    )
+    db.session.commit()
+
+    totp = pyotp.TOTP(user.otp_secret)
+
+    r = flask_client.post(
+        url_for("api.auth_mfa"),
+        json={
+            "mfa_token": totp.now(),
+            "mfa_key": "wrong mfa key",
+            "device": "Test Device",
+        },
+    )
+
+    assert r.status_code == 400
+    assert r.json["error"]