Example Flask with flask-restful and flask-jwt-extended

Author: joeavanzato

backend/app.py

@@ -0,0 +1,159 @@
+from flask import Flask, request, jsonify
+from pymongo import MongoClient
+from flask_restful import Resource, Api
+from bson.json_util import dumps
+from passlib.hash import sha256_crypt
+from flask_jwt_extended import create_access_token,create_refresh_token,get_jwt,get_jwt_identity, jwt_required, JWTManager
+import datetime
+import time
+# py -m flask --app app --debug run
+
+client = MongoClient()
+db = client.example_db
+demo_collection = db.demo_collection
+user_collection = db.user_collection
+token_collection = db.token_collection
+
+app = Flask(__name__, template_folder='templates')
+
+# For helping with Same-Origin requests from the React App
+@app.after_request
+def after_request(response):
+ response.headers.add('Access-Control-Allow-Origin', '*')
+ response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
+ response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
+ return response
+
+api = Api(app)
+
+jwt = JWTManager(app)
+app.config['JWT_SECRET_KEY'] = 'jwt-secret-string-example'
+app.config['JWT_ACCESS_TOKEN_EXPIRES'] = datetime.timedelta(minutes=240)
+app.config['JWT_REFRESH_TOKEN_EXPIRES'] = datetime.timedelta(days=30)
+
+@jwt.token_in_blocklist_loader
+def check_if_token_in_blacklist(jwt_header, jwt_payload: dict):
+ jti = jwt_payload['jti']
+ return is_jti_blacklisted(jti)
+
+def is_jti_blacklisted(jti):
+ if db.token_collection.count_documents({"jti": jti}) != 0:
+ return True
+ else:
+ return False
+
+app.config["INITIAL_USERNAME"] = "user@test.com"
+app.config["INITIAL_PASSWORD"] = sha256_crypt.encrypt("initial_password")
+
+# Create the Initial User specified in the configuration data
+def create_initial_user():
+ temp_dict = {"email": app.config["INITIAL_USERNAME"], "password": app.config["INITIAL_PASSWORD"]}
+ if db.user_collection.count_documents({"email": app.config["INITIAL_USERNAME"]}) != 0:
+ print("[*] Found Initial User")
+ filter = {"email": app.config["INITIAL_USERNAME"]}
+ record = user_collection.update_one(filter, {"$set": temp_dict})
+ else:
+ print("[*] Adding Initial User")
+ user_collection.update_one(temp_dict, {'$set':temp_dict}, upsert=True)
+create_initial_user()
+
+# Verify the initial user was created and password functionality is working as expected
+def verify_initial_user():
+ temp_doc = user_collection.find_one({'email': app.config["INITIAL_USERNAME"]}, {})
+ if sha256_crypt.verify("change_me", temp_doc['password']):
+ print("[*] Successfully Verified Initial User")
+verify_initial_user()
+
+
+# REST API Specifications for retrieving and setting scan configurations
+class Scan(Resource):
+ @jwt_required()
+ def get(self):
+ time.sleep(1)
+ scans = demo_collection.aggregate([{"$project": {'id': '$_id', 'name': 1, 'category': 1, 'test_object': 1, '_id': 0}}, ])
+ return dumps(scans)
+
+ @jwt_required()
+ def post(self):
+ data = request.get_json()
+ return jsonify({'data': data}), 201
+api.add_resource(Scan, '/api/demo')
+
+# REST API Specifications for creating JSON Web Tokens (JWTs)
+class Login(Resource):
+ def post(self):
+ email = request.json.get("email", None)
+ password = request.json.get("password", None)
+ if (email is None) or (password is None):
+ return {"msg": "Wrong Email or Password!"}, 401
+ temp_doc = user_collection.find_one({'email':email}, {})
+ if db.user_collection.count_documents({"email": email}) != 0:
+ if sha256_crypt.verify(password, temp_doc['password']):
+ print("[*] Successfully Verified User")
+ access_token = create_access_token(identity=email)
+ refresh_token = create_refresh_token(identity=email)
+ response = {"msg":"Successfully Authenticated", "access_token": access_token, "refresh_token": refresh_token}
+ return response
+ else:
+ return {"msg": "Wrong Email or Password!"}, 401
+ else:
+ return {"msg": "Wrong Email or Password!"}, 401
+api.add_resource(Login, '/api/login')
+
+# Revoke an Active Token via adding to Block List
+class Logout(Resource):
+ @jwt_required()
+ def post(self):
+ jti_data = get_jwt()
+ id = get_jwt_identity()
+ jti = jti_data['jti']
+ try:
+ if is_jti_blacklisted(jti):
+ pass
+ else:
+ revoked_token = RevokedToken(id, jti)
+ revoked_token.add()
+ return {'msg': 'Access Token has been Revoked'}
+ except Exception as e:
+ print(e)
+ return {'msg': 'Something Went Wrong Revoking Access Token'}, 500
+api.add_resource(Logout, '/api/logout')
+
+class TokenRefresh(Resource):
+ @jwt_required(refresh=True)
+ def post(self):
+ current_user = get_jwt_identity()
+ access_token = create_access_token(identity = current_user)
+ return {'access_token': access_token}
+api.add_resource(TokenRefresh, '/api/refresh')
+
+class ValidateToken(Resource):
+ @jwt_required()
+ def get(self):
+ try:
+ jti_data = get_jwt()
+ jti = jti_data['jti']
+ print(jti)
+ if is_jti_blacklisted(jti):
+ return {'msg': 'Revoked'}
+ else:
+ return {'msg': 'Not Revoked'}
+ except Exception as e:
+ print(e)
+ return {'msg': 'Invalid JWT'}
+api.add_resource(ValidateToken, '/api/validate')
+
+class Add_User(Resource):
+ @jwt_required()
+ def post(self):
+ pass
+
+
+class RevokedToken:
+ def __init__(self, ident, jti):
+ self.id = ident
+ self.jti = jti
+ self.temp_dict = {"id":self.id, "jti":self.jti}
+
+ def add(self):
+ token_collection.insert_one(self.temp_dict, {'$set':self.temp_dict})