Check that uploaded images are valid images

1 files changed, 23 insertions, 3 deletions

diff --git a/app/utils.py b/app/utils.py
index 21c7fd6..034c39c 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -20,7 +20,7 @@ from flask_user import *
 from flask_login import login_user, logout_user
 from app.models import *
 from app import app
-import random, string, os
+import random, string, os, imghdr
 
 def getExtension(filename):
 	return filename.rsplit(".", 1)[1].lower() if "." in filename else None
@@ -28,6 +28,10 @@ def getExtension(filename):
 def isFilenameAllowed(filename, exts):
 	return getExtension(filename) in exts
 
+ALLOWED_IMAGES = set(["jpeg", "png"])
+def isAllowedImage(data):
+	return imghdr.what(None, data) in ALLOWED_IMAGES
+
 def shouldReturnJson():
 	return "application/json" in request.accept_mimetypes and \
 			not "text/html" in request.accept_mimetypes
@@ -36,16 +40,32 @@ def randomString(n):
 	return ''.join(random.choice(string.ascii_lowercase + \
 			string.ascii_uppercase + string.digits) for _ in range(n))
 
-def doFileUpload(file, allowedExtensions, fileTypeName):
+def doFileUpload(file, fileType, fileTypeDesc):
 	if not file or file is None or file.filename == "":
 		flash("No selected file", "error")
 		return None
 
+	allowedExtensions = []
+	isImage = False
+	if fileType == "image":
+		allowedExtensions = ["jpg", "jpeg", "png"]
+		isImage = True
+	elif filetype == "zip":
+		allowedExtensions = ["zip"]
+	else:
+		raise Exception("Invalid fileType")
+
 	ext = getExtension(file.filename)
 	if ext is None or not ext in allowedExtensions:
-		flash("Please upload load " + fileTypeName, "error")
+		flash("Please upload load " + fileTypeDesc, "danger")
 		return None
 
+	if isImage and not isAllowedImage(file.stream.read()):
+		flash("Uploaded image isn't actually an image", "danger")
+		return None
+
+	file.stream.seek(0)
+
 	filename = randomString(10) + "." + ext
 	file.save(os.path.join("app/public/uploads", filename))
 	return "/uploads/" + filename