feat: Allow users to set custom profile pictures

mslib/mscolab/conf.py

@@ -76,6 +76,12 @@ class default_mscolab_settings:
     # used to generate the password token
     SECURITY_PASSWORD_SALT = secrets.token_urlsafe(16)
 
+    # Max allowed file size for uploading user profile image
+    MAX_IMAGE_SIZE = 1 * 1024 * 1024  # 1MB
+
+    # Allowed file extensions for uploading user profile image
+    ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
+
     STUB_CODE = """<?xml version="1.0" encoding="utf-8"?>
     <FlightTrack version="1.7.6">
       <ListOfWaypoints>

mslib/mscolab/server.py

@@ -197,6 +197,11 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def allowed_file(filename):
+    """ Check if the uploaded file is in the allowed set of extensions """
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in mscolab_settings.ALLOWED_EXTENSIONS
+
+
 def get_idp_entity_id(selected_idp):
     """
     Finds the entity_id from the configured IDPs
@@ -352,17 +357,37 @@ def get_user():
 
 
 @APP.route('/upload_profile_image', methods=["POST"])
+@verify_user
 def upload_profile_image():
     user_id = request.form['user_id']
     file = request.files['image']
-    if file:
-        success, message = fm.save_user_profile_image(user_id, file.read())
-        if success:
-            return jsonify({'message': message}), 200
+    if file and allowed_file(file.filename):
+        if file.mimetype.startswith('image/'):
+            try:
+                data = file.read()
+                if len(data) > mscolab_settings.MAX_IMAGE_SIZE:
+                    return jsonify({'message': 'File too large'}), 413
+                success, message = fm.save_user_profile_image(user_id, data)
+                if success:
+                    return jsonify({'message': message}), 200
+                else:
+                    return jsonify({'message': message}), 400
+            except Exception as e:
+                return jsonify({'message': str(e)}), 500
         else:
-            return jsonify({'message': message}), 400
+            return jsonify({'message': 'Invalid file type'}), 400
+    else:
+        return jsonify({'message': 'No file provided or invalid file type'}), 400
+
+
+@APP.route('/fetch_profile_image', methods=["GET"])
+def fetch_profile_image():
+    user_id = request.form['user_id']
+    image_data = fm.fetch_user_profile_image(user_id)
+    if image_data:
+        return Response(image_data, mimetype='image/png')
     else:
-        return jsonify({'message': 'No file provided'}), 400
+        return 404
 
 
 @APP.route("/delete_own_account", methods=["POST"])

mslib/msui/mscolab.py

@@ -52,7 +52,6 @@
 from mslib.msui import mscolab_admin_window as maw
 from mslib.msui import mscolab_version_history as mvh
 from mslib.msui import socket_control as sc
-from mslib.mscolab.file_manager import FileManager
 
 import PyQt5
 from PyQt5 import QtCore, QtGui, QtWidgets
@@ -699,25 +698,24 @@ def after_login(self, emailid, url, r):
 
     def fetch_gravatar(self, refresh=False):
         # Display custom profile picture if exists
-        from mslib.mscolab.server import APP
-        with APP.app_context():
-            fm = FileManager(APP.config["MSCOLAB_DATA_DIR"])
-            img_data = fm.fetch_user_profile_image(self.user["id"])
-            if img_data:
-                img_byte_arr = io.BytesIO(img_data)
-                pixmap = QPixmap()
-                pixmap.loadFromData(img_byte_arr.getvalue())
-
-                if hasattr(self, 'profile_dialog'):
-                    self.profile_dialog.gravatarLabel.setPixmap(pixmap)
-
-                # set icon for user options toolbutton
-                icon = QtGui.QIcon()
-                icon.addPixmap(pixmap, QtGui.QIcon.Normal, QtGui.QIcon.Off)
-                self.ui.userOptionsTb.setIcon(icon)
-                return
+        url = urljoin(self.mscolab_server_url, 'fetch_profile_image')
+        data = {'user_id': str(self.user["id"])}
+        response = requests.get(url, data=data)
+        if response.status_code == 200:
+            img_data = response.content
+            img_byte_arr = io.BytesIO(img_data)
+            pixmap = QPixmap()
+            pixmap.loadFromData(img_byte_arr.getvalue())
+            if hasattr(self, 'profile_dialog'):
+                self.profile_dialog.gravatarLabel.setPixmap(pixmap)
+            icon = QtGui.QIcon()
+            icon.addPixmap(pixmap, QtGui.QIcon.Normal, QtGui.QIcon.Off)
+            self.ui.userOptionsTb.setIcon(icon)
+        else:
+            self.display_default_gravatar(refresh)
 
-        # Display default gravatar if custom pfp is not set
+    def display_default_gravatar(self, refresh):
+        # Display default gravatar if custom profile image is not set
         email_hash = hashlib.md5(bytes(self.email.encode('utf-8')).lower()).hexdigest()
         email_in_config = self.email in config_loader(dataset="gravatar_ids")
         gravatar_img_path = fs.path.join(constants.GRAVATAR_DIR_PATH, f"{email_hash}.png")
@@ -835,7 +833,7 @@ def on_context_menu(point):
         self.fetch_gravatar()
 
     def upload_image(self):
-        file_name, _ = QFileDialog.getOpenFileName(self.prof_diag, "Open Image", "", "Image Files (*.png *.jpg *.bmp)")
+        file_name, _ = QFileDialog.getOpenFileName(self.prof_diag, "Open Image", "", "Image Files (*.png *.jpg *.jpeg)")
         if file_name:
             # Load, resize and display the image
             image = Image.open(file_name)
@@ -852,7 +850,10 @@ def upload_image(self):
             # Prepare data for upload
             img_byte_arr.seek(0)
             files = {'image': ('profile_image.png', img_byte_arr, 'image/png')}
-            data = {'user_id': str(self.user["id"])}
+            data = {
+                "user_id": str(self.user["id"]),
+                "token": self.token
+            }
 
             # Sending the request
             try: