consolidate avatar saving and validation

consolidate the avatar saving code in XSSImageCheck - Box, Team, User all had their own implementations, combined these into one `save_avatar` - new `avatar_validation` function Registration: - check avatar validation as part of form_validation function. Previously if a user provided a bad image the user/team would be created but then it would fail at avatar creation/saving....leaving them in a bad state and unable to play
moloch-- · Jul 12, 2024 · 0f50491 · 0f50491
1 parent 0b70b3a
commit 0f50491
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 102 deletions.
diff --git a/handlers/PublicHandlers.py b/handlers/PublicHandlers.py
@@ -64,7 +64,10 @@
     send_user_registered_webhook,
     send_user_validated_webhook,
 )
-from libs.XSSImageCheck import filter_avatars
+from libs.XSSImageCheck import (
+    filter_avatars,
+    avatar_validation,    
+)
 from models import azuread_app
 from models.EmailToken import EmailToken
 from models.GameLevel import GameLevel
@@ -510,7 +513,9 @@ def form_validation(self):
             raise ValidationError("Passwords do not match")
         if self.config.use_recaptcha and self.verify_recaptcha() is False:
             raise ValidationError("Invalid reCAPTCHA")
-
+        if hasattr(self.request, "files") and "avatar" in self.request.files:
+            avatar_validation(self.request.files["avatar"][0]["body"])
+
     def verify_recaptcha(self):
         """Checks recaptcha"""
         recaptcha_response = self.get_argument("g-recaptcha-response", None)

diff --git a/libs/XSSImageCheck.py b/libs/XSSImageCheck.py
@@ -11,12 +11,19 @@
 
 """
 
+import io
 import os
 from random import randint, sample
 from string import printable
+from pathlib import Path
+import imghdr
 
+from PIL import Image
+from resizeimage import resizeimage
 from tornado.options import options
 
+from libs.ValidationError import ValidationError
+
 MAX_AVATAR_SIZE = 1024 * 1024
 MIN_AVATAR_SIZE = 64
 IMG_FORMATS = ["png", "jpeg", "jpg", "gif", "bmp"]
@@ -87,3 +94,46 @@ def existing_avatars(dir):
             if user.avatar is not None:
                 avatars.append(user.avatar)
     return avatars
+
+
+def avatar_validation(image_data) -> str:
+    """Avatar validation check
+    
+    Returns image extension as str if checks pass
+    """
+    if MIN_AVATAR_SIZE < len(image_data) < MAX_AVATAR_SIZE:
+        ext = imghdr.what("", h=image_data)
+        if ext in IMG_FORMATS and not is_xss_image(image_data):
+            return ext
+        else:
+            raise ValidationError(
+                "Invalid image format, avatar must be: %s"
+                % (", ".join(IMG_FORMATS))
+            )                
+
+    else:
+        raise ValidationError(
+            "The image is too large must be %d - %d bytes"
+            % (MIN_AVATAR_SIZE, MAX_AVATAR_SIZE)
+        )
+
+def save_avatar(path: str, image_data: bytes) -> str:
+    """
+    Save avatar image to path
+    
+    Returns image path without avatar_dir
+    """
+    try:
+        base_path = Path(path)
+        image_path = os.path.join(options.avatar_dir, base_path)
+
+        if os.path.exists(image_path):
+            os.unlink(image_path)
+
+        image = Image.open(io.BytesIO(image_data))
+        cover = resizeimage.resize_cover(image, [500, 250])
+        cover.save(image_path, image.format)
+        return str(base_path)
+
+    except Exception as e:
+        raise ValidationError(e)
diff --git a/models/Box.py b/models/Box.py
@@ -22,24 +22,24 @@
 
 import binascii
 import enum
-import imghdr
-import io
 import os
 import xml.etree.cElementTree as ET
 from collections import OrderedDict
 from os import urandom
 from uuid import uuid4
 
-from PIL import Image
-from resizeimage import resizeimage
 from sqlalchemy import Column, ForeignKey
 from sqlalchemy.orm import backref, relationship
 from sqlalchemy.types import Boolean, Enum, Integer, String, Unicode
 from tornado.options import options
 
 from libs.StringCoding import decode, encode
 from libs.ValidationError import ValidationError
-from libs.XSSImageCheck import get_new_avatar, is_xss_image
+from libs.XSSImageCheck import (
+    get_new_avatar,
+    avatar_validation,
+    save_avatar
+)
 from models import dbsession
 from models.BaseModels import DatabaseObject
 from models.Category import Category
@@ -352,28 +352,12 @@ def avatar(self, image_data):
 
         if self.uuid is None:
             self.uuid = str(uuid4())
-        if len(image_data) < (1024 * 1024):
-            ext = imghdr.what("", h=image_data)
-            if ext in ["png", "jpeg", "gif", "bmp"] and not is_xss_image(image_data):
-                try:                    
-                    if self._avatar is not None:
-                        current_image_path = os.path.join(options.avatar_dir, avatar_path, self._avatar) if avatar_path == "upload" else os.path.join(options.avatar_dir, avatar_path)
-                        if os.path.exists(current_image_path):
-                            os.unlink(current_image_path)
-
-                    new_image_path = os.path.join(avatar_path, f"{self.uuid}.{ext}") if avatar_path == "upload" else avatar_path
-                    image = Image.open(io.BytesIO(image_data))
-                    cover = resizeimage.resize_cover(image, [500, 250])
-                    cover.save(os.path.join(options.avatar_dir,new_image_path), image.format)
-                    self._avatar = new_image_path
-                except Exception as e:
-                    raise ValidationError(e)
-            else:
-                raise ValidationError(
-                    "Invalid image format, avatar must be: .png .jpeg .gif or .bmp"
-                )
-        else:
-            raise ValidationError("The image is too large")
+
+        if avatar_path == "upload":
+            os.path.join("upload", f"{self.uuid}.{ext}")
+
+        ext = avatar_validation(image_data)
+        self._avatar = save_avatar(avatar_path)
 
     @property
     def ipv4s(self):

diff --git a/models/Team.py b/models/Team.py
@@ -21,17 +21,13 @@
 # pylint: disable=no-member
 
 
-import imghdr
-import io
 import os
 import xml.etree.cElementTree as ET
 from builtins import str
 from datetime import datetime
 from random import randint
 from uuid import uuid4
 
-from PIL import Image
-from resizeimage import resizeimage
 from sqlalchemy import Column, desc
 from sqlalchemy.orm import backref, relationship
 from sqlalchemy.types import Integer, String, Unicode
@@ -41,11 +37,9 @@
 from libs.StringCoding import encode
 from libs.ValidationError import ValidationError
 from libs.XSSImageCheck import (
-    IMG_FORMATS,
-    MAX_AVATAR_SIZE,
-    MIN_AVATAR_SIZE,
-    get_new_avatar,
-    is_xss_image,
+    get_new_avatar, 
+    avatar_validation,
+    save_avatar,
 )
 from models import dbsession
 from models.BaseModels import DatabaseObject
@@ -263,34 +257,8 @@ def avatar(self):
 
     @avatar.setter
     def avatar(self, image_data):
-        if MIN_AVATAR_SIZE < len(image_data) < MAX_AVATAR_SIZE:
-            ext = imghdr.what("", h=image_data)
-            if ext in IMG_FORMATS and not is_xss_image(image_data):
-                try:
-                    if self._avatar is not None and os.path.exists(
-                        options.avatar_dir + "/upload/" + self._avatar
-                    ):
-                        os.unlink(options.avatar_dir + "/upload/" + self._avatar)
-                    file_path = str(
-                        options.avatar_dir + "/upload/" + self.uuid + "." + ext
-                    )
-                    image = Image.open(io.BytesIO(image_data))
-                    cover = resizeimage.resize_cover(image, [500, 250])
-                    cover.save(file_path, image.format)
-                    self._avatar = "upload/" + self.uuid + "." + ext
-                except Exception as e:
-                    raise ValidationError(e)
-
-            else:
-                raise ValidationError(
-                    "Invalid image format, avatar must be: %s"
-                    % (", ".join(IMG_FORMATS))
-                )
-        else:
-            raise ValidationError(
-                "The image is too large must be %d - %d bytes"
-                % (MIN_AVATAR_SIZE, MAX_AVATAR_SIZE)
-            )
+        ext = avatar_validation(image_data)
+        self._avatar = save_avatar(os.path.join("upload", f"{self.uuid}.{ext}"),image_data)                
 
     @property
     def levels(self):

diff --git a/models/User.py b/models/User.py
@@ -25,8 +25,6 @@
 """
 
 
-import imghdr
-import io
 import os
 import random
 import string
@@ -39,8 +37,6 @@
 
 from past.builtins import basestring
 from pbkdf2 import PBKDF2
-from PIL import Image
-from resizeimage import resizeimage
 from sqlalchemy import Column, ForeignKey, desc, func
 from sqlalchemy.orm import backref, relationship, synonym
 from sqlalchemy.types import Boolean, DateTime, Integer, String, Unicode
@@ -50,12 +46,10 @@
 from libs.ValidationError import ValidationError
 from libs.WebhookHelpers import send_user_validated_webhook
 from libs.XSSImageCheck import (
-    IMG_FORMATS,
-    MAX_AVATAR_SIZE,
-    MIN_AVATAR_SIZE,
+    avatar_validation,
+    save_avatar,
     default_avatar,
-    get_new_avatar,
-    is_xss_image,
+    get_new_avatar,    
 )
 from models import dbsession
 from models.BaseModels import DatabaseObject
@@ -354,33 +348,8 @@ def avatar(self):
 
     @avatar.setter
     def avatar(self, image_data):
-        if MIN_AVATAR_SIZE < len(image_data) < MAX_AVATAR_SIZE:
-            ext = imghdr.what("", h=image_data)
-            if ext in IMG_FORMATS and not is_xss_image(image_data):
-                try:
-                    if self._avatar is not None and os.path.exists(
-                        options.avatar_dir + "/upload/" + self._avatar
-                    ):
-                        os.unlink(options.avatar_dir + "/upload/" + self._avatar)
-                    file_path = str(
-                        options.avatar_dir + "/upload/" + self.uuid + "." + ext
-                    )
-                    image = Image.open(io.BytesIO(image_data))
-                    cover = resizeimage.resize_cover(image, [500, 250])
-                    cover.save(file_path, image.format)
-                    self._avatar = "upload/" + self.uuid + "." + ext
-                except Exception as e:
-                    raise ValidationError(e)
-            else:
-                raise ValidationError(
-                    "Invalid image format, avatar must be: %s"
-                    % (", ".join(IMG_FORMATS))
-                )
-        else:
-            raise ValidationError(
-                "The image is too large must be %d - %d bytes"
-                % (MIN_AVATAR_SIZE, MAX_AVATAR_SIZE)
-            )
+        ext = avatar_validation(image_data)
+        self._avatar = save_avatar(os.path.join("upload", f"{self.uuid}.{ext}"),image_data)                
 
     def has_item(self, item_name):
         """Check to see if a team has purchased an item"""