Secure REST APIs with JWT Authentication: Building a Complete Authorization System

"""
JWT token encoding and decoding with security best practices.
"""

import jwt
import json
from datetime import datetime, timedelta, timezone
from typing import Dict, Optional
from functools import wraps
from flask import request, jsonify, current_app

class JWTHandler:
    """Handle JWT token creation and validation."""
    
    @staticmethod
    def create_token(
        user_id: int,
        email: str,
        role: str = 'user',
        expires_in: int = 3600
    ) -> str:
        """
        Create JWT token.
        
        Args:
            user_id: User ID to embed in token
            email: User email
            role: User role (user, admin, etc.)
            expires_in: Expiration time in seconds (default: 1 hour)
        
        Returns:
            Encoded JWT token
        """
        
        now = datetime.now(timezone.utc)
        expires = now + timedelta(seconds=expires_in)
        
        payload = {
            'user_id': user_id,
            'email': email,
            'role': role,
            'iat': int(now.timestamp()),  # Issued at
            'exp': int(expires.timestamp()),  # Expiration
            'type': 'access'  # Token type
        }
        
        token = jwt.encode(
            payload,
            current_app.config['JWT_SECRET_KEY'],
            algorithm='HS256'
        )
        
        return token
    
    @staticmethod
    def create_refresh_token(user_id: int, email: str) -> str:
        """
        Create refresh token (longer expiration).
        Valid for 30 days.
        """
        
        now = datetime.now(timezone.utc)
        expires = now + timedelta(days=30)
        
        payload = {
            'user_id': user_id,
            'email': email,
            'iat': int(now.timestamp()),
            'exp': int(expires.timestamp()),
            'type': 'refresh'
        }
        
        token = jwt.encode(
            payload,
            current_app.config['JWT_REFRESH_SECRET_KEY'],
            algorithm='HS256'
        )
        
        return token
    
    @staticmethod
    def verify_token(token: str) -> Optional[Dict]:
        """
        Verify and decode token.
        
        Returns:
            Decoded payload if valid
            None if invalid or expired
        """
        
        try:
            payload = jwt.decode(
                token,
                current_app.config['JWT_SECRET_KEY'],
                algorithms=['HS256']
            )
            
            # Verify token type
            if payload.get('type') != 'access':
                return None
            
            return payload
        
        except jwt.ExpiredSignatureError:
            return None  # Token expired
        except jwt.InvalidTokenError:
            return None  # Invalid signature
    
    @staticmethod
    def verify_refresh_token(token: str) -> Optional[Dict]:
        """Verify refresh token."""
        
        try:
            payload = jwt.decode(
                token,
                current_app.config['JWT_REFRESH_SECRET_KEY'],
                algorithms=['HS256']
            )
            
            if payload.get('type') != 'refresh':
                return None
            
            return payload
        
        except jwt.ExpiredSignatureError:
            return None
        except jwt.InvalidTokenError:
            return None
    
    @staticmethod
    def extract_token_from_request() -> Optional[str]:
        """
        Extract JWT from Authorization header.
        Expected format: Authorization: Bearer <token>
        """
        
        auth_header = request.headers.get('Authorization')
        if not auth_header:
            return None
        
        parts = auth_header.split()
        if len(parts) != 2 or parts[0].lower() != 'bearer':
            return None
        
        return parts[1]

"""
Authentication endpoints: register, login, refresh.
"""

from flask import Blueprint, request, jsonify
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash
import re

auth_bp = Blueprint('auth', __name__, url_prefix='/api/auth')
db = SQLAlchemy()

# ============= Models =============

class User(db.Model):
    """User model with security best practices."""
    
    __tablename__ = 'users'
    
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(255), unique=True, nullable=False, index=True)
    password_hash = db.Column(db.String(255), nullable=False)
    first_name = db.Column(db.String(100))
    last_name = db.Column(db.String(100))
    role = db.Column(db.String(20), default='user')  # user, admin
    is_active = db.Column(db.Boolean, default=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Relationships
    todos = db.relationship('Todo', backref='user', lazy='dynamic', cascade='all, delete-orphan')
    
    def set_password(self, password: str):
        """Hash and store password."""
        if not self._validate_password_strength(password):
            raise ValueError("Password doesn't meet security requirements")
        self.password_hash = generate_password_hash(password, method='pbkdf2:sha256')
    
    def check_password(self, password: str) -> bool:
        """Verify password."""
        return check_password_hash(self.password_hash, password)
    
    @staticmethod
    def _validate_password_strength(password: str) -> bool:
        """Ensure password meets security standards."""
        if len(password) < 8:
            return False
        if not re.search(r'[A-Z]', password):
            return False
        if not re.search(r'[a-z]', password):
            return False
        if not re.search(r'[0-9]', password):
            return False
        return True

# ============= Routes =============

@auth_bp.route('/register', methods=['POST'])
def register():
    """
    Register new user.
    
    Request body:
    {
        "email": "user@example.com",
        "password": "SecurePass123",
        "first_name": "John",
        "last_name": "Doe"
    }
    """
    
    try:
        data = request.get_json()
        
        # Validate input
        if not data:
            return {'error': 'Missing request body'}, 400
        
        email = data.get('email', '').lower().strip()
        password = data.get('password')
        first_name = data.get('first_name', '').strip()
        last_name = data.get('last_name', '').strip()
        
        # Validation
        errors = []
        
        if not email or '@' not in email:
            errors.append('Invalid email address')
        
        if not password or len(password) < 8:
            errors.append('Password must be at least 8 characters')
        
        if User.query.filter_by(email=email).first():
            errors.append('Email already registered')
        
        if errors:
            return {'errors': errors}, 422
        
        # Create user
        user = User(
            email=email,
            first_name=first_name,
            last_name=last_name
        )
        user.set_password(password)
        
        db.session.add(user)
        db.session.commit()
        
        return {
            'message': 'Registration successful',
            'user_id': user.id,
            'email': user.email
        }, 201
    
    except ValueError as e:
        return {'error': str(e)}, 422
    except Exception as e:
        db.session.rollback()
        return {'error': 'Registration failed'}, 500


@auth_bp.route('/login', methods=['POST'])
def login():
    """
    Login and get JWT tokens.
    
    Request body:
    {
        "email": "user@example.com",
        "password": "SecurePass123"
    }
    
    Response:
    {
        "access_token": "eyJ...",
        "refresh_token": "eyJ...",
        "user": {
            "id": 1,
            "email": "user@example.com",
            "first_name": "John"
        }
    }
    """
    
    try:
        data = request.get_json()
        
        if not data:
            return {'error': 'Missing credentials'}, 400
        
        email = data.get('email', '').lower().strip()
        password = data.get('password')
        
        if not email or not password:
            return {'error': 'Email and password required'}, 422
        
        # Find user
        user = User.query.filter_by(email=email).first()
        
        if not user or not user.check_password(password):
            # Don't reveal which field is wrong (security best practice)
            return {'error': 'Invalid credentials'}, 401
        
        if not user.is_active:
            return {'error': 'Account is inactive'}, 403
        
        # Create tokens
        access_token = JWTHandler.create_token(
            user_id=user.id,
            email=user.email,
            role=user.role,
            expires_in=3600  # 1 hour
        )
        
        refresh_token = JWTHandler.create_refresh_token(
            user_id=user.id,
            email=user.email
        )
        
        return {
            'access_token': access_token,
            'refresh_token': refresh_token,
            'user': {
                'id': user.id,
                'email': user.email,
                'first_name': user.first_name,
                'role': user.role
            }
        }, 200
    
    except Exception as e:
        return {'error': 'Login failed'}, 500


@auth_bp.route('/refresh', methods=['POST'])
def refresh():
    """
    Refresh access token using refresh token.
    
    Request body:
    {
        "refresh_token": "eyJ..."
    }
    """
    
    try:
        data = request.get_json()
        
        if not data or 'refresh_token' not in data:
            return {'error': 'Refresh token required'}, 422
        
        refresh_token = data.get('refresh_token')
        
        # Verify refresh token
        payload = JWTHandler.verify_refresh_token(refresh_token)
        
        if not payload:
            return {'error': 'Invalid or expired refresh token'}, 401
        
        # Create new access token
        user_id = payload.get('user_id')
        user = User.query.get(user_id)
        
        if not user or not user.is_active:
            return {'error': 'User not found or inactive'}, 404
        
        new_access_token = JWTHandler.create_token(
            user_id=user.id,
            email=user.email,
            role=user.role,
            expires_in=3600
        )
        
        return {
            'access_token': new_access_token
        }, 200
    
    except Exception as e:
        return {'error': 'Token refresh failed'}, 500


@auth_bp.route('/logout', methods=['POST'])
def logout():
    """
    Logout (client-side token deletion).
    With JWT, logout is typically just deleting the token client-side.
    Optionally, blacklist token server-side.
    """
    
    try:
        token = JWTHandler.extract_token_from_request()
        
        if not token:
            return {'error': 'No token provided'}, 422
        
        # Optional: Blacklist token (prevents reuse)
        # token_blacklist = redis.get(f'blacklist:{token}')
        # redis.set(f'blacklist:{token}', '1', ex=3600)
        
        return {'message': 'Logged out successfully'}, 200
    
    except Exception as e:
        return {'error': 'Logout failed'}, 500

"""
Decorators to protect routes and check permissions.
"""

from functools import wraps
from flask import current_app

def login_required(f):
    """Decorator to require valid JWT token."""
    
    @wraps(f)
    def decorated_function(*args, **kwargs):
        token = JWTHandler.extract_token_from_request()
        
        if not token:
            return {'error': 'Missing authorization token'}, 401
        
        payload = JWTHandler.verify_token(token)
        
        if not payload:
            return {'error': 'Invalid or expired token'}, 401
        
        # Add user info to request context
        from flask import g
        g.user_id = payload.get('user_id')
        g.user_email = payload.get('email')
        g.user_role = payload.get('role')
        
        return f(*args, **kwargs)
    
    return decorated_function


def admin_required(f):
    """Decorator to require admin role."""
    
    @wraps(f)
    def decorated_function(*args, **kwargs):
        token = JWTHandler.extract_token_from_request()
        
        if not token:
            return {'error': 'Missing authorization token'}, 401
        
        payload = JWTHandler.verify_token(token)
        
        if not payload:
            return {'error': 'Invalid or expired token'}, 401
        
        if payload.get('role') != 'admin':
            return {'error': 'Admin access required'}, 403
        
        from flask import g
        g.user_id = payload.get('user_id')
        g.user_email = payload.get('email')
        g.user_role = payload.get('role')
        
        return f(*args, **kwargs)
    
    return decorated_function

# ============= Usage in routes =============

@auth_bp.route('/profile', methods=['GET'])
@login_required
def get_profile():
    """Get current user profile."""
    
    from flask import g
    user = User.query.get(g.user_id)
    
    return {
        'user': {
            'id': user.id,
            'email': user.email,
            'first_name': user.first_name,
            'last_name': user.last_name,
            'role': user.role
        }
    }, 200


@auth_bp.route('/profile', methods=['PUT'])
@login_required
def update_profile():
    """Update user profile."""
    
    from flask import g
    data = request.get_json()
    
    user = User.query.get(g.user_id)
    user.first_name = data.get('first_name', user.first_name)
    user.last_name = data.get('last_name', user.last_name)
    
    db.session.commit()
    
    return {'message': 'Profile updated'}, 200


@auth_bp.route('/users', methods=['GET'])
@admin_required
def list_users():
    """Admin endpoint to list all users."""
    
    users = User.query.all()
    
    return {
        'users': [
            {
                'id': u.id,
                'email': u.email,
                'role': u.role,
                'created_at': u.created_at.isoformat()
            } for u in users
        ]
    }, 200

"""
Todo CRUD operations - all protected by JWT.
"""

class Todo(db.Model):
    """Todo model."""
    
    __tablename__ = 'todos'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False, index=True)
    title = db.Column(db.String(255), nullable=False)
    description = db.Column(db.Text)
    completed = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

todos_bp = Blueprint('todos', __name__, url_prefix='/api/todos')

@todos_bp.route('', methods=['GET'])
@login_required
def list_todos():
    """Get all todos for current user."""
    
    from flask import g
    
    page = request.args.get('page', 1, type=int)
    per_page = request.args.get('per_page', 20, type=int)
    completed = request.args.get('completed', None)
    
    query = Todo.query.filter_by(user_id=g.user_id)
    
    if completed is not None:
        query = query.filter_by(completed=completed.lower() == 'true')
    
    todos = query.paginate(page=page, per_page=per_page)
    
    return {
        'todos': [
            {
                'id': t.id,
                'title': t.title,
                'description': t.description,
                'completed': t.completed,
                'created_at': t.created_at.isoformat()
            } for t in todos.items
        ],
        'pagination': {
            'page': page,
            'per_page': per_page,
            'total': todos.total,
            'pages': todos.pages
        }
    }, 200


@todos_bp.route('', methods=['POST'])
@login_required
def create_todo():
    """Create new todo."""
    
    from flask import g
    data = request.get_json()
    
    if not data or 'title' not in data:
        return {'error': 'Title is required'}, 422
    
    todo = Todo(
        user_id=g.user_id,
        title=data.get('title'),
        description=data.get('description', '')
    )
    
    db.session.add(todo)
    db.session.commit()
    
    return {
        'id': todo.id,
        'title': todo.title,
        'completed': todo.completed
    }, 201


@todos_bp.route('/<int:todo_id>', methods=['GET'])
@login_required
def get_todo(todo_id):
    """Get specific todo."""
    
    from flask import g
    
    todo = Todo.query.filter_by(
        id=todo_id,
        user_id=g.user_id
    ).first()
    
    if not todo:
        return {'error': 'Todo not found'}, 404
    
    return {
        'id': todo.id,
        'title': todo.title,
        'description': todo.description,
        'completed': todo.completed,
        'created_at': todo.created_at.isoformat()
    }, 200


@todos_bp.route('/<int:todo_id>', methods=['PUT'])
@login_required
def update_todo(todo_id):
    """Update todo."""
    
    from flask import g
    data = request.get_json()
    
    todo = Todo.query.filter_by(
        id=todo_id,
        user_id=g.user_id
    ).first()
    
    if not todo:
        return {'error': 'Todo not found'}, 404
    
    todo.title = data.get('title', todo.title)
    todo.description = data.get('description', todo.description)
    todo.completed = data.get('completed', todo.completed)
    
    db.session.commit()
    
    return {'message': 'Todo updated'}, 200


@todos_bp.route('/<int:todo_id>', methods=['DELETE'])
@login_required
def delete_todo(todo_id):
    """Delete todo."""
    
    from flask import g
    
    todo = Todo.query.filter_by(
        id=todo_id,
        user_id=g.user_id
    ).first()
    
    if not todo:
        return {'error': 'Todo not found'}, 404
    
    db.session.delete(todo)
    db.session.commit()
    
    return {'message': 'Todo deleted'}, 200

# config.py
import os

class Config:
    """Flask configuration."""
    
    # Database
    SQLALCHEMY_DATABASE_URI = os.environ.get(
        'DATABASE_URL',
        'postgresql://user:password@localhost/flask_jwt_todo'
    )
    SQLALCHEMY_TRACK_MODIFICATIONS = False
    
    # JWT
    JWT_SECRET_KEY = os.environ.get('JWT_SECRET_KEY', 'dev-secret-key-change-in-production')
    JWT_REFRESH_SECRET_KEY = os.environ.get('JWT_REFRESH_SECRET_KEY', 'dev-refresh-secret-key')
    JWT_ALGORITHM = 'HS256'
    
    # Security
    SESSION_COOKIE_SECURE = True
    SESSION_COOKIE_HTTPONLY = True
    PERMANENT_SESSION_LIFETIME = 3600

# test_auth.py
import pytest
from app import create_app, db
from app.models import User

@pytest.fixture
def client():
    """Test client."""
    app = create_app('testing')
    with app.app_context():
        db.create_all()
        yield app.test_client()
        db.session.remove()
        db.drop_all()

def test_register(client):
    """Test user registration."""
    response = client.post('/api/auth/register', json={
        'email': 'test@example.com',
        'password': 'SecurePass123',
        'first_name': 'Test',
        'last_name': 'User'
    })
    
    assert response.status_code == 201
    assert response.json['email'] == 'test@example.com'

def test_login(client):
    """Test user login."""
    # Register first
    client.post('/api/auth/register', json={
        'email': 'test@example.com',
        'password': 'SecurePass123'
    })
    
    # Login
    response = client.post('/api/auth/login', json={
        'email': 'test@example.com',
        'password': 'SecurePass123'
    })
    
    assert response.status_code == 200
    assert 'access_token' in response.json
    assert 'refresh_token' in response.json

def test_protected_route(client):
    """Test protected route requires token."""
    response = client.get('/api/todos')
    assert response.status_code == 401

def test_protected_route_with_token(client):
    """Test protected route with valid token."""
    # Register and login
    client.post('/api/auth/register', json={
        'email': 'test@example.com',
        'password': 'SecurePass123'
    })
    
    login_response = client.post('/api/auth/login', json={
        'email': 'test@example.com',
        'password': 'SecurePass123'
    })
    
    token = login_response.json['access_token']
    
    # Access protected route
    response = client.get(
        '/api/todos',
        headers={'Authorization': f'Bearer {token}'}
    )
    
    assert response.status_code == 200

# ❌ BAD: Storing secrets in code
JWT_SECRET_KEY = "hardcoded-secret"

# ✅ GOOD: Use environment variables
JWT_SECRET_KEY = os.environ.get('JWT_SECRET_KEY')

# ❌ BAD: Storing passwords in plaintext
user.password = request.json['password']

# ✅ GOOD: Hash before storing
user.set_password(request.json['password'])

# ❌ BAD: Exposing sensitive data in response
{
    'password_hash': user.password_hash,
    'internal_id': user.internal_id
}

# ✅ GOOD: Return only necessary fields
{
    'user_id': user.id,
    'email': user.email
}