From 6b08abbada99183def718310f4854c31c558ad07 Mon Sep 17 00:00:00 2001
From: Elias Pschernig <elias@pschernig.at>
Date: Fri, 16 Nov 2007 16:07:45 +0000
Subject: [PATCH] Simplified whitelisting of python modules,

...and added information how to disable safe python
---
 data/ais/bruteforce.py |  2 +-
 data/ais/parse.py      | 15 +++++++++++----
 data/ais/safe.py       |  5 ++++-
 src/ai_python.cpp      | 10 ++++------
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/data/ais/bruteforce.py b/data/ais/bruteforce.py
index 55ca354d43b..8be65558fec 100644
--- a/data/ais/bruteforce.py
+++ b/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.
diff --git a/data/ais/parse.py b/data/ais/parse.py
index c89090fe322..dc9c7453a33 100644
--- a/data/ais/parse.py
+++ b/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
 
diff --git a/data/ais/safe.py b/data/ais/safe.py
index e3c9b931ff6..3f2d599ebc1 100644
--- a/data/ais/safe.py
+++ b/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)
diff --git a/src/ai_python.cpp b/src/ai_python.cpp
index 70b04a54c0a..8a2c8295602 100644
--- a/src/ai_python.cpp
+++ b/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"