Boost security

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

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,60 +24,152 @@
 
 
 @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/<username>', 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.
 
     Args:
         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__':

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__':