WebHost: working web-gen

Utils.py

@@ -24,7 +24,7 @@ import pickle
 import functools
 import io
 import collections
-
+import importlib
 from yaml import load, dump, safe_load
 
 try:
@@ -365,16 +365,25 @@ safe_builtins = {
 
 
 class RestrictedUnpickler(pickle.Unpickler):
+    def __init__(self, *args, **kwargs):
+        super(RestrictedUnpickler, self).__init__(*args, **kwargs)
+        self.options_module = importlib.import_module("Options")
+        self.net_utils_module = importlib.import_module("NetUtils")
+
     def find_class(self, module, name):
         if module == "builtins" and name in safe_builtins:
             return getattr(builtins, name)
+        # used by MultiServer -> savegame/multidata
         if module == "NetUtils" and name in {"NetworkItem", "ClientStatus", "Hint"}:
-            import NetUtils
+            return getattr(self.net_utils_module, name)
-            return getattr(NetUtils, name)
+        # Options are unpickled by WebHost -> Generate
-        if module == "Options":
+        if module.endswith("Options"):
-            import Options
+            if module == "Options":
-            obj = getattr(Options, name)
+                mod = self.options_module
-            if issubclass(obj, Options.Option):
+            else:
+                mod = importlib.import_module(module)
+            obj = getattr(mod, name)
+            if issubclass(obj, self.options_module.Option):
                 return obj
         # Forbid everything else.
         raise pickle.UnpicklingError("global '%s.%s' is forbidden" %

WebHostLib/generate.py

@@ -2,6 +2,7 @@ import os
 import tempfile
 import random
 import json
+import zipfile
 from collections import Counter
 
 from flask import request, flash, redirect, url_for, session, render_template
@@ -15,6 +16,7 @@ import pickle
 from .models import *
 from WebHostLib import app
 from .check import get_yaml_data, roll_options
+from .upload import upload_zip_to_db
 
 
 @app.route('/generate', methods=['GET', 'POST'])
@@ -69,7 +71,7 @@ def gen_game(gen_options, race=False, owner=None, sid=None):
         if race:
             random.seed()  # reset to time-based random source
 
-        seedname = "M" + (f"{random.randint(0, pow(10, seeddigits) - 1)}".zfill(seeddigits))
+        seedname = "W" + (f"{random.randint(0, pow(10, seeddigits) - 1)}".zfill(seeddigits))
 
         erargs = parse_arguments(['--multi', str(playercount)])
         erargs.seed = seed
@@ -84,7 +86,10 @@ def gen_game(gen_options, race=False, owner=None, sid=None):
         for player, (playerfile, settings) in enumerate(gen_options.items(), 1):
             for k, v in settings.items():
                 if v is not None:
-                    getattr(erargs, k)[player] = v
+                    if hasattr(erargs, k):
+                        getattr(erargs, k)[player] = v
+                    else:
+                        setattr(erargs, k, {player: v})
 
             if not erargs.name[player]:
                 erargs.name[player] = os.path.splitext(os.path.split(playerfile)[-1])[0]
@@ -92,7 +97,7 @@ def gen_game(gen_options, race=False, owner=None, sid=None):
 
         ERmain(erargs, seed)
 
-        return upload_to_db(target.name, owner, sid, race)
+        return upload_to_db(target.name, sid, owner, race)
     except BaseException as e:
         if sid:
             with db_session:
@@ -122,37 +127,19 @@ def wait_seed(seed: UUID):
     return render_template("waitSeed.html", seed_id=seed_id)
 
 
-def upload_to_db(folder, owner, sid, race:bool):
+def upload_to_db(folder, sid, owner, race):
-    slots = set()
-    spoiler = ""
-
-    multidata = None
     for file in os.listdir(folder):
         file = os.path.join(folder, file)
-        if file.endswith(".apbp"):
+        if file.endswith(".zip"):
-            player_text = file.split("_P", 1)[1]
+            with db_session:
-            player_name = player_text.split("_", 1)[1].split(".", 1)[0]
+                with zipfile.ZipFile(file) as zfile:
-            player_id = int(player_text.split(".", 1)[0].split("_", 1)[0])
+                    res = upload_zip_to_db(zfile, owner, {"race": race}, sid)
-            slots.add(Slot(data=open(file, "rb").read(),
+                if type(res) == "str":
-                             player_id=player_id, player_name = player_name, game = "A Link to the Past"))
+                    raise Exception(res)
-        elif file.endswith(".txt"):
+                elif res:
-            spoiler = open(file, "rt", encoding="utf-8-sig").read()
+                    seed = res
-        elif file.endswith(".archipelago"):
+                    gen = Generation.get(id=seed.id)
-            multidata = open(file, "rb").read()
+                    if gen is not None:
-    if multidata:
+                        gen.delete()
-        with db_session:
+                    return seed.id
-            if sid:
+    raise Exception("Generation zipfile not found.")
-                seed = Seed(multidata=multidata, spoiler=spoiler, slots=slots, owner=owner,
-                            id=sid, meta=json.dumps({"race": race, "tags": ["generated"]}))
-            else:
-                seed = Seed(multidata=multidata, spoiler=spoiler, slots=slots, owner=owner,
-                            meta=json.dumps({"race": race, "tags": ["generated"]}))
-            for patch in slots:
-                patch.seed = seed
-            if sid:
-                gen = Generation.get(id=sid)
-                if gen is not None:
-                    gen.delete()
-        return seed.id
-    else:
-        raise Exception("Multidata required (.archipelago), but not found.")

WebHostLib/upload.py

@@ -3,6 +3,7 @@ import lzma
 import json
 import base64
 import MultiServer
+import uuid
 
 from flask import request, flash, redirect, url_for, session, render_template
 from pony.orm import flush, select
@@ -17,6 +18,68 @@ accepted_zip_contents = {"patches": ".apbp",
 banned_zip_contents = (".sfc",)
 
 
+def upload_zip_to_db(zfile: zipfile.ZipFile, owner=None, meta={"race": False}, sid=None):
+    if not owner:
+        owner = session["_id"]
+    infolist = zfile.infolist()
+    slots = set()
+    spoiler = ""
+    multidata = None
+    for file in infolist:
+        if file.filename.endswith(banned_zip_contents):
+            return "Uploaded data contained a rom file, which is likely to contain copyrighted material. " \
+                   "Your file was deleted."
+        elif file.filename.endswith(".apbp"):
+            data = zfile.open(file, "r").read()
+            yaml_data = parse_yaml(lzma.decompress(data).decode("utf-8-sig"))
+            if yaml_data["version"] < 2:
+                return "Old format cannot be uploaded (outdated .apbp)", 500
+            metadata = yaml_data["meta"]
+            slots.add(Slot(data=data, player_name=metadata["player_name"],
+                           player_id=metadata["player_id"],
+                           game="A Link to the Past"))
+
+        elif file.filename.endswith(".apmc"):
+            data = zfile.open(file, "r").read()
+            metadata = json.loads(base64.b64decode(data).decode("utf-8"))
+            slots.add(Slot(data=data, player_name=metadata["player_name"],
+                           player_id=metadata["player_id"],
+                           game="Minecraft"))
+
+        elif file.filename.endswith(".zip"):
+            # Factorio mods need a specific name or they do not function
+            _, seed_name, slot_id, slot_name = file.filename.rsplit("_", 1)[0].split("-")
+            slots.add(Slot(data=zfile.open(file, "r").read(), player_name=slot_name,
+                           player_id=int(slot_id[1:]), game="Factorio"))
+
+        elif file.filename.endswith(".apz5"):
+            # .apz5 must be named specifically since they don't contain any metadata
+            _, seed_name, slot_id, slot_name = file.filename.split('.')[0].split('_', 3)
+            slots.add(Slot(data=zfile.open(file, "r").read(), player_name=slot_name,
+                           player_id=int(slot_id[1:]), game="Ocarina of Time"))
+
+        elif file.filename.endswith(".txt"):
+            spoiler = zfile.open(file, "r").read().decode("utf-8-sig")
+        elif file.filename.endswith(".archipelago"):
+            try:
+                multidata = zfile.open(file).read()
+                MultiServer.Context._decompress(multidata)
+            except:
+                flash("Could not load multidata. File may be corrupted or incompatible.")
+            else:
+                multidata = zfile.open(file).read()
+    if multidata:
+        flush()  # commit slots
+        seed = Seed(multidata=multidata, spoiler=spoiler, slots=slots, owner=owner, meta=json.dumps(meta),
+                    id=sid if sid else uuid.uuid4())
+        flush()  # create seed
+        for slot in slots:
+            slot.seed = seed
+        return seed
+    else:
+        flash("No multidata was found in the zip file, which is required.")
+
+
 @app.route('/uploads', methods=['GET', 'POST'])
 def uploads():
     if request.method == 'POST':
@@ -31,64 +94,12 @@ def uploads():
                 flash('No selected file')
             elif file and allowed_file(file.filename):
                 if file.filename.endswith(".zip"):
-                    slots = set()
-                    spoiler = ""
-                    multidata = None
                     with zipfile.ZipFile(file, 'r') as zfile:
-                        infolist = zfile.infolist()
+                        res = upload_zip_to_db(zfile)
-
+                        if type(res) == str:
-                        for file in infolist:
+                            return res
-                            if file.filename.endswith(banned_zip_contents):
+                        elif res:
-                                return "Uploaded data contained a rom file, which is likely to contain copyrighted material. Your file was deleted."
+                            return redirect(url_for("viewSeed", seed=res.id))
-                            elif file.filename.endswith(".apbp"):
-                                data = zfile.open(file, "r").read()
-                                yaml_data = parse_yaml(lzma.decompress(data).decode("utf-8-sig"))
-                                if yaml_data["version"] < 2:
-                                    return "Old format cannot be uploaded (outdated .apbp)", 500
-                                metadata = yaml_data["meta"]
-                                slots.add(Slot(data=data, player_name=metadata["player_name"],
-                                               player_id=metadata["player_id"],
-                                               game="A Link to the Past"))
-
-                            elif file.filename.endswith(".apmc"):
-                                data = zfile.open(file, "r").read()
-                                metadata = json.loads(base64.b64decode(data).decode("utf-8"))
-                                slots.add(Slot(data=data, player_name=metadata["player_name"],
-                                               player_id=metadata["player_id"],
-                                               game="Minecraft"))
-
-                            elif file.filename.endswith(".zip"):
-                                # Factorio mods needs a specific name or they do not function
-                                _, seed_name, slot_id, slot_name = file.filename.rsplit("_", 1)[0].split("-")
-                                slots.add(Slot(data=zfile.open(file, "r").read(), player_name=slot_name,
-                                              player_id=int(slot_id[1:]), game="Factorio"))
-
-                            elif file.filename.endswith(".apz5"):
-                                # .apz5 must be named specifically since they don't contain any metadata
-                                _, seed_name, slot_id, slot_name = file.filename.split('.')[0].split('_', 3)
-                                slots.add(Slot(data=zfile.open(file, "r").read(), player_name=slot_name,
-                                              player_id=int(slot_id[1:]), game="Ocarina of Time"))
-
-                            elif file.filename.endswith(".txt"):
-                                spoiler = zfile.open(file, "r").read().decode("utf-8-sig")
-                            elif file.filename.endswith(".archipelago"):
-                                try:
-                                    multidata = zfile.open(file).read()
-                                    MultiServer.Context._decompress(multidata)
-                                except:
-                                    flash("Could not load multidata. File may be corrupted or incompatible.")
-                                else:
-                                    multidata = zfile.open(file).read()
-                        if multidata:
-                            flush()  # commit slots
-                            seed = Seed(multidata=multidata, spoiler=spoiler, slots=slots, owner=session["_id"])
-                            flush()  # create seed
-                            for slot in slots:
-                                slot.seed = seed
-
-                            return redirect(url_for("viewSeed", seed=seed.id))
-                        else:
-                            flash("No multidata was found in the zip file, which is required.")
                 else:
                     try:
                         multidata = file.read()