Simplified whitelisting of python modules,

...and added information how to disable safe python

data/ais/bruteforce.py

@@ -1,6 +1,6 @@
 #!WPY
 
-#import wesnoth,random
+import wesnoth, random
 
 ##    Copyright 2006 by Michael Schmahl
 ##    This code is available under the latest version of the GNU Public License.

data/ais/parse.py

@@ -1,7 +1,8 @@
 import re, os, safe
 
-whitelisted = ["wesnoth", "heapq", "random"]
+whitelisted = ["wesnoth", "heapq", "random", "math", "string", "re"]
 rex = re.compile(r"^import\s+(.*)", re.M)
+modules = {}
 
 def include(matchob):
     """
@@ -11,7 +12,9 @@ def include(matchob):
     names = [x.strip() for x in matchob.group(1).split(",")]
     r = ""
     for name in names:
-        if name in whitelisted: continue
+        if name in whitelisted:
+            modules[name] = __import__(name)
+            continue
         for path in pathes:
             includefile = os.path.join(path, name)
             try:
@@ -36,8 +39,12 @@ def parse_file(name):
     code = rex.sub(include, code)
     return code
 
+# If you want to disable safe python, use this instead:
+#
+# def parse(name): return open(name).read(), {}
 def parse(name):
-    global already
+    global already, modules
     already = {}
-    return parse_file(name)
+    modules = {}
+    return parse_file(name), modules
 

data/ais/safe.py

@@ -124,7 +124,10 @@ def safe_run(code,context=None):
         _builtin_restore()
         raise
 
-def safe_exec(code,context = None):
+# If you want to disable safe python, use this instead:
+#
+# def safe_exec(code, context = None): exec code in context
+def safe_exec(code, context = None):
     """Check the code to be safe, then run it with only safe builtins on."""
     safe_check(code)
     safe_run(code,context)

src/ai_python.cpp

@@ -1791,12 +1791,10 @@ void python_ai::play_turn()
 		"\tbackup = sys.path[:]\n"
 		"\tsys.path.append(\"" + path + "/data/ais\")\n"
 		"\ttry:\n"
-		"\t\timport wesnoth, parse, safe, heapq, random\n"
-		"\t\tcode = parse.parse(\"" + script + "\")\n"
-		"\t\tsafe.safe_exec(code, {\n"
-		"\t\t\"wesnoth\" : wesnoth,\n"
-		"\t\t\"heapq\" : heapq,\n"
-		"\t\t\"random\" : random})\n"
+		"\t\timport parse, safe\n"
+		"\t\tparse.pathes = [\"" + path + "\"]\n"
+		"\t\tcode, context = parse.parse(\"" + script + "\")\n"
+		"\t\tsafe.safe_exec(code, context)\n"
 		"\texcept:\n"
 		"\t\terr = str(traceback.format_exc())\n"
 		"\t\traise\n"