Migrate to jwt token management

Author: narfasec

frontend/admin-one-vue-tailwind-master/src/router/index.js

@@ -169,16 +169,50 @@ const router = createRouter({
   }
 })
 
+// Helper function to validate JWT token format
+function isValidJWT(token) {
+  if (!token) return false;
+  
+  // JWT has 3 parts separated by dots
+  const parts = token.split('.');
+  if (parts.length !== 3) return false;
+  
+  // Each part should be base64url encoded
+  try {
+    // Decode header to check if it's a valid JWT
+    const header = JSON.parse(atob(parts[0].replace(/-/g, '+').replace(/_/g, '/')));
+    return header.typ === 'JWT' && (header.alg === 'HS256' || header.alg === 'RS256');
+  } catch (e) {
+    return false;
+  }
+}
+
+// Helper function to check if JWT is expired
+function isJWTExpired(token) {
+  try {
+    const parts = token.split('.');
+    const payload = JSON.parse(atob(parts[1].replace(/-/g, '+').replace(/_/g, '/')));
+    
+    // Check if token has expiration
+    if (payload.exp) {
+      const currentTime = Math.floor(Date.now() / 1000);
+      return payload.exp < currentTime;
+    }
+    
+    // If no expiration, consider it valid
+    return false;
+  } catch (e) {
+    return true; // If we can't parse, consider expired
+  }
+}
+
 router.beforeEach((to, from, next) => {
   // Check if the route requires authentication
   if (to.matched.some(record => record.meta.requiresAuth)) {
     const token = localStorage.getItem('token');
 
-    // Regex pattern to match the token
-    const regex = /^[A-Z0-9]{6}$/; // Simplified regex without quotes
-
-    // If the token is not valid, redirect to login
-    if (!token || !regex.test(token)) {
+    // Check if token is a valid JWT and not expired
+    if (!token || !isValidJWT(token) || isJWTExpired(token)) {
       next({
         path: '/login',
         query: { redirect: to.fullPath } // Save the intended route for after login

frontend/admin-one-vue-tailwind-master/src/stores/main.js

@@ -2,6 +2,17 @@ import { defineStore } from 'pinia'
 import { ref, computed } from 'vue'
 import axios from 'axios'
 
+// Helper function to decode JWT token
+function decodeJWT(token) {
+  try {
+    const parts = token.split('.');
+    const payload = JSON.parse(atob(parts[1].replace(/-/g, '+').replace(/_/g, '/')));
+    return payload;
+  } catch (e) {
+    return null;
+  }
+}
+
 export const useMainStore = defineStore('main', () => {
   const userName = ref(localStorage.getItem('username'))
   const userEmail = ref('doe.doe.doe@example.com')
@@ -34,6 +45,28 @@ export const useMainStore = defineStore('main', () => {
     }
   }
 
+  // Get user info from JWT token
+  function getUserFromJWT() {
+    const token = localStorage.getItem('token');
+    if (!token) return null;
+    
+    const payload = decodeJWT(token);
+    if (!payload) return null;
+    
+    return {
+      userId: payload.user_id,
+      compId: payload.comp_id,
+      isAdmin: payload.is_admin,
+      username: localStorage.getItem('username')
+    };
+  }
+
+  // Check if current user is admin
+  function isCurrentUserAdmin() {
+    const userInfo = getUserFromJWT();
+    return userInfo ? userInfo.isAdmin : false;
+  }
+
   function fetchSampleClients() {
     axios
       .get(`data-sources/clients.json?v=3`)
@@ -96,6 +129,7 @@ export const useMainStore = defineStore('main', () => {
     const url = "http://localhost:5000/clients"
     const clientsPayload = {
       user_id: Number(localStorage.getItem('userId')),
+      token: localStorage.getItem('token')
     };
     axios
       .post(url, clientsPayload)
@@ -231,6 +265,8 @@ export const useMainStore = defineStore('main', () => {
     getCompanyCashFlow,
     deleteEmployee,
     editCommission,
-    calculateSalesRevenue
+    calculateSalesRevenue,
+    getUserFromJWT,
+    isCurrentUserAdmin
   }
 })

server/api/auth/jwt_utils.py

@@ -0,0 +1,47 @@
+from flask import current_app
+import jwt
+import traceback
+
+
+def issue_token(user_id: int, comp_id: int, is_admin: bool) -> str:
+    """ Create a new token with user information """
+    payload = {
+        'user_id': user_id,
+        'comp_id': comp_id,
+        'is_admin': is_admin,
+    }
+    private_key = current_app.config.get('JWT_PRIVATE_PEM')
+    if private_key:
+        try:
+            token = jwt.encode(payload, key=private_key, algorithm='RS256')
+            str_token = token.decode('utf-8') if isinstance(token, (bytes, bytearray)) else token
+            return str_token
+        except Exception as e:
+            raise Exception('Issue RS256 token failed; falling back to HS256', e) from e
+            
+
+def validate_token(token: str):
+    """ Validate JWT token """
+    print("token", token)
+    if not token:
+        return False, None
+
+    try:
+        jwt.get_unverified_header(token)
+    except Exception as e:
+        print("Error decoding header", e, traceback.format_exc())
+        return False, None
+
+    try:
+        public_key_pem = current_app.config.get('JWT_PUBLIC_PEM', '')
+        if public_key_pem:
+            # Use PEM string directly - PyJWT 1.5.0 should handle this
+            payload = jwt.decode(
+                token,
+                key=public_key_pem,
+            )
+            return True, payload
+    except Exception as e:
+        print("Error decoding RS256 with public key", e)
+        pass
+    return False, None

server/api/auth/routes.py

@@ -1,8 +1,9 @@
-from flask import Flask, Blueprint, request, jsonify, current_app
+from flask import Blueprint, request, jsonify, current_app
 from db.db_connector import DBConnector
 from Crypto.Cipher import DES
 from Crypto.Util.Padding import pad, unpad
 import base64
+from .jwt_utils import issue_token, validate_token
 
 auth = Blueprint('auth', __name__)
 
@@ -44,23 +45,24 @@ def login():
         decrypted_password = str(decrypt_password(encrypted_password, DES_KEY))
         print(f'Password comparsion!! input: {password} vs decrypted_password: {decrypted_password}')
     if password == decrypted_password:
-        result = dbc.execute_query(query='update_user_activity', args={
+        dbc.execute_query(query='update_user_activity', args={
             'user_id': _id,
             'active': True
         })
         is_admin = dbc.execute_query(query='get_user_admin', args=_id)
         print(f'Admin --> {is_admin}')
         if is_admin == 1:
             is_admin = 'true'
-            token = current_app.config['ADMIN_AUTH_TOKEN']
         else:
             is_admin = 'false'
-            token = current_app.config['AUTH_TOKEN']
 
         comp_id = dbc.execute_query(query='get_user_comp_id', args=_id)
         if not isinstance(comp_id, int):
             return jsonify({'status': 'Bad request'}), 400
 
+        token: str = issue_token(user_id=_id, comp_id=comp_id, is_admin=(is_admin == 'true'))
+        print(token)
+
         return jsonify({'status': 'Ok', 'user_id': _id, 'token': token, 'is_admin': is_admin, 'comp_id': comp_id}), 200
 
     return jsonify({'status': 'Bad credentials'}), 403
@@ -87,8 +89,8 @@ def reset_password():
     user_id = dict_data['user_id']
     new_password = dict_data['new_password']
     token = dict_data['token']
-
-    if token != current_app.config["AUTH_TOKEN"] and token != current_app.config["ADMIN_AUTH_TOKEN"]:
+    is_valid, _payload = validate_token(token)
+    if not is_valid:
         return jsonify({'status': 'Unauthorised'}), 403
 
     encrypted_password = encrypt_password(new_password, DES_KEY)
@@ -133,14 +135,15 @@ def signup():
         'user_id': user_id,
         'comp_id': comp_id
     })
+    token: str = issue_token(user_id=user_id, comp_id=comp_id, is_admin=True)
     if isinstance(result, int):
         return jsonify(
             {
                 'status': 'Ok',
                 'comp_id': comp_id,
                 'user_id': user_id,
                 'is_admin': True,
-                'token': current_app.config["ADMIN_AUTH_TOKEN"]
+                'token': token
             }
         ), 200
     else:
@@ -151,7 +154,8 @@ def new_employee():
     ''' Create new employee function '''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorized'}), 403
     result = dbc.execute_query('create_user_employee', args={
         'username': dict_data['username'],
@@ -168,7 +172,8 @@ def retire():
     ''' Retire function, delete company and all employees '''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorized'}), 403
     result = dbc.execute_query(query='delete_sales_by_comp_id', args=dict_data['comp_id'])
     if result is False:
@@ -195,7 +200,8 @@ def delete_employee():
     ''' Delete employee function '''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorized'}), 403
     result = dbc.execute_query('delete_user_by_id', dict_data['employee_id'])
     if result is True:

server/api/clients/routes.py

@@ -1,6 +1,6 @@
-from flask import Flask, Blueprint, request, jsonify, current_app
+from flask import Blueprint, request, jsonify
 from db.db_connector import DBConnector
-import json
+from api.auth.jwt_utils import validate_token
 
 clients = Blueprint('clients', __name__)
 
@@ -9,7 +9,11 @@ def list_clients():
     ''' List clients function'''
     dbc = DBConnector()
     dict_data = request.get_json()
-    comp_id = dbc.execute_query(query='get_compnay_id_by_user', args=dict_data['user_id'])
+    token = dict_data['token']
+    is_valid, payload = validate_token(token)
+    if not is_valid or not payload.get('is_admin'):
+        return jsonify({'status': 'Unauthorised'}), 403
+    comp_id = dbc.execute_query(query='get_compnay_id_by_user', args=payload['user_id'])
     results = dbc.execute_query(query='get_clients_list', args=comp_id)
     if isinstance(results, list):
         return jsonify({'status': 'Ok', 'clients': results}), 200
@@ -41,7 +45,8 @@ def delete_client():
     dbc = DBConnector()
     dict_data = request.get_json()
     token = dict_data['token']
-    if token != current_app.config["ADMIN_AUTH_TOKEN"]:
+    is_valid, payload = validate_token(token)
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorised'}), 403
     result = dbc.execute_query(query='delete_client_by_id', args=dict_data['client_id'])
     if isinstance(result, int):

server/api/company/routes.py

@@ -1,9 +1,10 @@
 import os
-from flask import Flask, Blueprint, request, jsonify, current_app, abort, send_file
+from flask import Blueprint, request, jsonify, abort, send_file
 from db.db_connector import DBConnector
 from services.process_file import ProcessFile
 from services.process_cash_flow import ProcessCashFlow
 from services.process_sales     import ProcessSales
+from api.auth.jwt_utils import validate_token
 
 company = Blueprint('company', __name__)
 
@@ -12,7 +13,8 @@ def list_clients():
     ''' List Sales function'''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorised'}), 403
     results = dbc.execute_query(query='get_company_sales', args=dict_data['comp_id'])
     pcf = ProcessCashFlow(dict_data['comp_id'], 'PT', month=7) ### in this case we don't cate about the month, let's give one value just to simplify
@@ -28,7 +30,8 @@ def list_employees():
     ''' List employees function'''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorised'}), 403
     results = dbc.execute_query(query='get_employees_list', args=dict_data['comp_id'])
     if isinstance(results, list):
@@ -40,7 +43,8 @@ def list_products():
     ''' List products for given company '''
     dbc = DBConnector()
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN'] and dict_data['token'] != current_app.config['AUTH_TOKEN']:
+    is_valid, _payload = validate_token(dict_data.get('token'))
+    if not is_valid:
         return jsonify({'status': 'Unauthorised'}), 403
     results = dbc.execute_query(query='get_products_list', args=dict_data['comp_id'])
     if isinstance(results, list):
@@ -74,8 +78,9 @@ def update_commission():
     print(dict_data)
     seller_id = dict_data['seller_id']
     new_commission = dict_data['new_commission']
-    print(token + '======' + current_app.config['ADMIN_AUTH_TOKEN'])
-    if token != current_app.config['ADMIN_AUTH_TOKEN']:
+    # intentionally not comparing with static admin token anymore
+    is_valid, payload = validate_token(token)
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorized'}), 403
     dbc = DBConnector()
     dbc.execute_query(query='update_seller_commission', args={'seller_id': seller_id, 'new_commission':new_commission})
@@ -86,7 +91,8 @@ def upload_excel():
     ''' Update company products from csv or xlsx '''
     token = request.form.get('token')
     comp_id = request.form.get('comp_id')
-    if token != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(token)
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorized'}), 403
     if 'file' not in request.files:
         return jsonify({'error': 'No file part in the request'}), 400
@@ -104,7 +110,8 @@ def upload_excel():
 def cash_flow():
     ''' Calculate compnay's cash flow '''
     dict_data = request.get_json()
-    if dict_data['token'] != current_app.config['ADMIN_AUTH_TOKEN']:
+    is_valid, payload = validate_token(dict_data.get('token'))
+    if not is_valid or not payload.get('is_admin'):
         return jsonify({'status': 'Unauthorised'}), 403
 
     #### To Simplify, we will only use 3 fixed months of sales ####