Aug 26 - Enhance security 2

app/services/auth_service.py

@@ -84,14 +84,23 @@ def decrypt_totp_secret(ciphertext_b64: str, iv_b64: str) -> str:
     return aesgcm.decrypt(iv, ciphertext, None).decode()
 
 
-def generate_tokens(user_id: int) -> dict:
-    """Return access_token and refresh_token JWTs, each with a unique jti."""
+def generate_tokens(user_id: int, token_epoch: int = 0) -> dict:
+    """
+    Return access_token and refresh_token JWTs, each with a unique jti.
+
+    token_epoch stamps the user's current session generation into both tokens.
+    require_jwt and /refresh compare it against users.token_epoch and reject on
+    mismatch, so incrementing that column revokes every outstanding token.
+    Always pass user.token_epoch — the 0 default exists only so old call sites
+    fail visibly in tests rather than silently minting unrevokable tokens.
+    """
     now = datetime.now(timezone.utc).replace(tzinfo=None)
     secret = current_app.config['JWT_SECRET_KEY']
     access_payload = {
         'sub': str(user_id),
         'type': 'access',
         'jti': str(uuid.uuid4()),
+        'epoch': int(token_epoch or 0),
         'iat': now,
         'exp': now + current_app.config['JWT_ACCESS_TOKEN_EXPIRES'],
     }
@@ -99,6 +108,7 @@ def generate_tokens(user_id: int) -> dict:
         'sub': str(user_id),
         'type': 'refresh',
         'jti': str(uuid.uuid4()),
+        'epoch': int(token_epoch or 0),
         'iat': now,
         'exp': now + current_app.config['JWT_REFRESH_TOKEN_EXPIRES'],
     }
@@ -172,8 +182,45 @@ def blacklist_token(token: str, token_type: str) -> None:
         pass  # Never let blacklisting errors break the logout flow
 
 
+def load_user_for_token(payload) -> 'object | None':
+    """
+    Resolve the User a validated token refers to, or None if the token must be
+    rejected.
+
+    Two checks beyond signature validity:
+
+    1. The user still exists. Routes immediately dereference the result of
+       db.session.get(User, ...); without this a valid token for a deleted
+       account produced an AttributeError on None and a 500.
+    2. The token's epoch claim still matches users.token_epoch. A master-password
+       change increments that column, which revokes every token minted before it.
+       Tokens issued before the claim existed decode as 0 and match the column
+       default, so an upgrade does not sign existing sessions out.
+    """
+    from app.models.user import User
+    from app import db
+
+    try:
+        user_id = int(payload['sub'])
+    except (KeyError, TypeError, ValueError):
+        return None
+
+    user = db.session.get(User, user_id)
+    if user is None:
+        return None
+    if int(payload.get('epoch', 0) or 0) != int(user.token_epoch or 0):
+        return None
+    return user
+
+
 def require_jwt(f):
-    """Decorator: validates Bearer token and sets g.current_user_id."""
+    """
+    Decorator: validates the Bearer token and sets g.current_user_id.
+
+    Also sets g.current_user to the resolved User so handlers can reuse it
+    instead of issuing a second lookup (SQLAlchemy's identity map makes the
+    repeat cheap, but reusing it is clearer).
+    """
     @wraps(f)
     def decorated(*args, **kwargs):
         auth_header = request.headers.get('Authorization', '')
@@ -186,7 +233,14 @@ def require_jwt(f):
             return jsonify({'error': 'Token expired'}), 401
         except jwt.PyJWTError:
             return jsonify({'error': 'Invalid token'}), 401
-        g.current_user_id = int(payload['sub'])
+
+        user = load_user_for_token(payload)
+        if user is None:
+            # Deleted account, or a token predating a credential change.
+            return jsonify({'error': 'Session is no longer valid. Please log in again.'}), 401
+
+        g.current_user = user
+        g.current_user_id = user.id
         return f(*args, **kwargs)
     return decorated