diff --git a/handlers/PublicHandlers.py b/handlers/PublicHandlers.py index 42c28028..a6ec7995 100644 --- 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 index 1a750ab8..1ebf0fed 100644 --- 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) \ No newline at end of file diff --git a/models/Box.py b/models/Box.py index bcdddfc4..b640d964 100644 --- a/models/Box.py +++ b/models/Box.py @@ -22,16 +22,12 @@ 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 @@ -39,7 +35,11 @@ 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 index 2af8548a..0c39b532 100644 --- a/models/Team.py +++ b/models/Team.py @@ -21,8 +21,6 @@ # pylint: disable=no-member -import imghdr -import io import os import xml.etree.cElementTree as ET from builtins import str @@ -30,8 +28,6 @@ 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 index 4641587a..5dcb6f2b 100644 --- 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"""