124 lines
4.5 KiB
Python
124 lines
4.5 KiB
Python
import uuid
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from functools import wraps
|
|
|
|
import jwt
|
|
from flask import current_app, request, g, jsonify
|
|
from argon2 import PasswordHasher
|
|
from argon2.exceptions import VerifyMismatchError, VerificationError, InvalidHashError
|
|
|
|
|
|
def hash_auth_token(auth_hash: str) -> str:
|
|
"""Hash the client-derived PBKDF2 auth_hash with Argon2id before storing."""
|
|
ph = PasswordHasher(
|
|
time_cost=current_app.config['ARGON2_TIME_COST'],
|
|
memory_cost=current_app.config['ARGON2_MEMORY_COST'],
|
|
parallelism=current_app.config['ARGON2_PARALLELISM'],
|
|
)
|
|
return ph.hash(auth_hash)
|
|
|
|
|
|
def verify_auth_token(auth_hash: str, stored_hash: str) -> bool:
|
|
ph = PasswordHasher()
|
|
try:
|
|
return ph.verify(stored_hash, auth_hash)
|
|
except (VerifyMismatchError, VerificationError, InvalidHashError):
|
|
return False
|
|
|
|
|
|
def generate_tokens(user_id: int) -> dict:
|
|
"""Return access_token and refresh_token JWTs, each with a unique jti."""
|
|
now = datetime.utcnow()
|
|
secret = current_app.config['JWT_SECRET_KEY']
|
|
access_payload = {
|
|
'sub': str(user_id),
|
|
'type': 'access',
|
|
'jti': str(uuid.uuid4()),
|
|
'iat': now,
|
|
'exp': now + current_app.config['JWT_ACCESS_TOKEN_EXPIRES'],
|
|
}
|
|
refresh_payload = {
|
|
'sub': str(user_id),
|
|
'type': 'refresh',
|
|
'jti': str(uuid.uuid4()),
|
|
'iat': now,
|
|
'exp': now + current_app.config['JWT_REFRESH_TOKEN_EXPIRES'],
|
|
}
|
|
return {
|
|
'access_token': jwt.encode(access_payload, secret, algorithm='HS256'),
|
|
'refresh_token': jwt.encode(refresh_payload, secret, algorithm='HS256'),
|
|
}
|
|
|
|
|
|
def generate_mfa_token(user_id: int) -> str:
|
|
"""Short-lived (5-min) single-use token issued after password but before TOTP."""
|
|
now = datetime.utcnow()
|
|
payload = {
|
|
'sub': str(user_id),
|
|
'type': 'mfa',
|
|
'jti': str(uuid.uuid4()),
|
|
'iat': now,
|
|
'exp': now + timedelta(minutes=5),
|
|
}
|
|
return jwt.encode(payload, current_app.config['JWT_SECRET_KEY'], algorithm='HS256')
|
|
|
|
|
|
def decode_token(token: str, expected_type: str = 'access', check_blacklist: bool = True) -> dict:
|
|
"""Decode and validate a JWT. Raises jwt.PyJWTError on any failure."""
|
|
secret = current_app.config['JWT_SECRET_KEY']
|
|
payload = jwt.decode(token, secret, algorithms=['HS256'])
|
|
if payload.get('type') != expected_type:
|
|
raise jwt.InvalidTokenError('Wrong token type')
|
|
if check_blacklist:
|
|
from app.models.token_blacklist import TokenBlacklist
|
|
jti = payload.get('jti')
|
|
if jti and TokenBlacklist.is_blacklisted(jti):
|
|
raise jwt.InvalidTokenError('Token has been revoked')
|
|
return payload
|
|
|
|
|
|
def blacklist_token(token: str, token_type: str) -> None:
|
|
"""Add a JWT's jti to the blacklist. Silently ignores invalid tokens."""
|
|
try:
|
|
payload = decode_token(token, expected_type=token_type, check_blacklist=False)
|
|
jti = payload.get('jti')
|
|
if not jti:
|
|
return
|
|
exp = payload.get('exp')
|
|
expires_at = datetime.utcfromtimestamp(exp) if exp else datetime.utcnow() + timedelta(days=7)
|
|
from app.models.token_blacklist import TokenBlacklist
|
|
from app import db
|
|
# Avoid duplicate if already blacklisted
|
|
if not TokenBlacklist.query.filter_by(jti=jti).first():
|
|
entry = TokenBlacklist(
|
|
jti=jti,
|
|
user_id=int(payload.get('sub', 0)),
|
|
expires_at=expires_at,
|
|
)
|
|
db.session.add(entry)
|
|
db.session.commit()
|
|
# Opportunistic cleanup — runs in same transaction context
|
|
TokenBlacklist.cleanup_expired()
|
|
except Exception:
|
|
pass # Never let blacklisting errors break the logout flow
|
|
|
|
|
|
def require_jwt(f):
|
|
"""Decorator: validates Bearer token and sets g.current_user_id."""
|
|
@wraps(f)
|
|
def decorated(*args, **kwargs):
|
|
auth_header = request.headers.get('Authorization', '')
|
|
if not auth_header.startswith('Bearer '):
|
|
return jsonify({'error': 'Missing or invalid Authorization header'}), 401
|
|
token = auth_header[7:]
|
|
try:
|
|
payload = decode_token(token, expected_type='access')
|
|
except jwt.ExpiredSignatureError:
|
|
return jsonify({'error': 'Token expired'}), 401
|
|
except jwt.PyJWTError:
|
|
return jsonify({'error': 'Invalid token'}), 401
|
|
g.current_user_id = int(payload['sub'])
|
|
return f(*args, **kwargs)
|
|
return decorated
|