From 89f9b6d2708bfed0b1bf681b4c68818f5db85bcb Mon Sep 17 00:00:00 2001 From: Hugo H Date: Sun, 31 Aug 2025 11:22:41 +0100 Subject: [PATCH] Moved permissions checking to a function Permissions checking now happens in another function, code is now much cleaner --- main.py | 167 +++++++++++++++++++++++++++----------------------------- 1 file changed, 79 insertions(+), 88 deletions(-) diff --git a/main.py b/main.py index 0b9b94b..c516086 100644 --- a/main.py +++ b/main.py @@ -44,15 +44,9 @@ except Exception as e: app = Flask(__name__) -# Chat Details Endpoint: -# Get or change details about a chat using the chatId -# Arguments: token (required), details (required), model, name -@app.route('/api/chat/<_id>/details', methods = ['GET', 'POST']) -def getChatHistory(_id): - # Get user auth token - token = request.json['token'] +def checkUserPermission(token, permission): # Find the correct user token in user db - user = usersCollection.find_one({'tokens.token': token}, {"_id":1,"tokens":{"$elemMatch": {"token":token}}}) + user = usersCollection.find_one({'tokens.token': token}, {"_id":1,"tokens":{"$elemMatch": {"token":token}}, "permissions":1}) # If the user exists, continue, otherwise return fail if (user): # Convert _id to a string, python doesn't like ObjectId() @@ -61,44 +55,61 @@ def getChatHistory(_id): if (user['tokens'][0]['expiry'] > int(datetime.now().timestamp())): # Store the userId userId = user['_id'] - print(userId) - # Get the request details - details = request.json['details'] - # If the user is trying to GET data - if (request.method == 'GET'): - # Get the chat from the chatId - returnedChat = chatCollection.find_one({'_id': ObjectId(_id)}) - # Convert chatId into string - returnedChat['_id'] = str(returnedChat['_id']) - try: - returnedChat["permissions"][userId].index("view") - print("Chat " + _id + " has been found with token " + token) - # Check for detail type and return correct value from db - if (details == "history"): - return jsonify(returnedChat["messages"]) - elif (details == "users"): - return jsonify(returnedChat["permissions"]) - elif (details == "model"): - return jsonify(returnedChat["model"]) - elif (details == "name"): - return jsonify(returnedChat["name"]) - except: - return jsonify("Invalid Permissions") + if permission in user["permissions"]: + return True, userId + elif (permission == True): + return True, userId else: - try: - returnedChat["permissions"][userId].index("edit") - # Check for the detail type and add data to db - if (details == "model"): - model = request.json['model'] - chatCollection.update_one({'_id': ObjectId(_id)}, { "$set": { "model": model } }) - if (details == "name"): - name = request.json['name'] - chatCollection.update_one({'_id': ObjectId(_id)}, { "$set": { "name": name } }) - return jsonify("Success") - except: - return jsonify("Invalid Permissions") + return False, "Incorrect permissions" else: - return jsonify("User token is invalid") + return False, "Token is expired" + else: + return False, "Token doesn't exist" + +# Chat Details Endpoint: +# Get or change details about a chat using the chatId +# Arguments: token (required), details (required), model, name +@app.route('/api/chat/<_id>/details', methods = ['GET', 'POST']) +def getChatHistory(_id): + # Get user auth token + token = request.json['token'] + a, userId = checkUserPermission(token, True) + if (a == True): + # Get the request details + details = request.json['details'] + # If the user is trying to GET data + if (request.method == 'GET'): + # Get the chat from the chatId + returnedChat = chatCollection.find_one({'_id': ObjectId(_id)}) + # Convert chatId into string + returnedChat['_id'] = str(returnedChat['_id']) + try: + returnedChat["permissions"][userId].index("view") + print("Chat " + _id + " has been found with token " + token) + # Check for detail type and return correct value from db + if (details == "history"): + return jsonify(returnedChat["messages"]) + elif (details == "users"): + return jsonify(returnedChat["permissions"]) + elif (details == "model"): + return jsonify(returnedChat["model"]) + elif (details == "name"): + return jsonify(returnedChat["name"]) + except: + return jsonify("Invalid Permissions") + else: + try: + returnedChat["permissions"][userId].index("edit") + # Check for the detail type and add data to db + if (details == "model"): + model = request.json['model'] + chatCollection.update_one({'_id': ObjectId(_id)}, { "$set": { "model": model } }) + if (details == "name"): + name = request.json['name'] + chatCollection.update_one({'_id': ObjectId(_id)}, { "$set": { "name": name } }) + return jsonify("Success") + except: + return jsonify("Invalid Permissions") else: return jsonify("User token is invalid") @@ -109,44 +120,28 @@ def getChatHistory(_id): def createChat(): # Get user auth token token = request.json['token'] - # Find the correct user token in user db - user = usersCollection.find_one({'tokens.token': token}, {"_id":1,"tokens":{"$elemMatch": {"token":token}}}) - # If the user exists, continue, otherwise return fail - if (user): - # Convert _id to a string, python doesn't like ObjectId() - user['_id'] = str(user['_id']) - # Check if the token expiry is after the current date (Using unix timestamp, other mongodb Date datatype is a pain to use in python) - if (user['tokens'][0]['expiry'] > int(datetime.now().timestamp())): - # Store the userId - userId = user['_id'] - print(user) - print(user['permissions']) - if ("createChat" in user['permissions']): - print(userId) - name = request.json['name'] - model = request.json['model'] - chatCollection.insert_one( - { - "name":name, - "model":model, - "permissions": { - userId:[ - "owner", - "view", - "message", - "edit" - ] - }, - "messages": [ - - ] - } - ) - return jsonify("Success") - else: - return jsonify("Incorrect permissions") - else: - return jsonify("User token is invalid") + a, userId = checkUserPermission(token, "createChat") + if (a == True): + name = request.json['name'] + model = request.json['model'] + chatCollection.insert_one( + { + "name":name, + "model":model, + "permissions": { + userId:[ + "owner", + "view", + "message", + "edit" + ] + }, + "messages": [ + + ] + } + ) + return jsonify("Success") else: return jsonify("User token is invalid") @@ -167,13 +162,9 @@ def index(): if (token == 'none'): return render_template('login.html', appName=appName, githubUrl=github_auth_endpoint, githublogin=settings["github_oauth"]["enabled"], oauthlogin=settings["oauth_login"]) else: - user = usersCollection.find_one({'tokens.token': token}, {"_id":1,"tokens":{"$elemMatch": {"token":token}}}) - if (user): - user['_id'] = str(user['_id']) - if (user['tokens'][0]['expiry'] > int(datetime.now().timestamp())): + a, userId = checkUserPermission(token, True) + if (a == True): return render_template('home.html', appName=appName) - else: - render_template('logout.html', appName=appName) else: render_template('logout.html', appName=appName)