From dcf3c3d7b795b3b4bc5a4d725a4b95930d53faf7 Mon Sep 17 00:00:00 2001
From: Justin Campbell <campb303@w2vm4.ecn.purdue.edu>
Date: Sun, 8 Nov 2020 16:36:41 -0500
Subject: [PATCH] Implement login, token generation and route protection with
 JWTs

---
 api/api.py | 47 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/api/api.py b/api/api.py
index a9752f8..3d2da84 100644
--- a/api/api.py
+++ b/api/api.py
@@ -1,5 +1,9 @@
-from flask import Flask, request
+from flask import Flask, request, after_this_request
 from flask_restful import Api, Resource
+from flask_jwt_extended import (
+    JWTManager, create_access_token, create_refresh_token,
+    jwt_required, set_refresh_cookies
+)
 from werkzeug.security import check_password_hash
 import os, dotenv
 import ECNQueue
@@ -14,6 +18,30 @@
 api = Api(app)
 
 
+################################################################################
+# Configure Flask-JWT-Extended
+################################################################################
+
+# Set JWT secret key and create JWT manager
+app.config["JWT_SECRET_KEY"] = os.environ.get("JWT_SECRET_KEY")
+# Set identity claim field key to sub for JWT RFC complience
+# Flask-JWT-Extended uses 'identity' by default for compatibility reasons
+app.config["JWT_IDENTITY_CLAIM"] = "sub"
+# Set the key for error messages generated by Flask-JWT-Extended
+app.config["JWT_ERROR_MESSAGE_KEY"] = "message"
+
+# Look for JWTs in headers (for access) then cookies (for refresh)
+app.config["JWT_TOKEN_LOCATION"] = ["headers", "cookies"]
+# Restrict cookies to HTTPS in prod, allow HTTP in dev
+app.config["JWT_COOKIE_SECURE"] = False if os.environ.get("ENVIRONMENT") == "dev" else True
+# Restrict cookies using SameSite=strict flag
+app.config["JWT_COOKIE_SAMESITE"] = "strict"
+# Restrict refresh tokens to /token/refresh endpoint
+app.config['JWT_REFRESH_COOKIE_PATH'] = '/token/refresh'
+
+tokenManager = JWTManager(app)
+
+
 
 class Login(Resource):
     def post(self):
@@ -32,9 +60,22 @@ def post(self):
         if not check_password_hash(os.environ.get("SHARED_PASSWORD_HASH"), data["password"]):
             return ({ "message": "Password invalid"}, 401)
 
-        return ({ "message": "Login successful"}, 200)
+        access_token = create_access_token(data["username"])
+        refresh_token = create_refresh_token(data["username"])
+
+
+        # This decorator is needed because Flask-RESTful's 'resourceful routing`
+        # doesn't allow for direct modification to the Flask response object.
+        # See: https://flask-restful.readthedocs.io/en/latest/quickstart.html#resourceful-routing
+        @after_this_request
+        def _does_this_work(response):
+            set_refresh_cookies(response, refresh_token)
+            return response
+
+        return { "token": access_token }
 
 class Item(Resource):
+    @jwt_required
     def get(self, queue: str, number: int) -> str:
         """Returns the JSON representation of the item requested.
 
@@ -67,6 +108,7 @@ def get(self, queue: str, number: int) -> str:
         return ECNQueue.Item(queue, number).toJson()
 
 class Queue(Resource):
+    @jwt_required
     def get(self, queue: str) -> str:
         """Returns the JSON representation of the queue requested.
 
@@ -85,6 +127,7 @@ def get(self, queue: str) -> str:
         return queues
 
 class QueueList(Resource):
+    @jwt_required
     def get(self) -> list:
         """Returns a list of dictionaries with the number of items in each queue.