2FA Testing

src/api/login.py

@@ -32,6 +32,7 @@ async def login(login_form: LoginRequest, response: Response, request: Request):
     Can also return a `Set-Cookie` header with the session token. (See Config)
     """
     user = get_user_email_or_username(login_form.identifier)
+    # Check if User can be found
     if user is None:
         raise HTTPException(detail="User not found", status_code=404)
     # Check if Password Exists (or if OAuth SignIn)
@@ -43,8 +44,9 @@ async def login(login_form: LoginRequest, response: Response, request: Request):
 
     uid_email_key = "invallogin:" + user["email"]
 
+    # Get Failed Attempts from Redis
     failed_attempts = r.get(uid_email_key)
-    # Max Login Attempts enabed? and already failed attempts? and reached max?
+    # Max Login Attempts enabled? and already failed attempts? and reached max?
     if (
         SecurityConfig.max_login_attempts > 0
         and failed_attempts
@@ -62,6 +64,7 @@ async def login(login_form: LoginRequest, response: Response, request: Request):
         login_form.password.get_secret_value().encode("utf-8"),
         user["password"].encode("utf-8"),
     ):
+        # Wrong Password
         if SecurityConfig.max_login_attempts > 0:
             r.incrby(uid_email_key, 1)
             r.expire(uid_email_key, SecurityConfig.expire_unfinished_timeout * 60)

src/api/oauth_providers/__init__.py

@@ -1,5 +1,4 @@
 from fastapi import APIRouter
-import sys
 from tools.conf import SignupConfig
 
 router = APIRouter(
@@ -8,9 +7,6 @@
     dependencies=[],
 )
 
-if "pytest" in sys.modules:
-    SignupConfig.oauth_providers = []
-
 if "google" in SignupConfig.oauth_providers:
     from .google import router as ggl
 

src/api/twofactor.py

@@ -25,12 +25,15 @@ def enable_2fa_temp(user: dict) -> str:
     secret_bytes = pyotp.random_base32()
     totp = pyotp.TOTP(secret_bytes)
 
+    # Generate provisioning URL
     prov_url = totp.provisioning_uri(
         name=user["email"],
         issuer_name=AccountFeaturesConfig.issuer_name_2fa,
         image=AccountFeaturesConfig.issuer_image_url_2fa,
     )
+    # Set the secret in Redis (for confirmation)
     r.setex("2fa:" + bson.ObjectId(user["_id"]).__str__(), 60, secret_bytes)
+    # Return the provisioning URL for further processing
     return prov_url
 
 
@@ -62,6 +65,7 @@ async def enable_2fa_qr(user=Depends(get_dangerous_user_dep)):
     if not AccountFeaturesConfig.qr_code_endpoint_2fa:
         raise HTTPException(status_code=403, detail="QR Code endpoint is disabled.")
     prov_url = enable_2fa_temp(user)
+    # Generate QR Code for the Provisioning URL
     qr = QRCode(image_factory=qrcode.image.svg.SvgPathImage)
     qr.add_data(prov_url)
     return Response(qr.make_image().to_string(), media_type="image/svg+xml")
@@ -84,9 +88,11 @@ async def confirm_enable_2fa(
     ## Description
     This endpoint is used to confirm the enablement of 2FA for the user.
     """
+    # Retrieve Secret from redis
     secret = r.get("2fa:" + bson.ObjectId(user["_id"]).__str__())
     if not secret:
         raise HTTPException(status_code=400, detail="2FA activation expired")
+    # Initialize TOTP Generator from the Secret
     totp = pyotp.TOTP(secret)
     if totp.verify(str(code.code)):
         # Persist 2FA in the Database

src/api/twofactor_test.py

@@ -0,0 +1,55 @@
+from fastapi.testclient import TestClient
+import pytest
+import pyotp
+
+from .main import app
+
+client = TestClient(app)
+
+
+# Request 2FA
+@pytest.fixture
+def request2fa_fixture(fixturesessiontoken_user):
+    client.cookies.set("session", fixturesessiontoken_user[0])
+    response = client.post("/2fa/enable")
+    resp_js = response.json()
+    return resp_js["provision_uri"]
+
+
+def test_request_2fa(fixturesessiontoken_user):
+    client.cookies.set("session", fixturesessiontoken_user[0])
+    response = client.post("/2fa/enable")
+    resp_js = response.json()
+    assert response.status_code == 200
+    assert resp_js["provision_uri"] is not None
+
+
+def test_request_2fa_qr(fixturesessiontoken_user):
+    client.cookies.set("session", fixturesessiontoken_user[0])
+    response = client.get("/2fa/enable")
+    assert response.status_code == 200
+    assert response.headers.get("Content-Type") == "image/svg+xml"
+
+
+# Enable 2FA
+def test_enable_2fa_prov_url(fixturesessiontoken_user, request2fa_fixture):
+    otp_instance: pyotp.TOTP = pyotp.parse_uri(request2fa_fixture)
+    client.cookies.set("session", fixturesessiontoken_user[0])
+    response = client.post("/2fa/confirm-enable", json={"code": otp_instance.now()})
+    assert response.status_code == 204
+
+
+def test_enable_2fa_prov_url_invalid_code(fixturesessiontoken_user, request2fa_fixture):
+    client.cookies.set("session", fixturesessiontoken_user[0])
+    response = client.post("/2fa/confirm-enable", json={"code": 123456})
+    assert response.status_code == 400
+
+
+def test_request_2fa_no_login(fixturesessiontoken_user):
+    response = client.post("/2fa/enable")
+    assert response.status_code == 401
+
+
+def test_request_2fa_no_login_qr(fixturesessiontoken_user):
+    response = client.get("/2fa/enable")
+    assert response.status_code == 401

src/tools/testing_config.json

@@ -5,9 +5,7 @@
         "conf_code_complexity": 1,
         "enable_welcome_email": false,
         "oauth": {
-            "providers_enabled": [
-                "google"
-            ],
+            "providers_enabled": [],
             "base_url": "http://localhost:3250/"
         }
     },