diff --git a/examples/YOLOv8_server_api/model_downloader.py b/examples/YOLOv8_server_api/model_downloader.py index 2122d14..c143112 100644 --- a/examples/YOLOv8_server_api/model_downloader.py +++ b/examples/YOLOv8_server_api/model_downloader.py @@ -5,6 +5,7 @@ import requests from flask import Blueprint +from flask import current_app from flask import jsonify from flask import send_from_directory from flask_limiter import Limiter @@ -44,13 +45,21 @@ def download_model(model_name): response.headers['Last-Modified'], '%a, %d %b %Y %H:%M:%S GMT', ) - local_file_path = f"{MODELS_DIRECTORY}{model_name}" + local_file_path = Path(MODELS_DIRECTORY) / model_name + # Check local file's last modified time - local_last_modified = datetime.datetime.fromtimestamp( - Path(local_file_path).stat().st_mtime, - ) - if local_last_modified >= server_last_modified: - return jsonify({'message': 'Local model is up-to-date.'}), 304 + if local_file_path.exists(): + local_last_modified = datetime.datetime.fromtimestamp( + local_file_path.stat().st_mtime, + ) + + # If local file is up-to-date, return 304 Not Modified + if local_last_modified >= server_last_modified: + return jsonify( + { + 'message': 'Local model is up-to-date.', + }, + ), 304 # If not up-to-date, fetch the file and return it return send_from_directory( @@ -60,12 +69,6 @@ def download_model(model_name): except FileNotFoundError: return jsonify({'error': 'Model not found.'}), 404 except requests.RequestException as e: - return ( - jsonify( - { - 'error': 'Failed to fetch model information.', - 'details': str(e), - }, - ), - 500, - ) + # Log the detailed error internally + current_app.logger.error(f'Error fetching model information: {e}') + return jsonify({'error': 'Failed to fetch model information.'}), 500 diff --git a/examples/user_management/app.py b/examples/user_management/app.py index 68d0b3f..2dda615 100644 --- a/examples/user_management/app.py +++ b/examples/user_management/app.py @@ -1,5 +1,6 @@ from __future__ import annotations +from flask import escape from flask import Flask from flask import request @@ -23,21 +24,45 @@ @app.route('/add_user', methods=['POST']) -def add_user_route() -> str: +def add_user_route() -> tuple[str, int]: """ Endpoint to add a new user to the database. Returns: - str: Result message indicating success or failure of the operation. + tuple[str, int]: Result message indicating success + or failure of the operation and the status code. """ - username = request.form['username'] - password = request.form['password'] - result = add_user(username, password) - return result + try: + username = request.form['username'] + password = request.form['password'] + + # Validate input + if not username or not password: + return 'Username and password are required', 400 + + # Escape user input to prevent XSS + username = escape(username) + password = escape(password) + + result = add_user(username, password) + + # Check if the addition was successful + if result: + return 'User added successfully', 200 + else: + return 'Failed to add user', 500 + + except KeyError: + return 'Invalid input', 400 + except Exception as e: + # Log the exception for internal tracking + # And do not expose the exception message to the user + app.logger.error(f'Error adding user: {e}') + return 'An internal error occurred', 500 @app.route('/delete_user/', methods=['DELETE']) -def delete_user_route(username: str) -> str: +def delete_user_route(username: str) -> tuple[str, int]: """ Endpoint to delete an existing user from the database. @@ -45,38 +70,106 @@ def delete_user_route(username: str) -> str: username (str): Username of the user to be deleted. Returns: - str: Result message indicating success or failure of the operation. + tuple[str, int]: Result message indicating success + or failure of the operation and the status code. """ - result = delete_user(username) - return result + try: + # Validate input + if not username: + return 'Username is required', 400 + + # Escape user input to prevent XSS + username = escape(username) + + result = delete_user(username) + + # Check if the deletion was successful + if result: + return 'User deleted successfully', 200 + else: + return 'Failed to delete user', 500 + + except Exception as e: + # Log the exception for internal tracking + # And do not expose the exception message to the user + app.logger.error(f'Error deleting user {username}: {e}') + return 'An internal error occurred', 500 @app.route('/update_username', methods=['PUT']) -def update_username_route() -> str: +def update_username_route() -> tuple[str, int]: """ Endpoint to update a user's username in the database. Returns: - str: Result message indicating success or failure of the operation. + tuple[str, int]: Result message indicating success + or failure of the operation and the status code. """ - old_username = request.form['old_username'] - new_username = request.form['new_username'] - result = update_username(old_username, new_username) - return result + try: + old_username = request.form['old_username'] + new_username = request.form['new_username'] + + # Validate input + if not old_username or not new_username: + return 'Old username and new username are required', 400 + + # Escape user input to prevent XSS + old_username = escape(old_username) + new_username = escape(new_username) + + result = update_username(old_username, new_username) + + # Check if the update was successful + if result: + return 'Username updated successfully', 200 + else: + return 'Failed to update username', 500 + + except KeyError: + return 'Invalid input', 400 + except Exception as e: + # Log the exception for internal tracking + # And do not expose the exception message to the user + app.logger.error(f'Error updating username: {e}') + return 'An internal error occurred', 500 @app.route('/update_password', methods=['PUT']) -def update_password_route() -> str: +def update_password_route() -> tuple[str, int]: """ Endpoint to update a user's password in the database. Returns: - str: Result message indicating success or failure of the operation. + tuple[str, int]: Result message indicating success + or failure of the operation and the status code. """ - username = request.form['username'] - new_password = request.form['new_password'] - result = update_password(username, new_password) - return result + try: + username = request.form['username'] + new_password = request.form['new_password'] + + # Validate input + if not username or not new_password: + return 'Username and new password are required', 400 + + # Escape user input to prevent XSS + username = escape(username) + new_password = escape(new_password) + + result = update_password(username, new_password) + + # Check if the update was successful + if result: + return 'Password updated successfully', 200 + else: + return 'Failed to update password', 500 + + except KeyError: + return 'Invalid input', 400 + except Exception as e: + # Log the exception for internal tracking + # And do not expose the exception message to the user + app.logger.error(f'Error updating password: {e}') + return 'An internal error occurred', 500 if __name__ == '__main__': diff --git a/src/wechat_notifier.py b/src/wechat_notifier.py index 0392a79..e8f2544 100644 --- a/src/wechat_notifier.py +++ b/src/wechat_notifier.py @@ -126,7 +126,14 @@ def main(): message = 'Hello, WeChat!' image = np.zeros((100, 100, 3), dtype=np.uint8) # Example image (black) response = notifier.send_notification(user_id, message, image=image) - print(response) + + # Remove sensitive data from response before logging + sanitised_response = { + k: v for k, v in response.items() if k not in [ + 'access_token', 'corp_secret', + ] + } + print(sanitised_response) if __name__ == '__main__':