bug fixes routes: edit, logout

Author: ksh7

api/__init__.py

@@ -8,7 +8,7 @@
 from flask import Flask
 from flask_cors import CORS
 
-from .routes import rest_api, jwt
+from .routes import rest_api
 from .models import db
 
 app = Flask(__name__)
@@ -17,13 +17,13 @@
 
 db.init_app(app)
 rest_api.init_app(app)
-jwt.init_app(app)
 CORS(app)
 
 """
    Custom responses
 """
 
+
 @app.after_request
 def after_request(response):
     """
@@ -33,7 +33,6 @@ def after_request(response):
     if int(response.status_code) >= 400:
         response_data = json.loads(response.get_data())
         if "errors" in response_data:
-            _err_keys = response_data.keys()
             response_data = {"success": False,
                              "msg": list(response_data["errors"].items())[0][1]}
             response.set_data(json.dumps(response_data))

api/models.py

@@ -6,18 +6,19 @@
 from datetime import datetime
 
 import json
-from json import JSONEncoder
 
 from werkzeug.security import generate_password_hash, check_password_hash
 from flask_sqlalchemy import SQLAlchemy
 
 db = SQLAlchemy()
 
+
 class Users(db.Model):
     id = db.Column(db.Integer(), primary_key=True)
     username = db.Column(db.String(32), nullable=False)
     email = db.Column(db.String(64), nullable=False)
     password = db.Column(db.Text())
+    jwt_auth_active = db.Column(db.Boolean())
     date_joined = db.Column(db.DateTime(), default=datetime.utcnow)
 
     def __repr__(self):
@@ -39,8 +40,11 @@ def update_email(self, new_email):
     def update_username(self, new_username):
         self.username = new_username
 
-    def update_email(self, new_email):
-        self.email = new_email
+    def check_jwt_auth_active(self):
+        return self.jwt_auth_active
+
+    def set_jwt_auth_active(self, set_status):
+        self.jwt_auth_active = set_status
 
     @classmethod
     def get_by_id(cls, id):
@@ -51,7 +55,7 @@ def get_by_email(cls, email):
         return cls.query.filter_by(email=email).first()
 
     def toDICT(self):
-        
+
         cls_dict = {}
         cls_dict['_id'] = self.id
         cls_dict['username'] = self.username
@@ -61,15 +65,16 @@ def toDICT(self):
 
     def toJSON(self):
 
-        return json.dumps( self.toDICT() )
+        return json.dumps(self.toDICT())
+
 
 class JWTTokenBlocklist(db.Model):
     id = db.Column(db.Integer(), primary_key=True)
-    jti = db.Column(db.String(36), nullable=False)
+    jwt_token = db.Column(db.String(), nullable=False)
     created_at = db.Column(db.DateTime(), nullable=False)
 
     def __repr__(self):
-        return f"JTI {self.jti}"
+        return f"Expired Token: {self.jwt_token}"
 
     def save(self):
         db.session.add(self)

api/routes.py

@@ -3,29 +3,19 @@
 Copyright (c) 2019 - present AppSeed.us
 """
 
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
 
-import json
+from functools import wraps
 
 from flask import request
-from flask_restx import Api, Resource, fields, abort
-from flask_jwt_extended import JWTManager, jwt_required, get_jwt, create_access_token, get_jwt_identity
+from flask_restx import Api, Resource, fields
+
+import jwt
 
 from .models import db, Users, JWTTokenBlocklist
+from .config import BaseConfig
 
 rest_api = Api(version="1.0", title="Users API")
-jwt = JWTManager()
-
-"""
-    Helper function for revoking JWT token
-"""
-
-
-@jwt.token_in_blocklist_loader
-def check_if_token_revoked(jwt_header, jwt_payload):
-    jti = jwt_payload["jti"]
-    token = db.session.query(JWTTokenBlocklist.id).filter_by(jti=jti).scalar()
-    return token is not None
 
 
 """
@@ -49,6 +39,47 @@ def check_if_token_revoked(jwt_header, jwt_payload):
 logout_model = rest_api.model('LogoutModel', {"token": fields.String(required=True)})
 
 
+"""
+   Helper function for JWT token required
+"""
+
+def token_required(f):
+
+    @wraps(f)
+    def decorator(*args, **kwargs):
+
+        token = None
+
+        if "authorization" in request.headers:
+            token = request.headers["authorization"]
+
+        if not token:
+            return {"success": False, "msg": "Valid JWT token is missing"}, 400
+
+        try:
+            data = jwt.decode(token, BaseConfig.SECRET_KEY, algorithms=["HS256"])
+            current_user = Users.get_by_email(data["email"])
+
+            if not current_user:
+                return {"success": False,
+                        "msg": "Sorry. Wrong auth token. This user does not exist."}, 400
+
+            token_expired = db.session.query(JWTTokenBlocklist.id).filter_by(jwt_token=token).scalar()
+
+            if token_expired is not None:
+                return {"success": False, "msg": "Token revoked."}, 400
+
+            if not current_user.check_jwt_auth_active():
+                return {"success": False, "msg": "Token expired."}, 400
+
+        except:
+            return {"success": False, "msg": "Token is invalid"}, 400
+
+        return f(current_user, *args, **kwargs)
+
+    return decorator
+
+
 """
     Flask-Restx routes
 """
@@ -80,7 +111,7 @@ def post(self):
         new_user.save()
 
         return {"success": True,
-                "userID" : new_user.id,
+                "userID": new_user.id,
                 "msg": "The user was successfully registered"}, 200
 
 
@@ -109,11 +140,14 @@ def post(self):
                     "msg": "Sorry. Wrong credentials."}, 400
 
         # create access token uwing JWT
-        access_token = create_access_token(identity=_email)
+        token = jwt.encode({'email': _email, 'exp': datetime.utcnow() + timedelta(minutes=30)}, BaseConfig.SECRET_KEY)
+
+        user_exists.set_jwt_auth_active(True)
+        user_exists.save()
 
         return {"success": True,
-                "token": access_token,
-                "user" : user_exists.toJSON() }, 200
+                "token": token,
+                "user": user_exists.toJSON()}, 200
 
 
 @rest_api.route('/api/users/edit')
@@ -123,28 +157,21 @@ class EditUser(Resource):
     """
 
     @rest_api.expect(user_edit_model)
-    @jwt_required()
-    def post(self):
-
-        user_email = get_jwt_identity()
-        current_user = Users.get_by_email(user_email)
-
-        if not current_user:
-            return {"success": False,
-                    "msg": "Sorry. Wrong auth token. This user does not exist."}, 400
+    @token_required
+    def post(self, current_user):
 
         req_data = request.get_json()
 
         _new_username = req_data.get("username")
         _new_email = req_data.get("email")
 
         if _new_username:
-            current_user.update_username(_new_username)
+            self.update_username(_new_username)
 
         if _new_email:
-            current_user.update_email(_new_email)
+            self.update_email(_new_email)
 
-        current_user.save()
+        self.save()
 
         return {"success": True}, 200
 
@@ -156,18 +183,17 @@ class LogoutUser(Resource):
     """
 
     @rest_api.expect(logout_model, validate=True)
-    @jwt_required()
-    def post(self):
+    @token_required
+    def post(self, current_user):
 
-        user_email = get_jwt_identity()
-        current_user = Users.get_by_email(user_email)
-
-        if not current_user:
-            return {"success": False,
-                    "msg": "Sorry. Wrong auth token"}, 400
+        req_data = request.get_json()
+        _jwt_token = req_data.get("token")
 
-        jwt_block = JWTTokenBlocklist(jti=get_jwt()["jti"], created_at=datetime.now(timezone.utc))
+        jwt_block = JWTTokenBlocklist(jwt_token=_jwt_token, created_at=datetime.now(timezone.utc))
         jwt_block.save()
 
+        self.set_jwt_auth_active(False)
+        self.save()
+
         return {"success": True,
                 "msg": "JWT Token revoked successfully"}, 200