Token usage security updates

Author: azvoleff

gefapi/__init__.py

@@ -763,6 +763,58 @@ def api_docs():
 
 jwt = JWTManager(app)
 
+# Token blocklist storage for revoked access tokens
+# Uses Redis for production scalability, falls back to in-memory for testing
+_revoked_tokens = set()
+
+
+def get_revoked_tokens_storage():
+    """Get the revoked tokens storage (Redis or in-memory fallback)."""
+    try:
+        import redis
+
+        redis_url = SETTINGS.get("CELERY_BROKER_URL")
+        if redis_url:
+            return redis.from_url(redis_url)
+    except Exception as e:
+        logger.debug(f"Redis not available for token blocklist, using in-memory: {e}")
+    return None
+
+
+def add_token_to_blocklist(jti: str, expires_in_seconds: int = 3600) -> None:
+    """Add a token JTI to the blocklist."""
+    redis_client = get_revoked_tokens_storage()
+    if redis_client:
+        try:
+            # Store in Redis with expiration matching token lifetime
+            redis_client.setex(f"blocklist:{jti}", expires_in_seconds, "revoked")
+            return
+        except Exception as e:
+            logger.warning(f"Failed to add token to Redis blocklist: {e}")
+    # Fallback to in-memory (note: not shared across workers)
+    _revoked_tokens.add(jti)
+
+
+def is_token_in_blocklist(jti: str) -> bool:
+    """Check if a token JTI is in the blocklist."""
+    redis_client = get_revoked_tokens_storage()
+    if redis_client:
+        try:
+            return redis_client.exists(f"blocklist:{jti}") > 0
+        except Exception as e:
+            logger.warning(f"Failed to check Redis blocklist: {e}")
+    # Fallback to in-memory
+    return jti in _revoked_tokens
+
+
+@jwt.token_in_blocklist_loader
+def check_if_token_in_blocklist(jwt_header, jwt_payload):
+    """Check if JWT access token has been revoked."""
+    jti = jwt_payload.get("jti")
+    if not jti:
+        return False
+    return is_token_in_blocklist(jti)
+
 
 from gefapi.models import User  # noqa:E402
 from gefapi.services import UserService  # noqa:E402
@@ -845,7 +897,27 @@ def refresh_token():
 @jwt_required()
 def logout():
     logger.info("[JWT]: User logout...")
-    refresh_token_string = request.json.get("refresh_token", None)
+    from flask_jwt_extended import get_jwt
+
+    # Revoke the current access token by adding its JTI to blocklist
+    try:
+        jwt_data = get_jwt()
+        jti = jwt_data.get("jti")
+        if jti:
+            # Get remaining token lifetime for blocklist expiration
+            exp = jwt_data.get("exp", 0)
+            import time
+
+            remaining_seconds = max(int(exp - time.time()), 0) + 60  # Add buffer
+            add_token_to_blocklist(jti, remaining_seconds)
+            logger.info(f"[JWT]: Access token {jti[:8]}... added to blocklist")
+    except Exception as e:
+        logger.warning(f"[JWT]: Failed to revoke access token: {e}")
+
+    # Revoke refresh token if provided
+    refresh_token_string = None
+    if request.json:
+        refresh_token_string = request.json.get("refresh_token", None)
 
     if refresh_token_string:
         # Import here to avoid circular imports
@@ -878,6 +950,47 @@ def user_lookup_callback(_jwt_header, jwt_data):
     return User.query.filter_by(id=identity).one_or_none()
 
 
+@jwt.expired_token_loader
+def expired_token_callback(jwt_header, jwt_payload):
+    """Handle expired JWT tokens with consistent error response."""
+    logger.debug("[JWT]: Expired token detected")
+    return jsonify(
+        {"status": 401, "detail": "Token has expired", "error": "token_expired"}
+    ), 401
+
+
+@jwt.invalid_token_loader
+def invalid_token_callback(error_message):
+    """Handle invalid JWT tokens with consistent error response."""
+    logger.warning(f"[JWT]: Invalid token: {error_message}")
+    return jsonify(
+        {"status": 401, "detail": "Invalid token", "error": "invalid_token"}
+    ), 401
+
+
+@jwt.unauthorized_loader
+def missing_token_callback(error_message):
+    """Handle missing JWT tokens with consistent error response."""
+    logger.debug(f"[JWT]: Missing token: {error_message}")
+    return jsonify(
+        {
+            "status": 401,
+            "detail": "Authorization token required",
+            "error": "authorization_required",
+        }
+    ), 401
+
+
+@jwt.revoked_token_loader
+def revoked_token_callback(jwt_header, jwt_payload):
+    """Handle revoked JWT tokens with consistent error response."""
+    jti = jwt_payload.get("jti", "unknown")
+    logger.info(f"[JWT]: Revoked token access attempt: {jti[:8]}...")
+    return jsonify(
+        {"status": 401, "detail": "Token has been revoked", "error": "token_revoked"}
+    ), 401
+
+
 @app.errorhandler(403)
 def forbidden(e):
     return error(status=403, detail="Forbidden")

gefapi/config/base.py

@@ -72,6 +72,9 @@
     "JWT_ACCESS_TOKEN_EXPIRES": timedelta(seconds=60 * 60 * 1),
     "JWT_REFRESH_TOKEN_EXPIRES": timedelta(days=30),  # 30 days for refresh tokens
     "JWT_TOKEN_LOCATION": ["headers"],
+    "JWT_IDENTITY_CLAIM": "sub",  # Standard JWT subject claim for identity
+    "JWT_BLOCKLIST_ENABLED": True,  # Enable token blocklist for revocation
+    "JWT_BLOCKLIST_TOKEN_CHECKS": ["access"],  # Check access tokens against blocklist
     "TRUSTED_PROXY_COUNT": int(os.getenv("TRUSTED_PROXY_COUNT", "0")),
     "INTERNAL_NETWORKS": [
         net.strip().strip("\"'")  # Remove quotes and whitespace

gefapi/models/refresh_token.py

@@ -20,7 +20,9 @@ class RefreshToken(db.Model):
     token = db.Column(db.String(255), nullable=False, index=True)
     expires_at = db.Column(db.DateTime, nullable=False, index=True)
     created_at = db.Column(
-        db.DateTime, default=datetime.datetime.utcnow, nullable=False
+        db.DateTime(timezone=True),
+        default=lambda: datetime.datetime.now(datetime.UTC),
+        nullable=False,
     )
     is_revoked = db.Column(db.Boolean, default=False, nullable=False)
     device_info = db.Column(db.String(500))  # Store user agent, IP, etc.
@@ -48,19 +50,62 @@ def generate_token():
     @staticmethod
     def default_expiry():
         """Default expiry time (30 days from now)"""
-        return datetime.datetime.utcnow() + datetime.timedelta(days=30)
-
-    def is_valid(self):
-        """Check if token is valid (not expired and not revoked)"""
-        return not self.is_revoked and self.expires_at > datetime.datetime.utcnow()
+        return datetime.datetime.now(datetime.UTC) + datetime.timedelta(days=30)
+
+    def is_valid(self, verify_client_ip=False, current_ip=None):
+        """Check if token is valid (not expired and not revoked).
+
+        Args:
+            verify_client_ip: If True, optionally verify the client IP matches
+            current_ip: Current client IP address for verification
+
+        Returns:
+            bool: True if token is valid, False otherwise
+        """
+        if self.is_revoked or self.expires_at <= datetime.datetime.now(datetime.UTC):
+            return False
+
+        # Optional client IP verification for additional security
+        if verify_client_ip and current_ip and self.device_info:
+            stored_ip = self._extract_ip_from_device_info()
+            if stored_ip and stored_ip != current_ip:
+                # Log suspicious activity but don't fail by default
+                # This provides visibility without breaking existing clients
+                import logging
+
+                logger = logging.getLogger(__name__)
+                logger.warning(
+                    f"Token IP mismatch: stored={stored_ip}, current={current_ip}, "
+                    f"token_id={self.id}"
+                )
+
+        return True
+
+    def _extract_ip_from_device_info(self):
+        """Extract IP address from stored device_info string."""
+        if not self.device_info:
+            return None
+        # device_info format: "IP: x.x.x.x | UA: ..."
+        if self.device_info.startswith("IP: "):
+            parts = self.device_info.split(" | ")
+            if parts:
+                return parts[0].replace("IP: ", "").strip()
+        return None
+
+    def get_client_fingerprint(self):
+        """Get a fingerprint of the client that created this token."""
+        return {
+            "ip_address": self._extract_ip_from_device_info(),
+            "device_info": self.device_info,
+        }
 
     def revoke(self):
         """Revoke the refresh token"""
         self.is_revoked = True
 
     def update_last_used(self):
         """Update last used timestamp"""
-        self.last_used_at = datetime.datetime.utcnow()
+        self.last_used_at = datetime.datetime.now(datetime.UTC)
 
     def serialize(self):
         """Return object data in easily serializable format"""