Enhance upload module with docstring, implement username sanitization, and improve error handling for file uploads.

This commit is contained in:
voldemort 2025-07-23 16:46:00 +07:00
parent 78d59b295c
commit 7f84542cfb

View File

@ -1,23 +1,36 @@
from flask import Blueprint, request, jsonify, session """Module to handle the upload process."""
import os import os
import logging import logging
import re
from flask import Blueprint, request, jsonify, session
upload_bp = Blueprint("upload_bp", __name__) upload_bp = Blueprint("upload_bp", __name__)
def sanitize_username(username):
"""Sanitize the username."""
return re.sub(r"[^a-zA-Z0-9_\-]", "_", username).lower()
@upload_bp.route("/upload/<cdmtype>", methods=["POST"]) @upload_bp.route("/upload/<cdmtype>", methods=["POST"])
def upload(cdmtype): def upload(cdmtype):
"""Handle the upload process."""
try: try:
username = session.get("username") username = session.get("username")
if not username: if not username:
return jsonify({"message": "False", "error": "No username in session"}), 400 return jsonify({"message": "False", "error": "No username in session"}), 400
safe_username = sanitize_username(username)
# Validate CDM type # Validate CDM type
if cdmtype not in ["PR", "WV"]: if cdmtype not in ["PR", "WV"]:
return jsonify({"message": "False", "error": "Invalid CDM type"}), 400 return jsonify({"message": "False", "error": "Invalid CDM type"}), 400
# Set up user directory paths # Set up user directory paths
base_path = os.path.join(os.getcwd(), "configs", "CDMs", username) base_path = os.path.join(
os.getcwd(), "configs", "CDMs", "users_uploaded", safe_username
)
pr_path = os.path.join(base_path, "PR") pr_path = os.path.join(base_path, "PR")
wv_path = os.path.join(base_path, "WV") wv_path = os.path.join(base_path, "WV")
@ -32,11 +45,13 @@ def upload(cdmtype):
# Determine correct save path based on cdmtype # Determine correct save path based on cdmtype
filename = uploaded_file.filename filename = uploaded_file.filename
save_path = os.path.join(pr_path if cdmtype == "PR" else wv_path, filename) assert filename is not None
target_path = pr_path if cdmtype == "PR" else wv_path
save_path = os.path.join(target_path, filename)
uploaded_file.save(save_path) uploaded_file.save(save_path)
return jsonify({"message": "Success", "file_saved_to": save_path}) return jsonify({"message": "Success", "file_saved_to": save_path})
except Exception as e: except (OSError, IOError, ValueError, AttributeError) as e:
logging.exception("Upload failed") logging.exception("Upload failed: %s", {e})
return jsonify({"message": "False", "error": "Server error"}), 500 return jsonify({"message": "False", "error": "Server error"}), 500