Change JWT encryption from HMAC to RSA with rotating keys

.devcontainer/Dockerfile

@@ -2,7 +2,7 @@ FROM alpine:latest
 
 # Install common tools
 RUN apk add --no-cache bash git \
-    python3 py3-pip
+    python3 py3-pip openssl
 
 # Setup default user
 ARG USERNAME=vscode

.gitignore

@@ -2,7 +2,8 @@
 logs/
 *.log
 
-# Ignore all pem files
+# Ignore JWT keys
+keys/
 *.pem
 
 # Upload folder

Dockerfile

@@ -4,7 +4,8 @@ FROM python:3.13-slim
 # Set a specific working directory in the container
 WORKDIR /app
 
-# Install dependencies separately for better caching
+# Install dependencies
+RUN apt update && apt install -y openssl
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
@@ -23,4 +24,4 @@ HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
     CMD curl --fail http://localhost:5000/ || exit 1
 
 # Command to run the app
-CMD ["python", "app.py"]
+ENTRYPOINT ["./scripts/entrypoint.sh"]

config/settings.py

@@ -20,7 +20,6 @@ class Config:
     MYSQL_CURSORCLASS = "DictCursor"
 
     # JWT configuration
-    JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY")
     JWT_ACCESS_TOKEN_EXPIRY = timedelta(hours=1)
     JWT_REFRESH_TOKEN_EXPIRY = timedelta(days=30)
 

entrypoint.sh

@@ -0,0 +1,20 @@
+#!/bin/sh
+
+set -e  # Exit if any command fails
+
+KEY_DIR="/app/keys"
+PRIVATE_KEY="$KEY_DIR/private_key.pem"
+PUBLIC_KEY="$KEY_DIR/public_key.pem"
+
+mkdir -p "$KEY_DIR"
+
+if [ ! -f "$PRIVATE_KEY" ] || [ ! -f "$PUBLIC_KEY" ]; then
+  echo "🔐 Generating RSA key pair..."
+  openssl genpkey -algorithm RSA -out "$PRIVATE_KEY" -pkeyopt rsa_keygen_bits:2048
+  openssl rsa -pubout -in "$PRIVATE_KEY" -out "$PUBLIC_KEY"
+else
+  echo "✅ Keys already exist. Skipping generation."
+fi
+
+# Start the API
+exec python3 app.py

jwtoken/decorators.py

@@ -0,0 +1,24 @@
+from functools import wraps
+
+from flask import jsonify, request
+
+from utility.jwtoken.common import extract_token_from_header
+
+from .exceptions import TokenError
+from .tokens import verify_token
+
+
+def token_required(f):
+    """Decorator to protect routes by requiring a valid token."""
+
+    @wraps(f)
+    def decorated(*args, **kwargs):
+        try:
+            token = extract_token_from_header()
+            decoded = verify_token(token, required_type="access")
+            request.person_id = decoded["person_id"]
+            return f(*args, **kwargs)
+        except TokenError as e:
+            return jsonify(message=e.message), e.status_code
+
+    return decorated

jwtoken/exceptions.py

@@ -0,0 +1,7 @@
+class TokenError(Exception):
+    """Custom exception for token-related errors."""
+
+    def __init__(self, message, status_code):
+        super().__init__(message)
+        self.status_code = status_code
+        self.message = message

jwtoken/tokens.py

@@ -0,0 +1,61 @@
+from datetime import datetime, timedelta, timezone
+
+import jwt
+
+from utility.jwtoken.common import get_active_kid, load_private_key, load_public_key
+
+from .exceptions import TokenError
+
+JWT_ACCESS_TOKEN_EXPIRY = timedelta(hours=1)
+JWT_REFRESH_TOKEN_EXPIRY = timedelta(days=30)
+
+
+def generate_access_token(person_id: int) -> str:
+    kid = get_active_kid()
+    private_key = load_private_key(kid)
+
+    payload = {
+        "person_id": person_id,
+        "exp": datetime.now(timezone.utc) + JWT_ACCESS_TOKEN_EXPIRY,  # Expiration
+        "iat": datetime.now(timezone.utc),  # Issued at
+        "token_type": "access",
+    }
+    headers = {"kid": kid}
+    return jwt.encode(payload, private_key, algorithm="RS256", headers=headers)
+
+
+def generate_refresh_token(person_id: int) -> str:
+    """Generate a long-lived refresh token for a user."""
+    kid = get_active_kid()
+    private_key = load_private_key(kid)
+
+    payload = {
+        "person_id": person_id,
+        "exp": datetime.now(timezone.utc) + JWT_REFRESH_TOKEN_EXPIRY,  # Expiration
+        "iat": datetime.now(timezone.utc),  # Issued at
+        "token_type": "refresh",
+    }
+    headers = {"kid": kid}
+
+    return jwt.encode(payload, private_key, algorithm="RS256", headers=headers)
+
+
+def verify_token(token: str, required_type: str) -> dict:
+    """Verify and decode a JWT token."""
+    try:
+        unverified_header = jwt.get_unverified_header(token)
+        kid = unverified_header.get("kid")
+        if not kid:
+            raise TokenError("KID missing in token header", 401)
+
+        public_key = load_public_key(kid)
+
+        decoded = jwt.decode(token, public_key, algorithms=["RS256"])
+        if decoded.get("token_type") != required_type:
+            raise jwt.InvalidTokenError("Invalid token type")
+        return decoded
+
+    except jwt.ExpiredSignatureError:
+        raise TokenError("Token has expired", 401)
+    except jwt.InvalidTokenError:
+        raise TokenError("Invalid token", 401)

routes/authentication.py

@@ -3,15 +3,11 @@
 from pymysql import MySQLError
 
 from config.ratelimit import limiter
-from jwt_helper import (
-    TokenError,
-    extract_token_from_header,
-    generate_access_token,
-    generate_refresh_token,
-    verify_token,
-)
+from jwtoken.exceptions import TokenError
+from jwtoken.tokens import generate_access_token, generate_refresh_token, verify_token
 from utility.database import database_cursor
 from utility.encryption import encrypt_email, hash_email, hash_password, verify_password
+from utility.jwtoken.common import extract_token_from_header
 from utility.validation import validate_email, validate_password
 
 authentication_blueprint = Blueprint("authentication", __name__)

routes/picture.py

@@ -3,7 +3,7 @@
 from flask import Blueprint, jsonify, request, send_from_directory
 
 from config.settings import Config
-from jwt_helper import token_required
+from jwtoken.decorators import token_required
 from utility.database import database_cursor
 
 picture_blueprint = Blueprint("picture", __name__)

tests/conftest.py

@@ -0,0 +1,7 @@
+import pytest
+
+
+@pytest.fixture
+def sample_person_id() -> int:
+    """Provide a sample person ID for testing"""
+    return 12345

tests/test_jwt/conftest.py → tests/test_jwtoken/conftest.py

@@ -1,7 +1,8 @@
+import jwt
 import pytest
 from flask import Flask
 
-from jwt_helper import generate_access_token
+from jwtoken.tokens import generate_access_token
 
 
 @pytest.fixture
@@ -27,3 +28,8 @@ def sample_token():
 def sample_access_token(sample_person_id):
     """Provide a sample access token for testing"""
     return generate_access_token(sample_person_id)
+
+
+@pytest.fixture
+def sample_kid(sample_access_token):
+    return jwt.get_unverified_header(sample_access_token).get("kid")

tests/test_jwt/test_extract_token_from_header.py → tests/test_jwtoken/test_extract_token_from_header.py

@@ -1,7 +1,8 @@
 import pytest
 from flask import Flask
 
-from jwt_helper import TokenError, extract_token_from_header
+from jwtoken.exceptions import TokenError
+from utility.jwtoken.common import extract_token_from_header
 
 app = Flask(__name__)