Shift allow attrs check to be after evaluating parent node (to get

type), and fix tests based on that.

Author: danthedeckie

simpleeval.py

@@ -726,17 +726,8 @@ def _eval_subscript(self, node):
         return container[key]
 
     def _eval_attribute(self, node):
-        if self.allowed_attrs is not None:
-            allowed_attrs = self.allowed_attrs.get(type(node.value.value), TypeNotSpecified)
-            if allowed_attrs == TypeNotSpecified:
-                raise FeatureNotAvailable(
-                    f"Sorry, attribute access not allowed on '{type(node.value.value)}'"
-                )
-            if node.attr not in allowed_attrs:
-                raise FeatureNotAvailable(
-                    f"Sorry, '{node.attr}' access not allowed on '{type(node.value.value)}'"
-                )
-
+        # DISALLOW_PREFIXES & DISALLOW_METHODS are global, there's never any access to
+        # attrs with these names, so we can bail early:
         for prefix in DISALLOW_PREFIXES:
             if node.attr.startswith(prefix):
                 raise FeatureNotAvailable(
@@ -748,9 +739,27 @@ def _eval_attribute(self, node):
             raise FeatureNotAvailable(
                 "Sorry, this method is not available. " "({0})".format(node.attr)
             )
-        # eval node
+
+        # Evaluate "node" - the thing that we're trying to access an attr of first:
         node_evaluated = self._eval(node.value)
 
+        # If we've opted in to the 'allowed_attrs' checking per type, then since we now
+        # know what kind of node we've got, we can check if we're permitted to access this
+        # attr name on this node:
+        if self.allowed_attrs is not None:
+            type_to_check = type(node_evaluated)
+
+            allowed_attrs = self.allowed_attrs.get(type_to_check, TypeNotSpecified)
+            if allowed_attrs == TypeNotSpecified:
+                raise FeatureNotAvailable(
+                    f"Sorry, attribute access not allowed on '{type_to_check}'"
+                    f" (attempted to access `.{node.attr}`)"
+                )
+            if node.attr not in allowed_attrs:
+                raise FeatureNotAvailable(
+                    f"Sorry, '{node.attr}' access not allowed on '{type_to_check}'"
+                )
+
         # Maybe the base object is an actual object, not just a dict
         try:
             return getattr(node_evaluated, node.attr)

test_simpleeval.py

@@ -1353,11 +1353,32 @@ def test_no_operators(self):
 
 
 class TestAllowedAttributes(DRYTest):
+    def setUp(self):
+        self.saved_disallow_methods = simpleeval.DISALLOW_METHODS
+        simpleeval.DISALLOW_METHODS = []
+        super().setUp()
+
+    def tearDown(self) -> None:
+        simpleeval.DISALLOW_METHODS = self.saved_disallow_methods
+        return super().tearDown()
+
     def test_allowed_attrs_(self):
         self.s.allowed_attrs = BASIC_ALLOWED_ATTRS
         self.t("5 + 5", 10)
         self.t('"   hello  ".strip()', "hello")
 
+    def test_allowed_extra_attr(self):
+        class Foo:
+            def bar(self):
+                return 42
+
+        assert Foo().bar() == 42
+
+        extended_attrs = BASIC_ALLOWED_ATTRS.copy()
+        extended_attrs[Foo] = {"bar"}
+
+        simple_eval("foo.bar()", names={"foo": Foo()}, allowed_attrs=extended_attrs)
+
     def test_breakout_via_generator(self):
         # Thanks decorator-factory
         class Foo:
@@ -1369,11 +1390,11 @@ def bar(self):
 
         evil = "foo.bar().gi_frame.f_globals['__builtins__'].exec('raise RuntimeError(\"Oh no\")')"
 
-        x = simpleeval.DISALLOW_METHODS
-        simpleeval.DISALLOW_METHODS = []
-        with self.assertRaises(FeatureNotAvailable):
-            simple_eval(evil, names={"foo": Foo()}, allowed_attrs=BASIC_ALLOWED_ATTRS)
-        simpleeval.DISALLOW_METHODS = x
+        extended_attrs = BASIC_ALLOWED_ATTRS.copy()
+        extended_attrs[Foo] = {"bar"}
+
+        with self.assertRaisesRegex(FeatureNotAvailable, r".*attempted to access `\.gi_frame`.*"):
+            simple_eval(evil, names={"foo": Foo()}, allowed_attrs=extended_attrs)
 
 
 if __name__ == "__main__":  # pragma: no cover