Fix: base64 encode everything captured from user input

1 files changed, 22 insertions, 5 deletions

diff --git a/templates/python.py b/templates/python.py
index 5768534..4b5f343 100644
--- a/templates/python.py
+++ b/templates/python.py
@@ -31,17 +31,28 @@ class Context:
     def __setitem__(self, key, value):
         self._vars[key] = value
 
+# Decode a base64 encoded string. Result is binary or unicode string.
+
+def decode_bytes(s):
+    return base64.b64decode(s)
+
+def decode_str(s):
+    return base64.b64decode(s).decode()
+
 # Test data files that were embedded in the source document. Base64
 # encoding is used to allow arbitrary data.
 _files = {}
 {% for file in files %}
-_files['''{{ file.filename }}'''] = '''{{ file.contents | base64 }}'''
+# {{ file.filename }}
+filename = decode_str('{{ file.filename | base64 }}')
+contents = decode_bytes('{{ file.contents | base64 }}')
+_files[filename] = contents
 {% endfor %}
 
 
 # Retrieve an embedded test data file using filename.
 def get_file(filename):
-    return base64.b64decode(_files[filename])
+    return _files[filename]
 
 # Check two values for equality and give error if they are not equal
 def assert_eq(a, b):
@@ -68,15 +79,21 @@ os.chdir(_datadir)
 
 {% for scenario in scenarios %}
 ######################################
-print('''scenario: {{ scenario.title }}''')
+# Scenario: {{ scenario.title }}
+title = decode_str('{{ scenario.title | base64 }}')
+print('scenario: {}'.format(title))
 _scendir = tempfile.mkdtemp(dir=_datadir)
 os.chdir(_scendir)
 ctx = Context()
 {% for step in scenario.steps %}
-print('''  step: {{ step.kind | lower }} {{ step.text }}''')
+# Step: {{ step.text }}
+step = decode_str('{{ step.text | base64 }}')
+print('  step: {{ step.kind | lower }} {}'.format(step))
 args = {}
 {% for part in step.parts %}{% if part.CapturedText is defined -%}
-args['''{{ part.CapturedText.name }}'''] = '''{{ part.CapturedText.text }}'''
+name = decode_str('{{ part.CapturedText.name | base64 }}')
+text = decode_str('{{ part.CapturedText.text | base64 }}')
+args[name] = text
 {% endif -%}
 {% endfor -%}
 {{ step.function }}(ctx, **args)