From e565b6f6f551ca51f6967375e5f12eceafbbdc4d Mon Sep 17 00:00:00 2001
From: Tuukka Mustonen <tuukka.mustonen@f-secure.com>
Date: Fri, 25 Oct 2019 17:44:51 +0300
Subject: [PATCH] Add 'strict' mode

This raises TemplateExecutionException if reference does not exist in content.
It's even stricter than Velocity's "strict mode", that accepts undefined
variables (in special cases) in #if blocks.
---
 airspeed/__init__.py   | 35 ++++++++++++++++++++++++++++-------
 tests/airspeed_test.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/airspeed/__init__.py b/airspeed/__init__.py
index cf35738..c3ac1a9 100755
--- a/airspeed/__init__.py
+++ b/airspeed/__init__.py
@@ -88,18 +88,20 @@ def __init__(self, content, filename="<string>"):
         self.filename = filename
         self.root_element = None
 
-    def merge(self, namespace, loader=None):
+    def merge(self, namespace, loader=None, options=None):
         output = StoppableStream()
-        self.merge_to(namespace, output, loader)
+        self.merge_to(namespace, output, loader, options)
         return output.getvalue()
 
     def ensure_compiled(self):
         if not self.root_element:
             self.root_element = TemplateBody(self.filename, self.content)
 
-    def merge_to(self, namespace, fileobj, loader=None):
+    def merge_to(self, namespace, fileobj, loader=None, options=None):
         if loader is None:
-            loader = NullLoader()
+            loader = NullLoader(options=options)
+        elif options:
+            loader.options.update(options)
         self.ensure_compiled()
         self.root_element.evaluate(fileobj, namespace, loader)
 
@@ -157,7 +159,15 @@ def element_name(self):
             self.element.__class__.__name__).strip()
 
 
+class UndefinedVariable(Exception):
+    def __init__(self, name):
+        super(Exception, self).__init__("Variable '{}' is not defined or is None!".format(name))
+
+
 class NullLoader:
+    def __init__(self, options=None):
+        self.options = options or {}
+
     def load_text(self, name):
         raise TemplateError("no loader available for '%s'" % name)
 
@@ -166,12 +176,13 @@ def load_template(self, name):
 
 
 class CachingFileLoader:
-    def __init__(self, basedir, debugging=False):
+    def __init__(self, basedir, debugging=False, options=None):
         self.basedir = basedir
         self.known_templates = {}  # name -> (template, file_mod_time)
         self.debugging = debugging
         if debugging:
             print("creating caching file loader with basedir:", basedir)
+        self.options = options or {}
 
     def filename_of(self, name):
         return os.path.join(self.basedir, name)
@@ -604,7 +615,10 @@ def calculate(self, current_object, loader, top_namespace):
             if methods_for_type and self.name in methods_for_type:
                 result = lambda *args: methods_for_type[self.name](current_object, *args)
         if result is None:
-            return None  # TODO: an explicit 'not found' exception?
+            if loader.options.get('strict', False):
+                raise UndefinedVariable(self.name)
+            else:
+                return None
         if self.parameters is not None:
             result = result(*self.parameters.calculate(top_namespace, loader))
         elif self.index is not None:
@@ -715,7 +729,14 @@ def parse(self):
     def evaluate_raw(self, stream, namespace, loader):
         value = None
         if self.expression is not None:
-            value = self.expression.calculate(namespace, loader)
+            try:
+                value = self.expression.calculate(namespace, loader)
+            except UndefinedVariable:
+                # Allow silent variables to not raise
+                if self.silent:
+                    pass
+                else:
+                    raise
         if value is None:
             if self.silent and self.expression is not None:
                 value = ''
diff --git a/tests/airspeed_test.py b/tests/airspeed_test.py
index a4e8491..a742252 100644
--- a/tests/airspeed_test.py
+++ b/tests/airspeed_test.py
@@ -14,6 +14,8 @@
 
 import six
 
+STRICT = {'strict': True}
+
 
 class TemplateTestCase(TestCase):
     def assertRaisesExecutionError(self, exctype, func, *args, **kwargs):
@@ -38,6 +40,38 @@ def test_dollar_left_untouched(self):
         template = airspeed.Template("Hello $")
         self.assertEquals("Hello $", template.merge({}))
 
+    def test_strict_mode(self):
+        def nok(tpl, content=None):
+            with self.assertRaises(airspeed.TemplateExecutionError):
+                airspeed.Template(tpl).merge(content or {}, options=STRICT)
+
+        def ok(tpl, out, content=None):
+            self.assertEquals(out, airspeed.Template(tpl).merge(content or {}, options=STRICT))
+
+        nok("$undefined")
+        nok("${undefined}")
+        ok("$!undefined $!{undefined}", " ")
+        ok("$!undefined", "")
+        ok("$!defined", "1", {"defined": 1})
+        ok("#set($foo = 1)", "")
+        nok("#set($foo = $bar)")
+        ok("#set($foo = $bar)", "", {"bar": 1})
+        ok("#if(false)$undefined#end", "")
+
+        # These would work with Velocity's *standard* strict mode (to allow checking if defined)
+        nok("#if ($foo)#end")
+        nok("#if ( ! $foo)#end")
+        nok("#if ($foo && $foo.bar)#end")
+        nok("#if ($foo && $foo == 'bar')#end")
+        nok("#if ($foo1 || $foo2)#end")
+
+        # Workaround, to check for undefined, is via helper variable
+        content = {}
+        content['__exists'] = lambda x: x in content
+        ok("#if($__exists('undefined'))yes#end", "", content)
+        content['undefined'] = 1
+        ok("#if($__exists('undefined'))yes#end", "yes", content)
+
     def test_unmatched_name_does_not_get_substituted(self):
         template = airspeed.Template("Hello $name")
         self.assertEquals("Hello $name", template.merge({}))