WebHost: Some refactors and additional checks when uploading files. (#2549)

WebHostLib/upload.py

@@ -19,7 +19,22 @@ from worlds.Files import AutoPatchRegister
 from . import app
 from .models import Seed, Room, Slot, GameDataPackage
 
-banned_zip_contents = (".sfc", ".z64", ".n64", ".sms", ".gb")
+banned_extensions = (".sfc", ".z64", ".n64", ".nes", ".smc", ".sms", ".gb", ".gbc", ".gba")
+allowed_options_extensions = (".yaml", ".json", ".yml", ".txt", ".zip")
+allowed_generation_extensions = (".archipelago", ".zip")
+
+
+def allowed_options(filename: str) -> bool:
+    return filename.endswith(allowed_options_extensions)
+
+
+def allowed_generation(filename: str) -> bool:
+    return filename.endswith(allowed_generation_extensions)
+
+
+def banned_file(filename: str) -> bool:
+    return filename.endswith(banned_extensions)
+
 
 def process_multidata(compressed_multidata, files={}):
     decompressed_multidata = MultiServer.Context.decompress(compressed_multidata)
@@ -61,8 +76,8 @@ def upload_zip_to_db(zfile: zipfile.ZipFile, owner=None, meta={"race": False}, s
     if not owner:
         owner = session["_id"]
     infolist = zfile.infolist()
-    if all(file.filename.endswith((".yaml", ".yml")) or file.is_dir() for file in infolist):
-        flash(Markup("Error: Your .zip file only contains .yaml files. "
+    if all(allowed_options(file.filename) or file.is_dir() for file in infolist):
+        flash(Markup("Error: Your .zip file only contains options files. "
                      'Did you mean to <a href="/generate">generate a game</a>?'))
         return
 
@@ -73,7 +88,7 @@ def upload_zip_to_db(zfile: zipfile.ZipFile, owner=None, meta={"race": False}, s
     # Load files.
     for file in infolist:
         handler = AutoPatchRegister.get_handler(file.filename)
-        if file.filename.endswith(banned_zip_contents):
+        if banned_file(file.filename):
             return "Uploaded data contained a rom file, which is likely to contain copyrighted material. " \
                    "Your file was deleted."
 
@@ -136,35 +151,34 @@ def upload_zip_to_db(zfile: zipfile.ZipFile, owner=None, meta={"race": False}, s
         flash("No multidata was found in the zip file, which is required.")
 
 
-@app.route('/uploads', methods=['GET', 'POST'])
+@app.route("/uploads", methods=["GET", "POST"])
 def uploads():
-    if request.method == 'POST':
-        # check if the post request has the file part
-        if 'file' not in request.files:
-            flash('No file part')
+    if request.method == "POST":
+        # check if the POST request has a file part.
+        if "file" not in request.files:
+            flash("No file part in POST request.")
         else:
-            file = request.files['file']
-            # if user does not select file, browser also
-            # submit an empty part without filename
-            if file.filename == '':
-                flash('No selected file')
-            elif file and allowed_file(file.filename):
-                if zipfile.is_zipfile(file):
-                    with zipfile.ZipFile(file, 'r') as zfile:
+            uploaded_file = request.files["file"]
+            # If the user does not select file, the browser will still submit an empty string without a file name.
+            if uploaded_file.filename == "":
+                flash("No selected file.")
+            elif uploaded_file and allowed_generation(uploaded_file.filename):
+                if zipfile.is_zipfile(uploaded_file):
+                    with zipfile.ZipFile(uploaded_file, "r") as zfile:
                         try:
                             res = upload_zip_to_db(zfile)
                         except VersionException:
                             flash(f"Could not load multidata. Wrong Version detected.")
                         else:
-                            if type(res) == str:
+                            if res is str:
                                 return res
                             elif res:
                                 return redirect(url_for("view_seed", seed=res.id))
                 else:
-                    file.seek(0)  # offset from is_zipfile check
+                    uploaded_file.seek(0)  # offset from is_zipfile check
                     # noinspection PyBroadException
                     try:
-                        multidata = file.read()
+                        multidata = uploaded_file.read()
                         slots, multidata = process_multidata(multidata)
                     except Exception as e:
                         flash(f"Could not load multidata. File may be corrupted or incompatible. ({e})")
@@ -182,7 +196,3 @@ def user_content():
     rooms = select(room for room in Room if room.owner == session["_id"])
     seeds = select(seed for seed in Seed if seed.owner == session["_id"])
     return render_template("userContent.html", rooms=rooms, seeds=seeds)
-
-
-def allowed_file(filename):
-    return filename.endswith(('.archipelago', ".zip"))