diff --git a/backend/Generator/main.py b/backend/Generator/main.py index 04aed79f..42799dac 100644 --- a/backend/Generator/main.py +++ b/backend/Generator/main.py @@ -22,6 +22,8 @@ import os import fitz import mammoth +import zipfile +from werkzeug.utils import secure_filename class MCQGenerator: @@ -368,19 +370,63 @@ def extract_text_from_docx(self, file_path): return result.value def process_file(self, file): - file_path = os.path.join(self.upload_folder, file.filename) + if not file.filename: + raise ValueError("Empty filename") + + filename = secure_filename(file.filename) + if not filename: + raise ValueError("Invalid filename") + + target_dir = os.path.abspath(self.upload_folder) + file_path = os.path.abspath(os.path.join(target_dir, filename)) + + if os.path.commonpath([target_dir, file_path]) != target_dir: + raise ValueError("Path traversal detected") + file.save(file_path) content = "" - if file.filename.endswith('.txt'): - with open(file_path, 'r') as f: - content = f.read() - elif file.filename.endswith('.pdf'): - content = self.extract_text_from_pdf(file_path) - elif file.filename.endswith('.docx'): - content = self.extract_text_from_docx(file_path) + try: + if '.' not in filename: + raise ValueError("Unsupported file extension") + ext = filename.rsplit('.', 1)[1].lower() + + with open(file_path, 'rb') as f: + header = f.read(512) + + if ext == 'pdf': + if not header.startswith(b"%PDF"): + raise ValueError("Invalid file content: PDF signature not found") + content = self.extract_text_from_pdf(file_path) + + elif ext == 'docx': + if not header.startswith(b"PK\x03\x04"): + raise ValueError("Invalid file content: DOCX signature not found") + + try: + with zipfile.ZipFile(file_path, 'r') as zf: + namelist = zf.namelist() + if '[Content_Types].xml' not in namelist or 'word/document.xml' not in namelist: + raise ValueError("Invalid file content: Missing DOCX internal structures") + except zipfile.BadZipFile: + raise ValueError("Invalid file content: Not a valid ZIP archive") + + content = self.extract_text_from_docx(file_path) + + elif ext == 'txt': + try: + header.decode('utf-8') + except UnicodeDecodeError: + raise ValueError("Invalid file content: TXT must be valid UTF-8 text") + with open(file_path, 'r', encoding='utf-8') as f: + content = f.read() + else: + raise ValueError("Unsupported file extension") + + finally: + if os.path.exists(file_path): + os.remove(file_path) - os.remove(file_path) return content class QuestionGenerator: diff --git a/backend/server.py b/backend/server.py index 683c1241..2658512c 100644 --- a/backend/server.py +++ b/backend/server.py @@ -495,12 +495,19 @@ def upload_file(): if file.filename == '': return jsonify({"error": "No selected file"}), 400 - content = file_processor.process_file(file) - - if content: - return jsonify({"content": content}) - else: - return jsonify({"error": "Unsupported file type or error processing file"}), 400 + try: + content = file_processor.process_file(file) + + if content: + return jsonify({"content": content}) + else: + return jsonify({"error": "Unsupported file type or error processing file"}), 400 + except ValueError as e: + app.logger.warning("Upload validation failed: %s", e) + return jsonify({"error": str(e)}), 400 + except Exception as e: + app.logger.exception("Error processing upload: %s", e) + return jsonify({"error": "Internal server error"}), 500 @app.route("/", methods=["GET"]) def hello():