PassKeeper/tests/test_session_revocation.py

"""
Regression tests for session revocation (review finding #3).

Changing the master password used to leave every outstanding access and refresh
token valid. The response said "Please log in again" but nothing enforced it, so
a stolen refresh token kept working for its full 7-day lifetime after the victim
changed the password it was obtained under.

Every JWT now carries an `epoch` claim checked against users.token_epoch.

Also covers the sub-finding that require_jwt never confirmed the user still
existed: a valid token for a deleted account dereferenced None and returned 500.
"""
import jwt as pyjwt

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


def _rotate_password(client, token, ids):
    """Complete a full, valid password change."""
    res = client.post('/api/auth/change-password', headers=auth_headers(token), json={
        'current_auth_hash': 'AUTH-HASH-V1',
        'new_auth_hash': 'AUTH-HASH-V2',
        'new_enc_key_salt': 'SALT-V2',
        'items': [{'id': i, 'enc_data': f'NEW-{i}', 'iv': f'IV-{i}'} for i in ids],
    })
    assert res.status_code == 200, res.get_json()
    return res


def test_access_token_is_revoked_by_password_change(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token)]

    assert client.get('/api/vault', headers=auth_headers(token)).status_code == 200
    _rotate_password(client, token, ids)

    res = client.get('/api/vault', headers=auth_headers(token))
    assert res.status_code == 401, 'access token survived the password change'


def test_refresh_token_is_revoked_by_password_change(client, app):
    """
    The more damaging half: a refresh token is valid for 7 days and can mint
    fresh access tokens indefinitely.
    """
    token, refresh = make_user(client)
    ids = [add_item(client, token)]
    _rotate_password(client, token, ids)

    res = client.post('/api/auth/refresh', json={'refresh_token': refresh})
    assert res.status_code == 401, 'refresh token survived the password change'


def test_epoch_increments_and_new_login_works(client, app):
    token, _ = make_user(client)
    ids = [add_item(client, token)]
    _rotate_password(client, token, ids)

    user = User.query.filter_by(email='user@example.com').first()
    assert user.token_epoch == 1

    res = login(client, auth_hash='AUTH-HASH-V2')
    assert res.status_code == 200
    new_token = res.get_json()['access_token']
    assert client.get('/api/vault', headers=auth_headers(new_token)).status_code == 200


def test_recovery_also_revokes_prior_sessions(client, app):
    """An attacker holding a token must not survive the victim recovering."""
    import hashlib
    import hmac

    verifier = 'd' * 64
    token, _ = make_user(client)
    ids = [add_item(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

    nonce = client.get('/api/auth/recovery/data?email=user@example.com').get_json()['nonce']
    proof = hmac.new(verifier.encode(), nonce.encode(), hashlib.sha256).hexdigest()
    client.get('/api/auth/recovery/items?email=user@example.com',
               headers={'X-Recovery-Proof': 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': [{'id': i, 'enc_data': f'NEW-{i}', 'iv': f'IV-{i}'} for i in ids],
    })
    assert res.status_code == 200, res.get_json()

    assert client.get('/api/vault', headers=auth_headers(token)).status_code == 401
    # The tokens handed back by /recover must carry the NEW epoch and work.
    fresh = res.get_json()['access_token']
    assert client.get('/api/vault', headers=auth_headers(fresh)).status_code == 200


def test_token_for_deleted_account_is_401_not_500(client, app):
    """Previously this dereferenced None inside the handler and returned 500."""
    token, _ = make_user(client)
    user = User.query.filter_by(email='user@example.com').first()
    db.session.delete(user)
    db.session.commit()

    for path in ('/api/vault', '/api/auth/me', '/api/sharing/keys', '/api/emergency'):
        res = client.get(path, headers=auth_headers(token))
        assert res.status_code == 401, f'{path} returned {res.status_code}'


def test_forged_epoch_claim_is_rejected(client, app):
    """
    The epoch is inside the signed payload, so tampering invalidates the
    signature. Re-signing with the wrong key must also fail.
    """
    token, _ = make_user(client)
    payload = pyjwt.decode(token, options={'verify_signature': False})
    payload['epoch'] = 99
    forged = pyjwt.encode(payload, 'not-the-real-signing-key', algorithm='HS256')

    assert client.get('/api/vault', headers=auth_headers(forged)).status_code == 401


def test_tokens_predating_the_epoch_claim_still_work(client, app):
    """
    Deploying this must not sign existing sessions out: tokens minted before the
    claim existed decode with epoch 0, matching the column default.
    """
    token, _ = make_user(client)
    payload = pyjwt.decode(token, options={'verify_signature': False})
    del payload['epoch']                       # simulate a pre-upgrade token
    legacy = pyjwt.encode(payload, app.config['JWT_SECRET_KEY'], algorithm='HS256')

    assert client.get('/api/vault', headers=auth_headers(legacy)).status_code == 200


def test_logout_still_revokes_via_blacklist(client, app):
    """Epoch checking must not have displaced the existing jti blacklist."""
    token, refresh = make_user(client)
    assert client.post('/api/auth/logout', headers=auth_headers(token),
                       json={'refresh_token': refresh}).status_code == 200

    assert client.get('/api/vault', headers=auth_headers(token)).status_code == 401
    assert client.post('/api/auth/refresh',
                       json={'refresh_token': refresh}).status_code == 401