PassKeeper/tests/test_mfa_gate.py

"""
Regression tests for the MFA bypass (review finding #1).

The bug had two halves:
  a) /login returned enc_key_salt in the MFA-pending response, before the second
     factor was verified.
  b) The recovery challenge was keyed on enc_key_salt, so anyone holding it
     could forge a proof and pull the whole encrypted vault from the
     unauthenticated /recovery/items — bypassing MFA entirely.

Chained, an attacker with only the master password could exfiltrate or take over
the account. These tests pin both halves shut.
"""
import hashlib
import hmac
import json

import pyotp

from app import db
from app.models.user import User
from app.services.auth_service import encrypt_totp_secret
from tests.conftest import add_item, auth_headers, login, make_user


def _enable_mfa(email='user@example.com'):
    """Turn on TOTP directly in the DB and return the plaintext secret."""
    user = User.query.filter_by(email=email).first()
    secret = pyotp.random_base32()
    enc, iv = encrypt_totp_secret(secret)
    user.totp_secret, user.totp_iv, user.totp_enabled = enc, iv, True
    db.session.commit()
    return secret


def test_login_withholds_enc_key_salt_until_mfa_is_verified(client, app):
    make_user(client)
    secret = _enable_mfa()

    res = login(client)
    body = res.get_json()

    assert res.status_code == 200
    assert body['mfa_required'] is True
    # The heart of finding #1a.
    assert 'enc_key_salt' not in body, (
        'enc_key_salt leaked before the second factor was verified'
    )

    res = client.post('/api/auth/mfa/verify', json={
        'mfa_token': body['mfa_token'], 'totp_code': pyotp.TOTP(secret).now(),
    })
    verified = res.get_json()
    assert res.status_code == 200
    # Released only now, once both factors are proven.
    assert verified['enc_key_salt'] == 'SALT-V1'


def test_login_without_mfa_still_returns_enc_key_salt(client):
    """The withholding must apply only to the MFA-pending path."""
    make_user(client)
    body = login(client).get_json()
    assert 'mfa_required' not in body
    assert body['enc_key_salt'] == 'SALT-V1'


# ── Finding #1b: the recovery proof must not be forgeable from enc_key_salt ──

def _setup_recovery(client, token, verifier):
    res = client.post('/api/auth/recovery/setup', headers=auth_headers(token), json={
        'recovery_enc_salt': 'RECOVERY-BLOB', 'recovery_iv': 'RECOVERY-IV',
        'recovery_verifier': verifier,
    })
    assert res.status_code == 200, res.get_json()


def test_recovery_proof_cannot_be_forged_from_enc_key_salt(client, app):
    """
    The attack: password known, second factor not. Previously the attacker could
    key the HMAC with enc_key_salt and walk away with every encrypted item.
    """
    token, _ = make_user(client)
    add_item(client, token)
    verifier = 'a' * 64
    _setup_recovery(client, token, verifier)

    nonce = client.get('/api/auth/recovery/data?email=user@example.com').get_json()['nonce']

    forged = hmac.new(b'SALT-V1', nonce.encode(), hashlib.sha256).hexdigest()
    res = client.get('/api/auth/recovery/items?email=user@example.com',
                     headers={'X-Recovery-Proof': forged})
    assert res.status_code == 401, 'enc_key_salt still forges a valid recovery proof'


def test_recovery_proof_from_verifier_is_accepted(client, app):
    """The legitimate holder of the recovery code must still get through."""
    token, _ = make_user(client)
    add_item(client, token)
    verifier = 'b' * 64
    _setup_recovery(client, token, verifier)

    data = client.get('/api/auth/recovery/data?email=user@example.com').get_json()
    assert data['proof_scheme'] == 'verifier'

    proof = hmac.new(verifier.encode(), data['nonce'].encode(), hashlib.sha256).hexdigest()
    res = client.get('/api/auth/recovery/items?email=user@example.com',
                     headers={'X-Recovery-Proof': proof})
    assert res.status_code == 200, res.get_json()
    assert len(res.get_json()['items']) == 1


def test_legacy_account_falls_back_to_enc_key_salt_proof(client, app):
    """
    Recovery codes created before recovery_verifier must keep working, and be
    reported as legacy so the UI can prompt a regeneration.
    """
    token, _ = make_user(client)
    add_item(client, token)
    user = User.query.filter_by(email='user@example.com').first()
    user.recovery_enc_salt, user.recovery_iv = 'BLOB', 'IV'
    user.recovery_verifier = None          # pre-migration state
    db.session.commit()

    status = client.get('/api/auth/recovery/status', headers=auth_headers(token)).get_json()
    assert status['recovery_configured'] is True
    assert status['recovery_is_legacy'] is True

    data = client.get('/api/auth/recovery/data?email=user@example.com').get_json()
    assert data['proof_scheme'] == 'legacy'

    proof = hmac.new(b'SALT-V1', data['nonce'].encode(), hashlib.sha256).hexdigest()
    res = client.get('/api/auth/recovery/items?email=user@example.com',
                     headers={'X-Recovery-Proof': proof})
    assert res.status_code == 200


def test_recovery_setup_rejects_malformed_verifier(client):
    token, _ = make_user(client)
    for bad in ('', 'short', 'g' * 64, 'A' * 63):
        res = client.post('/api/auth/recovery/setup', headers=auth_headers(token), json={
            'recovery_enc_salt': 'B', 'recovery_iv': 'IV', 'recovery_verifier': bad,
        })
        assert res.status_code == 400, f'accepted malformed verifier {bad!r}'