add jwt extended features

Author: pclavier92

Diff for: __pycache__/blacklist.cpython-37.pyc

@@ -1,33 +1,88 @@
 import os
 
-from flask import Flask
+from flask import Flask, jsonify
 from flask_restful import Api
-from flask_jwt import JWT
+from flask_jwt_extended import JWTManager
 
-from security import authenticate, identity
-from resources.user import UserRegister
+from blacklist import BLACKLIST
+from resources.user import User, UserRegister, UserLogin, UserLogout, TokenRefresh
 from resources.item import Item, ItemList
 from resources.store import Store, StoreList
 
 app = Flask(__name__)
-app.secret_key = 'some_secret_key'
+app.secret_key = 'some_secret_key' # app.config['JWT_SECRET_KEY']
 # Could be sqlite, mysql, postgresql, oracle ... without changing any of the code
 app.config['SQLALCHEMY_DATABASE_URI'] = os.environ.get('DATABASE_URL', 'sqlite:///data.db') # second argument default value
 # Turns off Flask-SQLAlchemy tracking but no for SQLAlchemy
 app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
+# Propagates internal errors to the user (not secure, but good for debugin)
+app.config['PROPAGATE_EXCEPTIONS'] = True 
+app.config['JWT_BLACKLIST_ENABLED'] = True
+app.config['JWT_BLACKLIST_TOKEN_CHECKS'] = ['access', 'refresh']
 api = Api(app)
 
-# @app.before_first_request
-# def create_tables():
-#     db.create_all()
+@app.before_first_request
+def create_tables():
+    db.create_all()
 
-jwt = JWT(app, authenticate, identity) # /auth
+jwt = JWTManager(app) # not creating /auth
 
+@jwt.user_claims_loader
+def add_claims_to_jwt(identity):
+    # Insted of hard-coding, you should read from a config file or db
+    if identity == 1:
+        return {'is_admin': True}
+    return {'is_admin': False}
+
+@jwt.token_in_blacklist_loader
+def check_if_token_in_blacklist(decrypted_token):
+    return decrypted_token['jti'] in BLACKLIST
+
+@jwt.expired_token_loader
+def expired_token_callback():
+    return jsonify({
+        'message': 'The token has expired',
+        'error': 'token_expired'
+    }), 401
+
+@jwt.invalid_token_loader
+def invalid_token_callback(error):
+    return jsonify({
+        'message': 'Signature verification failed',
+        'error': 'invalid_token'
+    }), 401
+
+@jwt.unauthorized_loader
+def unauthorized_callback(error):
+    return jsonify({
+        'message': 'Unauthorized to access this resource',
+        'error': 'unauthorization_required'
+    }), 401
+
+@jwt.needs_fresh_token_loader
+def needs_fresh_token_callback(error):
+    return jsonify({
+        'message': 'Nedd fresh token to access this resource',
+        'error': 'needs_fresh_token'
+    }), 401
+
+
+@jwt.revoked_token_loader
+def revoked_token_callback(error):
+    return jsonify({
+        'message': 'You dont have permission any more',
+        'error': 'revoked_token'
+    }), 401
+
+api.add_resource(User, '/user/<int:user_id>')
 api.add_resource(UserRegister,'/register')
+api.add_resource(UserLogin, '/login')
+api.add_resource(UserLogout, '/logout')
 api.add_resource(ItemList,'/items')
 api.add_resource(Item,'/item/<string:name>')
 api.add_resource(StoreList,'/stores')
 api.add_resource(Store,'/store/<string:name>')
+api.add_resource(TokenRefresh, '/refresh')
 
 if __name__ == '__main__':
     from db import db

Diff for: app.py

@@ -0,0 +1 @@
+BLACKLIST = set()

Diff for: blacklist.py

@@ -16,11 +16,16 @@ def __init__(self, name, price, store_id):
         self.store_id = store_id
 
     def json(self):
-        return {'name': self.name, 'price': self.price}
+        return {
+            'id': self.id,
+            'name': self.name, 
+            'price': self.price, 
+            'store_id': self.store_id
+        }
 
     @classmethod
     def get_all(cls):
-        return ItemModel.query.all()
+        return cls.query.all()
 
     @classmethod
     def find_by_name(cls,name):

Diff for: data.db

@@ -13,12 +13,15 @@ def __init__(self, name):
         self.name = name
 
     def json(self):
-        itemsJson = [item.json() for item in self.items.all()]
-        return {'name': self.name, 'items': itemsJson}
+        return {
+            'id': self.id,     
+            'name': self.name, 
+            'items': [item.json() for item in self.items.all()]
+        }
 
     @classmethod
     def get_all(cls):
-        return StoreModel.query.all()
+        return cls.query.all()
 
     @classmethod
     def find_by_name(cls,name):

Diff for: models/__pycache__/item.cpython-37.pyc

@@ -11,6 +11,12 @@ def __init__(self, username, password):
         self.username = username
         self.password = password
 
+    def json(self):
+        return {
+            'id': self.id,
+            'username': self.username 
+        }
+
     @classmethod
     def find_by_username(cls, username):
         return cls.query.filter_by(username=username).first()
@@ -21,4 +27,8 @@ def find_by_id(cls, _id):
 
     def save_to_db(self):
         db.session.add(self)
+        db.session.commit()
+
+    def delete(self):
+        db.session.delete(self)
         db.session.commit()

Diff for: models/__pycache__/store.cpython-37.pyc

@@ -1,6 +1,6 @@
 Flask
 Flask-RESTful
 Flask-SQLAlchemy
-Flask-JWT
+Flask-JWT-Extended
 uwsgi
 psycopg2

Diff for: models/__pycache__/user.cpython-37.pyc

@@ -1,5 +1,11 @@
 from flask_restful import Resource, reqparse
-from flask_jwt import jwt_required
+from flask_jwt_extended import (
+    jwt_required, 
+    get_jwt_claims, 
+    jwt_optional, 
+    get_jwt_identity,
+    fresh_jwt_required
+)
 from models.item import ItemModel
 
 class Item(Resource):
@@ -26,7 +32,7 @@ def get(self, name):
             return item.json()
         return {'message': 'Item not found'},404
 
-    @jwt_required()
+    @jwt_required
     def post(self, name):
         if ItemModel.find_by_name(name):
             return {'error': "An item with name {} already exists".format(name)}, 400
@@ -38,7 +44,7 @@ def post(self, name):
             return {'message': 'Error inserting item'}, 500
         return item.json(), 201
 
-    @jwt_required()
+    @jwt_required
     def put(self, name):
         data = Item.parser.parse_args()
         item = ItemModel.find_by_name(name)
@@ -50,15 +56,26 @@ def put(self, name):
         item.save_to_db()
         return item.json(), 201
 
-    @jwt_required()
+    @fresh_jwt_required
     def delete(self, name):
+        claims = get_jwt_claims()
+        if not claims['is_admin']:
+            return {'message': 'Admin privilege required'}, 401
         item = ItemModel.find_by_name(name)
         if item:
             item.delete()
         return {'message': 'Item deleted'}
 
 class ItemList(Resource):
+    @jwt_optional
     def get(self):
+        user_id = get_jwt_identity()
         items = ItemModel.get_all()
-        return {'items': [i.json() for i in items ]}
-        # list(map(lambda i: i.json(), items))
+        # If logged in
+        if user_id:
+            return {'items': [i.json() for i in items ]}, 200
+            # list(map(lambda i: i.json(), items))
+        return {
+            'items': [i.name for i in items],
+            'message': 'More info if logged in'
+        }, 200

Diff for: models/item.py

@@ -1,22 +1,79 @@
 from flask_restful import Resource, reqparse
-from models.user import UserModel 
+from flask_jwt_extended import (
+    create_access_token, 
+    create_refresh_token,
+    jwt_refresh_token_required, 
+    get_jwt_identity,
+    jwt_required,
+    get_raw_jwt
+)
+from werkzeug.security import safe_str_cmp
+from models.user import UserModel
+from blacklist import BLACKLIST 
 
-class UserRegister(Resource):
-    parser = reqparse.RequestParser()
-    parser.add_argument('username',
-        type=str,
-        required=True,
-        help="This field cannot be left blank"
-    )
-    parser.add_argument('password',
-        type=str,
-        required=True,
-        help="This field cannot be left blank"
-    )
+_user_parser = reqparse.RequestParser()
+_user_parser.add_argument('username',
+    type=str,
+    required=True,
+    help="This field cannot be left blank"
+)
+_user_parser.add_argument('password',
+    type=str,
+    required=True,
+    help="This field cannot be left blank"
+)
 
+class UserRegister(Resource):
     def post(self):
-        data = UserRegister.parser.parse_args() 
+        data = _user_parser.parse_args() 
         if UserModel.find_by_username(data['username']):
             return {'message': 'Username already exists'}, 400
         UserModel(**data).save_to_db()
-        return {'message': 'User created successfully'}, 201
+        return {'message': 'User created successfully'}, 201
+
+class User(Resource):
+    @classmethod
+    def get(cls, user_id):
+        user = UserModel.find_by_id(user_id)
+        if not user:
+            return {'message': 'User not found'}, 404
+        return user.json()
+
+    @classmethod
+    def delete(cls, user_id):
+        user = UserModel.find_by_id(user_id)
+        if not user:
+            return {'message': 'User not found'}, 404
+        user.delete()
+        return {'message': 'User deleted'}, 200
+
+class UserLogin(Resource):
+    @classmethod
+    def post(cls):
+        data = _user_parser.parse_args()
+        user = UserModel.find_by_username(data['username'])
+        # This is what the authenticate() function used to do
+        if user and safe_str_cmp(user.password, data['password']):
+            # This is what the identity() function used to do
+            access_token = create_access_token(identity=user.id, fresh=True)
+            refresh_token = create_refresh_token(user.id)
+            return {
+                'access_token': access_token,
+                'refresh_token': refresh_token
+            }, 200
+        return {'message': 'Invalid credentials'}, 401
+
+class UserLogout(Resource):
+    @jwt_required
+    def post(self):
+        jti = get_raw_jwt()['jti']
+        BLACKLIST.add(jti)
+        return {'message': 'Successfully logged out'}, 200
+
+
+class TokenRefresh(Resource):
+    @jwt_refresh_token_required
+    def post(self):
+        current_user = get_jwt_identity()
+        new_token = create_access_token(identity=current_user, fresh=False)
+        return {'access_token': new_token}, 200