125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
"""
|
|
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
|