Update src/web_app.py

src/web_app.py

@@ -6,9 +6,11 @@ Supports image upload, video upload, and live webcam streaming
 from flask import Flask, render_template, request, jsonify, send_file
 from werkzeug.utils import secure_filename
 from pathlib import Path
+import os
 import cv2
 import numpy as np
 import base64
+import logging
 
 from face_detection_yolov12 import YOLOv12FaceDetector, detect_from_video
 
@@ -21,7 +23,12 @@ UPLOAD_FOLDER = PROJECT_ROOT / "data" / "uploads"
 MODELS_DIR = PROJECT_ROOT / "models"
 ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'gif', 'mp4', 'avi', 'mov', 'mkv'}
 MAX_FILE_SIZE = 500 * 1024 * 1024  # 500MB
-
+ALLOWED_MODELS = {
+    'yolov12n-face.pt',
+    'yolov12s-face.pt',
+    'yolov12m-face.pt',
+    'yolov12l-face.pt'
+}
 UPLOAD_FOLDER.mkdir(exist_ok=True)
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
@@ -33,12 +40,33 @@ detector_cache = {}
 
 def get_detector(model_name):
     """Get or create detector instance (cached)"""
-    if model_name not in detector_cache:
-        model_path = MODELS_DIR / model_name
-        if not model_path.exists():
-            raise FileNotFoundError(f"Model not found: {model_path}")
-        detector_cache[model_name] = YOLOv12FaceDetector(str(model_path))
-    return detector_cache[model_name]
+    safe_name = secure_filename(model_name)
+
+    if safe_name not in ALLOWED_MODELS:
+        logging.error(f"Attempt to load unsupported model: {safe_name}")
+        raise ValueError(f"Unsupported model: {safe_name}")
+
+    if safe_name not in detector_cache:
+        model_path = MODELS_DIR / safe_name
+
+        try:
+            final_path = model_path.resolve()
+            safe_root = MODELS_DIR.resolve()
+            if not str(final_path).startswith(str(safe_root)):
+                logging.error(f"Security Alert: Symlink attack detected! {final_path}")
+                raise ValueError("Invalid model path (Symlink violation)")
+                
+        except Exception as e:
+            logging.error(f"Error resolving model path: {str(e)}")
+            raise FileNotFoundError(f"Model path error: {str(e)}")
+
+        if not final_path.exists():
+            logging.error(f"Model file not found: {final_path}")
+            raise FileNotFoundError(f"Model not found: {final_path}")
+
+        detector_cache[safe_name] = YOLOv12FaceDetector(str(final_path))
+        
+    return detector_cache[safe_name]
 
 
 def allowed_file(filename):
@@ -80,6 +108,9 @@ def detect_image():
         
         # Get model selection
         model = request.form.get('model', 'yolov12l-face.pt')
+        if model not in ALLOWED_MODELS:
+            app.logger.info(f"Invalid model '{model}' requested. Fallback to default.")
+            model = 'yolov12l-face.pt'
         
         # Get detector
         detector = get_detector(model)
@@ -130,7 +161,8 @@ def detect_image():
         return jsonify(response)
     
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        logging.exception("Error during image detection")
+        return jsonify({'error': 'Internal server error during image detection'}), 500
 
 
 @app.route('/api/detect-video', methods=['POST'])
@@ -149,7 +181,10 @@ def detect_video():
         
         # Get model selection
         model = request.form.get('model', 'yolov12m-face.pt')
-        
+        if model not in ALLOWED_MODELS:
+            app.logger.info(f"Invalid model '{model}' requested. Fallback to default.")
+            model = 'yolov12m-face.pt'
+            
         # Save uploaded file
         filename = secure_filename(file.filename)
         input_path = UPLOAD_FOLDER / f"input_{filename}"
@@ -175,7 +210,9 @@ def detect_video():
         return jsonify(response)
     
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        # Log the full exception server-side without exposing details to the client
+        app.logger.exception("Error while processing video detection request")
+        return jsonify({'error': 'Internal server error'}), 500
 
 
 @app.route('/api/download/<filename>', methods=['GET'])
@@ -188,7 +225,9 @@ def download_file(filename):
         
         return send_file(filepath, as_attachment=True)
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        # Log the full exception server-side without exposing details to the client
+        app.logger.exception("Error while processing download request for %s", filename)
+        return jsonify({'error': 'Internal server error'}), 500
 
 
 @app.route('/api/models', methods=['GET'])
@@ -268,11 +307,14 @@ if __name__ == '__main__':
     print("   GET  /api/models          - Get available models")
     print("   GET  /api/health          - Health check")
     print("\n" + "="*70 + "\n")
+
+    # Determine debug mode from environment (default: disabled)
+    debug_mode = os.getenv("FLASK_ENV") == "development"
     
     # Run Flask app
     app.run(
         host='0.0.0.0',
         port=7860,
-        debug=True,
+        debug=debug_mode,
         use_reloader=False
-    )
+    )