PassKeeper/tests/test_key_rotation.py

"""
Regression tests for silent vault destruction (review finding #2).

change_password and /recover both rotate enc_key_salt, which invalidates every
ciphertext under the previous vault key. The client re-encrypts each item and
sends it back — but nothing checked that the payload actually covered every
item. A short payload rotated the key anyway and left the missing items
permanently undecryptable, with no error and a success entry in the audit log.
"""
import hashlib
import hmac

from app import db
from app.models.audit_log import AuditLog
from app.models.user import User
from app.models.vault_item import VaultItem
from tests.conftest import add_item, auth_headers, login, make_user

VERIFIER = 'c' * 64


def _payload(ids):
    return [{'id': i, 'enc_data': f'NEW-{i}', 'iv': f'NEWIV-{i}'} for i in ids]


def _change_password(client, token, items, allow_partial=None):
    body = {
        'current_auth_hash': 'AUTH-HASH-V1',
        'new_auth_hash': 'AUTH-HASH-V2',
        'new_enc_key_salt': 'SALT-V2',
        'items': items,
    }
    if allow_partial is not None:
        body['allow_partial'] = allow_partial
    return client.post('/api/auth/change-password', headers=auth_headers(token), json=body)


# -- change_password ---------------------------------------------------------

def test_full_payload_rotates_everything(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(3)]

    res = _change_password(client, token, _payload(ids))
    assert res.status_code == 200, res.get_json()

    for i in ids:
        assert db.session.get(VaultItem, i).enc_data == f'NEW-{i}'
    user = User.query.filter_by(email='user@example.com').first()
    assert user.enc_key_salt == 'SALT-V2'


def test_short_payload_is_refused_and_nothing_changes(client, app):
    """The exact data-loss bug: one item omitted from the re-encryption."""
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(3)]

    res = _change_password(client, token, _payload(ids[:2]))   # third omitted
    assert res.status_code == 409, res.get_json()
    body = res.get_json()
    assert body['code'] == 'incomplete_reencryption'
    assert (body['expected'], body['received']) == (3, 2)

    # Nothing may have been committed: salt unchanged, ciphertext untouched.
    user = User.query.filter_by(email='user@example.com').first()
    assert user.enc_key_salt == 'SALT-V1', 'key rotated despite refusal'
    assert user.token_epoch == 0
    for i in ids:
        assert db.session.get(VaultItem, i).enc_data == 'CT'

    # The old password must still work.
    assert login(client).status_code == 200


def test_empty_payload_against_populated_vault_is_refused(client, app):
    token, _ = make_user(client)
    for _ in range(4):
        add_item(client, token)

    res = _change_password(client, token, [])
    assert res.status_code == 409
    assert res.get_json()['received'] == 0
    assert User.query.filter_by(email='user@example.com').first().enc_key_salt == 'SALT-V1'


def test_refusal_is_audited(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(2)]
    _change_password(client, token, _payload(ids[:1]))

    entry = (AuditLog.query
             .filter_by(action='auth.change_password_failed')
             .order_by(AuditLog.id.desc()).first())
    assert entry is not None
    assert '1 of 2' in entry.detail


def test_another_users_item_does_not_count_toward_coverage(client, app):
    """A foreign id must not pad the payload up to the expected count."""
    token_a, _ = make_user(client, 'a@example.com', 'HASH-A', 'SALT-A')
    token_b, _ = make_user(client, 'b@example.com', 'HASH-B', 'SALT-B')
    a_ids = [add_item(client, token_a) for _ in range(2)]
    b_id = add_item(client, token_b)

    res = client.post('/api/auth/change-password', headers=auth_headers(token_a), json={
        'current_auth_hash': 'HASH-A', 'new_auth_hash': 'HASH-A2',
        'new_enc_key_salt': 'SALT-A2',
        'items': _payload([a_ids[0], b_id]),   # b_id does not belong to user A
    })
    assert res.status_code == 409
    assert (res.get_json()['expected'], res.get_json()['received']) == (2, 1)
    assert db.session.get(VaultItem, b_id).enc_data == 'CT'


# -- /recover ----------------------------------------------------------------

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


def _proof(client, email='user@example.com'):
    nonce = client.get(f'/api/auth/recovery/data?email={email}').get_json()['nonce']
    return hmac.new(VERIFIER.encode(), nonce.encode(), hashlib.sha256).hexdigest()


def _consume_challenge(client, proof):
    """Fetch items exactly as recover.js does, which rotates the challenge."""
    return client.get('/api/auth/recovery/items?email=user@example.com',
                      headers={'X-Recovery-Proof': proof})


def test_recover_refuses_short_payload(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(3)]
    _setup_recovery(client, token)
    proof = _proof(client)
    _consume_challenge(client, proof)

    res = client.post('/api/auth/recover', json={
        'email': 'user@example.com', 'new_auth_hash': 'AUTH-HASH-V2',
        'new_enc_key_salt': 'SALT-V2', 'recovery_proof': proof,
        'items': _payload(ids[:2]),
    })
    assert res.status_code == 409, res.get_json()
    user = User.query.filter_by(email='user@example.com').first()
    assert user.enc_key_salt == 'SALT-V1'
    assert user.recovery_enc_salt == 'BLOB', 'recovery code consumed despite refusal'


def test_recover_allows_partial_when_explicitly_confirmed(client, app):
    """
    The escape hatch exists because refusing outright would leave a locked-out
    user with no way into their account at all.
    """
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(3)]
    _setup_recovery(client, token)
    proof = _proof(client)
    _consume_challenge(client, proof)

    res = client.post('/api/auth/recover', json={
        'email': 'user@example.com', 'new_auth_hash': 'AUTH-HASH-V2',
        'new_enc_key_salt': 'SALT-V2', 'recovery_proof': proof,
        'items': _payload(ids[:2]), 'allow_partial': True,
    })
    assert res.status_code == 200, res.get_json()

    entry = (AuditLog.query.filter_by(action='auth.recovery_success')
             .order_by(AuditLog.id.desc()).first())
    assert 'PARTIAL' in entry.detail, 'partial recovery not flagged in the audit log'


def test_recover_full_payload_succeeds_and_consumes_the_code(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(2)]
    _setup_recovery(client, token)
    proof = _proof(client)
    _consume_challenge(client, proof)

    res = client.post('/api/auth/recover', json={
        'email': 'user@example.com', 'new_auth_hash': 'AUTH-HASH-V2',
        'new_enc_key_salt': 'SALT-V2', 'recovery_proof': proof,
        'items': _payload(ids),
    })
    assert res.status_code == 200, res.get_json()

    user = User.query.filter_by(email='user@example.com').first()
    assert user.enc_key_salt == 'SALT-V2'
    assert user.recovery_enc_salt is None
    assert user.recovery_verifier is None
    assert login(client, auth_hash='AUTH-HASH-V2').status_code == 200


def test_change_password_ignores_allow_partial(client, app):
    """
    Recovery has a partial-completion escape hatch; changing the password must
    not. The current password keeps working, so there is never a reason to
    accept permanent data loss here — the server refuses even if a client asks.
    """
    token, _ = make_user(client)
    ids = [add_item(client, token) for _ in range(3)]

    res = _change_password(client, token, _payload(ids[:1]), allow_partial=True)
    assert res.status_code == 409, 'server honoured allow_partial on change-password'
    assert User.query.filter_by(email='user@example.com').first().enc_key_salt == 'SALT-V1'