July 4 - Implement TOTP 2FA

tests/test_mfa.py

@@ -0,0 +1,124 @@
+"""
+tests/test_mfa.py
+-----------------
+End-to-end tests for phase35 two-factor authentication (main app).
+
+Drives the real login → challenge flow through the test client and asserts:
+
+  * an MFA-enabled account is NOT authenticated until the code step passes
+  * a correct TOTP completes login; a wrong code does not
+  * a recovery code works once and is then consumed
+  * a non-MFA account logs in directly (no challenge, no regression)
+  * the mfa utility verifies/rejects codes correctly
+"""
+
+import pyotp
+import pytest
+
+
+@pytest.fixture
+def client(app):
+    """Fresh schema + client, rate limiter disabled for deterministic runs."""
+    from app import db, limiter
+    limiter.enabled = False
+    with app.app_context():
+        db.drop_all()
+        db.create_all()
+        yield app.test_client()
+        db.session.remove()
+    limiter.enabled = True
+
+
+def _make_user(username='admin1', role='admin', mfa=False):
+    from app import db
+    from app.models.user import User
+    from app.utils import mfa as mfa_util
+
+    u = User(username=username, full_name='Ada Admin',
+             email=f'{username}@example.com', role=role, active=True,
+             password_set=True)
+    u.set_password('pw-correct')
+    secret = None
+    if mfa:
+        secret = mfa_util.new_secret()
+        u.mfa_secret = secret
+        u.mfa_enabled = True
+        # store two known recovery codes (hashed)
+        u.mfa_recovery_codes = [
+            __import__('werkzeug.security', fromlist=['generate_password_hash'])
+            .generate_password_hash(c) for c in ('aaaa-bbbb', 'cccc-dddd')
+        ]
+    db.session.add(u)
+    db.session.commit()
+    return u.id, secret
+
+
+def _is_authenticated(client):
+    """Profile page is login-gated: 200 == authenticated, 302 == not."""
+    return client.get('/auth/profile').status_code == 200
+
+
+def test_mfa_user_must_pass_second_factor(client):
+    uid, secret = _make_user(mfa=True)
+
+    # Correct password → redirected to the challenge, NOT yet authenticated.
+    resp = client.post('/auth/login',
+                       data={'username': 'admin1', 'password': 'pw-correct'})
+    assert resp.status_code == 302
+    assert '/auth/mfa' in resp.headers['Location']
+    assert not _is_authenticated(client)
+
+    # Wrong code keeps us out.
+    bad = client.post('/auth/mfa', data={'code': '000000'})
+    assert not _is_authenticated(client)
+
+    # Correct TOTP completes login.
+    code = pyotp.TOTP(secret).now()
+    ok = client.post('/auth/mfa', data={'code': code})
+    assert ok.status_code == 302
+    assert _is_authenticated(client)
+
+
+def test_recovery_code_logs_in_and_is_consumed(client):
+    uid, secret = _make_user(mfa=True)
+
+    client.post('/auth/login', data={'username': 'admin1', 'password': 'pw-correct'})
+    # Use a recovery code.
+    r = client.post('/auth/mfa', data={'code': 'aaaa-bbbb', 'recovery': '1'})
+    assert r.status_code == 302
+    assert _is_authenticated(client)
+
+    # The consumed code no longer works on a fresh login.
+    client.get('/auth/logout')
+    client.post('/auth/login', data={'username': 'admin1', 'password': 'pw-correct'})
+    reuse = client.post('/auth/mfa', data={'code': 'aaaa-bbbb', 'recovery': '1'})
+    assert not _is_authenticated(client)
+    # The other, unused code still works.
+    good = client.post('/auth/mfa', data={'code': 'cccc-dddd', 'recovery': '1'})
+    assert _is_authenticated(client)
+
+
+def test_non_mfa_user_logs_in_directly(client):
+    _make_user(username='pm1', role='project_manager', mfa=False)
+    resp = client.post('/auth/login',
+                       data={'username': 'pm1', 'password': 'pw-correct'})
+    assert resp.status_code == 302
+    assert '/auth/mfa' not in resp.headers['Location']
+    assert _is_authenticated(client)
+
+
+def test_wrong_password_never_reaches_challenge(client):
+    _make_user(mfa=True)
+    client.post('/auth/login', data={'username': 'admin1', 'password': 'WRONG'})
+    # No pending challenge, not authenticated.
+    assert client.get('/auth/mfa').status_code == 302   # bounced back to login
+    assert not _is_authenticated(client)
+
+
+def test_mfa_util_verifies_and_rejects():
+    from app.utils import mfa
+    s = mfa.new_secret()
+    assert mfa.verify_totp(s, pyotp.TOTP(s).now()) is True
+    assert mfa.verify_totp(s, '000000') is False
+    assert mfa.verify_totp(s, 'not-a-code') is False
+    assert mfa.verify_totp('', '123456') is False