Fix user creation

Author: azvoleff

gefapi/__init__.py

@@ -66,6 +66,7 @@
 
 
 from gefapi.services import UserService  # noqa:E402
+from gefapi.models import User  # noqa:E402
 
 
 @app.route("/auth", methods=["POST"])
@@ -83,6 +84,12 @@ def create_token():
     return jsonify({"access_token": access_token, "user_id": user.id})
 
 
+@jwt.user_lookup_loader
+def user_lookup_callback(_jwt_header, jwt_data):
+    identity = jwt_data["sub"]
+    return User.query.filter_by(id=identity).one_or_none()
+
+
 @app.errorhandler(403)
 def forbidden(e):
     return error(status=403, detail="Forbidden")

gefapi/routes/api/v1/gef_api_router.py

@@ -6,7 +6,7 @@
 
 import dateutil.parser
 from flask import Response, json, jsonify, request, send_from_directory
-from flask_jwt_extended import get_jwt_identity, jwt_required
+from flask_jwt_extended import jwt_required, current_user
 
 from gefapi.errors import (
     EmailError,
@@ -33,18 +33,6 @@
 logger = logging.getLogger()
 
 
-@jwt_required()
-def get_identity():
-    user = None
-    try:
-        id = get_jwt_identity()
-        user = UserService.get_user(id)
-    except Exception as e:
-        logger.error(str(e))
-        logger.error("[JWT]: Error getting user for %s" % (id))
-    return user
-
-
 # SCRIPT CREATION
 @endpoints.route("/script", strict_slashes=False, methods=["POST"])
 @jwt_required()
@@ -57,7 +45,7 @@ def create_script():
     sent_file = request.files.get("file")
     if sent_file.filename == "":
         sent_file.filename = "script"
-    user = get_identity()
+    user = current_user
     try:
         user = ScriptService.create_script(sent_file, user)
     except InvalidFile as e:
@@ -80,7 +68,7 @@ def get_scripts():
     include = request.args.get("include")
     include = include.split(",") if include else []
     try:
-        scripts = ScriptService.get_scripts(get_identity())
+        scripts = ScriptService.get_scripts(current_user)
     except Exception as e:
         logger.error("[ROUTER]: " + str(e))
         return error(status=500, detail="Generic Error")
@@ -95,7 +83,7 @@ def get_script(script):
     include = request.args.get("include")
     include = include.split(",") if include else []
     try:
-        script = ScriptService.get_script(script, get_identity())
+        script = ScriptService.get_script(script, current_user)
     except ScriptNotFound as e:
         logger.error("[ROUTER]: " + e.message)
         return error(status=404, detail=e.message)
@@ -111,7 +99,7 @@ def publish_script(script):
     """Publish a script"""
     logger.info("[ROUTER]: Publishing script " + script)
     try:
-        script = ScriptService.publish_script(script, get_identity())
+        script = ScriptService.publish_script(script, current_user)
     except ScriptNotFound as e:
         logger.error("[ROUTER]: " + e.message)
         return error(status=404, detail=e.message)
@@ -127,7 +115,7 @@ def unpublish_script(script):
     """Unpublish a script"""
     logger.info("[ROUTER]: Unpublishsing script " + script)
     try:
-        script = ScriptService.unpublish_script(script, get_identity())
+        script = ScriptService.unpublish_script(script, current_user)
     except ScriptNotFound as e:
         logger.error("[ROUTER]: " + e.message)
         return error(status=404, detail=e.message)
@@ -143,7 +131,7 @@ def download_script(script):
     """Download a script"""
     logger.info("[ROUTER]: Download script " + script)
     try:
-        script = ScriptService.get_script(script, get_identity())
+        script = ScriptService.get_script(script, current_user)
 
         temp_dir = tempfile.TemporaryDirectory().name
         script_file = script.slug + ".tar.gz"
@@ -191,7 +179,7 @@ def update_script(script):
     sent_file = request.files.get("file")
     if sent_file.filename == "":
         sent_file.filename = "script"
-    user = get_identity()
+    user = current_user
     # if user.role != 'ADMIN' and user.email != '[email protected]':
     #     return error(status=403, detail='Forbidden')
     try:
@@ -216,7 +204,7 @@ def update_script(script):
 def delete_script(script):
     """Delete a script"""
     logger.info("[ROUTER]: Deleting script: " + script)
-    identity = get_identity()
+    identity = current_user
     if identity.role != "ADMIN" and identity.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -236,7 +224,7 @@ def delete_script(script):
 def run_script(script):
     """Run a script"""
     logger.info("[ROUTER]: Running script: " + script)
-    user = get_identity()
+    user = current_user
     try:
         params = request.args.to_dict() if request.args else {}
         if request.get_json(silent=True):
@@ -270,9 +258,7 @@ def get_executions():
     exclude = request.args.get("exclude")
     exclude = exclude.split(",") if exclude else []
     try:
-        executions = ExecutionService.get_executions(
-            get_identity(), user_id, updated_at
-        )
+        executions = ExecutionService.get_executions(current_user, user_id, updated_at)
     except Exception as e:
         logger.error("[ROUTER]: " + str(e))
         return error(status=500, detail="Generic Error")
@@ -291,7 +277,7 @@ def get_execution(execution):
     exclude = request.args.get("exclude")
     exclude = exclude.split(",") if exclude else []
     try:
-        execution = ExecutionService.get_execution(execution, get_identity())
+        execution = ExecutionService.get_execution(execution, current_user)
     except ExecutionNotFound as e:
         logger.error("[ROUTER]: " + e.message)
         return error(status=404, detail=e.message)
@@ -308,7 +294,7 @@ def update_execution(execution):
     """Update an execution"""
     logger.info("[ROUTER]: Updating execution " + execution)
     body = request.get_json()
-    user = get_identity()
+    user = current_user
     if user.role != "ADMIN" and user.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -367,7 +353,7 @@ def create_execution_log(execution):
     """Create log of an execution"""
     logger.info("[ROUTER]: Creating execution log for " + execution)
     body = request.get_json()
-    user = get_identity()
+    user = current_user
     if user.role != "ADMIN" and user.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -389,13 +375,15 @@ def create_user():
     logger.info("[ROUTER]: Creating user")
     body = request.get_json()
     if request.headers.get("Authorization", None) is not None:
+        logger.debug("[ROUTER]: Authorization header found")
 
         @jwt_required()
         def identity():
             pass
 
         identity()
-    identity = get_identity()
+    logger.debug("[ROUTER]: Getting identity")
+    identity = current_user
     if identity:
         user_role = body.get("role", "USER")
         if identity.role == "USER" and user_role == "ADMIN":
@@ -420,7 +408,7 @@ def get_users():
     logger.info("[ROUTER]: Getting all users")
     include = request.args.get("include")
     include = include.split(",") if include else []
-    identity = get_identity()
+    identity = current_user
     if identity.role != "ADMIN" and identity.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -438,7 +426,7 @@ def get_user(user):
     logger.info("[ROUTER]: Getting user" + user)
     include = request.args.get("include")
     include = include.split(",") if include else []
-    identity = get_identity()
+    identity = current_user
     if identity.role != "ADMIN" and identity.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -457,7 +445,7 @@ def get_user(user):
 def get_me():
     """Get me"""
     logger.info("[ROUTER]: Getting my user")
-    user = get_identity()
+    user = current_user
     return jsonify(data=user.serialize()), 200
 
 
@@ -467,7 +455,7 @@ def update_profile():
     """Update an user"""
     logger.info("[ROUTER]: Updating profile")
     body = request.get_json()
-    identity = get_identity()
+    identity = current_user
     try:
         password = body.get("password", None)
         repeat_password = body.get("repeatPassword", None)
@@ -501,7 +489,7 @@ def update_profile():
 def delete_profile():
     """Delete Me"""
     logger.info("[ROUTER]: Delete me")
-    identity = get_identity()
+    identity = current_user
     try:
         user = UserService.delete_user(str(identity.id))
     except UserNotFound as e:
@@ -540,7 +528,7 @@ def update_user(user):
     """Update an user"""
     logger.info("[ROUTER]: Updating user" + user)
     body = request.get_json()
-    identity = get_identity()
+    identity = current_user
     if identity.role != "ADMIN" and identity.email != "[email protected]":
         return error(status=403, detail="Forbidden")
     try:
@@ -559,7 +547,7 @@ def update_user(user):
 def delete_user(user):
     """Delete an user"""
     logger.info("[ROUTER]: Deleting user" + user)
-    identity = get_identity()
+    identity = current_user
     if user == "[email protected]":
         return error(status=403, detail="Forbidden")
     if identity.role != "ADMIN" and identity.email != "[email protected]":