3 files changed, 11 insertions, 8 deletions

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from schema import And, Schema
 
 from Options import Toggle, Choice, DefaultOnToggle, Range, PerGameCommonOptions, StartInventoryPool, OptionDict
-from worlds.lingo.items import TRAP_ITEMS
+from .items import TRAP_ITEMS
 
 
 class ShuffleDoors(Choice):
@@ -78,13 +78,16 @@ def get_progressive_item_id(name: str):
 def load_static_data_from_file():
     global PAINTING_ENTRANCES, PAINTING_EXITS
 
+    from . import datatypes
+    from Utils import safe_builtins
+
     class RenameUnpickler(pickle.Unpickler):
         def find_class(self, module, name):
-            renamed_module = module
-            if module == "datatypes":
-                renamed_module = "worlds.lingo.datatypes"
-
-            return super(RenameUnpickler, self).find_class(renamed_module, name)
+            if module in ("worlds.lingo.datatypes", "datatypes"):
+                return getattr(datatypes, name)
+            elif module == "builtins" and name in safe_builtins:
+                return getattr(safe_builtins, name)
+            raise pickle.UnpicklingError(f"global '{module}.{name}' is forbidden")
 
     file = pkgutil.get_data(__name__, os.path.join("data", "generated.dat"))
     pickdata = RenameUnpickler(BytesIO(file)).load()
@@ -1,8 +1,8 @@
 import os
 import unittest
 
-from worlds.lingo.static_logic import HASHES
-from worlds.lingo.utils.pickle_static_data import hash_file
+from ..static_logic import HASHES
+from ..utils.pickle_static_data import hash_file
 
 
 class TestDatafile(unittest.TestCase):