04/16 Upload codebase
This commit is contained in:
@@ -0,0 +1,241 @@
|
||||
import re
|
||||
import time
|
||||
|
||||
from flask import Blueprint, request, jsonify, g
|
||||
from app import db, limiter
|
||||
from app.models.user import User
|
||||
from app.services.auth_service import (
|
||||
hash_auth_token,
|
||||
verify_auth_token,
|
||||
generate_tokens,
|
||||
generate_mfa_token,
|
||||
decode_token,
|
||||
blacklist_token,
|
||||
require_jwt,
|
||||
)
|
||||
|
||||
auth_bp = Blueprint('auth', __name__)
|
||||
|
||||
EMAIL_RE = re.compile(r'^[^@\s]+@[^@\s]+\.[^@\s]+$')
|
||||
|
||||
|
||||
@auth_bp.route('/register', methods=['POST'])
|
||||
@limiter.limit('10 per minute')
|
||||
def register():
|
||||
data = request.get_json(silent=True) or {}
|
||||
email = (data.get('email') or '').strip().lower()
|
||||
auth_hash = data.get('auth_hash', '')
|
||||
enc_key_salt = data.get('enc_key_salt', '')
|
||||
|
||||
if not email or not EMAIL_RE.match(email):
|
||||
return jsonify({'error': 'Invalid email address'}), 400
|
||||
if not auth_hash:
|
||||
return jsonify({'error': 'auth_hash is required'}), 400
|
||||
if not enc_key_salt:
|
||||
return jsonify({'error': 'enc_key_salt is required'}), 400
|
||||
|
||||
if User.query.filter_by(email=email).first():
|
||||
return jsonify({'error': 'Email already registered'}), 409
|
||||
|
||||
master_hash = hash_auth_token(auth_hash)
|
||||
user = User(email=email, master_hash=master_hash, enc_key_salt=enc_key_salt)
|
||||
db.session.add(user)
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({'message': 'Account created successfully'}), 201
|
||||
|
||||
|
||||
@auth_bp.route('/login', methods=['POST'])
|
||||
@limiter.limit('10 per minute')
|
||||
def login():
|
||||
data = request.get_json(silent=True) or {}
|
||||
email = (data.get('email') or '').strip().lower()
|
||||
auth_hash = data.get('auth_hash', '')
|
||||
|
||||
time.sleep(0.1) # mitigate timing-based user enumeration
|
||||
|
||||
if not email or not auth_hash:
|
||||
return jsonify({'error': 'Email and auth_hash are required'}), 400
|
||||
|
||||
user = User.query.filter_by(email=email).first()
|
||||
if not user or not verify_auth_token(auth_hash, user.master_hash):
|
||||
return jsonify({'error': 'Invalid email or password'}), 401
|
||||
|
||||
from datetime import datetime
|
||||
user.last_login = datetime.utcnow()
|
||||
db.session.commit()
|
||||
|
||||
# MFA gate: if enabled, issue a short-lived mfa_token instead of full tokens
|
||||
if user.totp_enabled:
|
||||
mfa_token = generate_mfa_token(user.id)
|
||||
return jsonify({
|
||||
'mfa_required': True,
|
||||
'mfa_token': mfa_token,
|
||||
'enc_key_salt': user.enc_key_salt,
|
||||
}), 200
|
||||
|
||||
tokens = generate_tokens(user.id)
|
||||
return jsonify({
|
||||
'access_token': tokens['access_token'],
|
||||
'refresh_token': tokens['refresh_token'],
|
||||
'enc_key_salt': user.enc_key_salt,
|
||||
}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/logout', methods=['POST'])
|
||||
def logout():
|
||||
"""Blacklist both the access token (from header) and refresh token (from body)."""
|
||||
auth_header = request.headers.get('Authorization', '')
|
||||
if auth_header.startswith('Bearer '):
|
||||
blacklist_token(auth_header[7:], 'access')
|
||||
|
||||
data = request.get_json(silent=True) or {}
|
||||
refresh_token = data.get('refresh_token', '')
|
||||
if refresh_token:
|
||||
blacklist_token(refresh_token, 'refresh')
|
||||
|
||||
return jsonify({'message': 'Logged out'}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/refresh', methods=['POST'])
|
||||
@limiter.limit('30 per minute')
|
||||
def refresh():
|
||||
data = request.get_json(silent=True) or {}
|
||||
refresh_token = data.get('refresh_token', '')
|
||||
if not refresh_token:
|
||||
return jsonify({'error': 'refresh_token is required'}), 400
|
||||
|
||||
try:
|
||||
payload = decode_token(refresh_token, expected_type='refresh')
|
||||
except Exception:
|
||||
return jsonify({'error': 'Invalid or expired refresh token'}), 401
|
||||
|
||||
# Rotate: blacklist old refresh token and issue fresh pair
|
||||
blacklist_token(refresh_token, 'refresh')
|
||||
tokens = generate_tokens(int(payload['sub']))
|
||||
return jsonify({
|
||||
'access_token': tokens['access_token'],
|
||||
'refresh_token': tokens['refresh_token'],
|
||||
}), 200
|
||||
|
||||
|
||||
# ── MFA / TOTP endpoints ─────────────────────────────────────────────────────
|
||||
|
||||
@auth_bp.route('/mfa/setup', methods=['GET'])
|
||||
@require_jwt
|
||||
def mfa_setup():
|
||||
"""Generate a new TOTP secret and return QR code (as base64 PNG data URI)."""
|
||||
user = User.query.get(g.current_user_id)
|
||||
if user.totp_enabled:
|
||||
return jsonify({'error': 'MFA is already enabled'}), 400
|
||||
|
||||
import pyotp
|
||||
import qrcode
|
||||
import io
|
||||
import base64
|
||||
|
||||
secret = pyotp.random_base32()
|
||||
uri = pyotp.TOTP(secret).provisioning_uri(
|
||||
name=user.email,
|
||||
issuer_name='PassKeeper',
|
||||
)
|
||||
|
||||
img = qrcode.make(uri)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format='PNG')
|
||||
qr_b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
return jsonify({
|
||||
'secret': secret,
|
||||
'qr_code': f'data:image/png;base64,{qr_b64}',
|
||||
'uri': uri,
|
||||
}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/mfa/enable', methods=['POST'])
|
||||
@require_jwt
|
||||
def mfa_enable():
|
||||
"""Enable MFA after verifying the first TOTP code."""
|
||||
user = User.query.get(g.current_user_id)
|
||||
if user.totp_enabled:
|
||||
return jsonify({'error': 'MFA is already enabled'}), 400
|
||||
|
||||
data = request.get_json(silent=True) or {}
|
||||
secret = (data.get('secret') or '').strip()
|
||||
totp_code = (data.get('totp_code') or '').strip()
|
||||
|
||||
if not secret or not totp_code:
|
||||
return jsonify({'error': 'secret and totp_code are required'}), 400
|
||||
|
||||
import pyotp
|
||||
if not pyotp.TOTP(secret).verify(totp_code, valid_window=1):
|
||||
return jsonify({'error': 'Invalid verification code'}), 400
|
||||
|
||||
user.totp_secret = secret
|
||||
user.totp_enabled = True
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({'message': 'MFA enabled successfully'}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/mfa/disable', methods=['POST'])
|
||||
@require_jwt
|
||||
def mfa_disable():
|
||||
"""Disable MFA after verifying the current TOTP code."""
|
||||
user = User.query.get(g.current_user_id)
|
||||
if not user.totp_enabled:
|
||||
return jsonify({'error': 'MFA is not enabled'}), 400
|
||||
|
||||
data = request.get_json(silent=True) or {}
|
||||
totp_code = (data.get('totp_code') or '').strip()
|
||||
|
||||
import pyotp
|
||||
if not pyotp.TOTP(user.totp_secret).verify(totp_code, valid_window=1):
|
||||
return jsonify({'error': 'Invalid verification code'}), 400
|
||||
|
||||
user.totp_secret = None
|
||||
user.totp_enabled = False
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({'message': 'MFA disabled'}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/mfa/verify', methods=['POST'])
|
||||
@limiter.limit('10 per minute')
|
||||
def mfa_verify():
|
||||
"""Complete MFA login: verify TOTP code and exchange mfa_token for real tokens."""
|
||||
data = request.get_json(silent=True) or {}
|
||||
mfa_token = data.get('mfa_token', '')
|
||||
totp_code = (data.get('totp_code') or '').strip()
|
||||
|
||||
if not mfa_token or not totp_code:
|
||||
return jsonify({'error': 'mfa_token and totp_code are required'}), 400
|
||||
|
||||
try:
|
||||
payload = decode_token(mfa_token, expected_type='mfa', check_blacklist=True)
|
||||
except Exception:
|
||||
return jsonify({'error': 'Invalid or expired MFA token'}), 401
|
||||
|
||||
user = User.query.get(int(payload['sub']))
|
||||
if not user or not user.totp_enabled:
|
||||
return jsonify({'error': 'MFA not configured for this account'}), 400
|
||||
|
||||
import pyotp
|
||||
if not pyotp.TOTP(user.totp_secret).verify(totp_code, valid_window=1):
|
||||
return jsonify({'error': 'Invalid verification code'}), 400
|
||||
|
||||
# One-time use: blacklist the mfa_token
|
||||
blacklist_token(mfa_token, 'mfa')
|
||||
|
||||
tokens = generate_tokens(user.id)
|
||||
return jsonify({
|
||||
'access_token': tokens['access_token'],
|
||||
'refresh_token': tokens['refresh_token'],
|
||||
}), 200
|
||||
|
||||
|
||||
@auth_bp.route('/mfa/status', methods=['GET'])
|
||||
@require_jwt
|
||||
def mfa_status():
|
||||
user = User.query.get(g.current_user_id)
|
||||
return jsonify({'totp_enabled': user.totp_enabled}), 200
|
||||
Reference in New Issue
Block a user