fix(iast): forbbidenfruit package conflict (backport #4546) (#4549)

Remove forbiddenfruit as dependency and rollback wrapt changes where forbiddenfruit was called.
IAST: Patch builtins only when IAST is enabled.


Co-authored-by: Alberto Vara <alberto.vara@datadoghq.com>

ddtrace/appsec/iast/_patch.py

@@ -0,0 +1,85 @@
+import ctypes
+import gc
+
+from ddtrace.internal.logger import get_logger
+from ddtrace.vendor.wrapt import FunctionWrapper
+from ddtrace.vendor.wrapt import resolve_path
+
+
+log = get_logger(__name__)
+
+
+def try_wrap_function_wrapper(module, name, wrapper):
+    try:
+        wrap_object(module, name, FunctionWrapper, (wrapper,))
+    except (ImportError, AttributeError):
+        log.debug("IAST patching. Module %s.%s not exists", module, name)
+
+
+def apply_patch(parent, attribute, replacement):
+    try:
+        setattr(parent, attribute, replacement)
+    except (TypeError, AttributeError):
+        patch_builtins(parent, attribute, replacement)
+
+
+def wrap_object(module, name, factory, args=(), kwargs={}):
+    (parent, attribute, original) = resolve_path(module, name)
+    wrapper = factory(original, *args, **kwargs)
+    apply_patch(parent, attribute, wrapper)
+    return wrapper
+
+
+def patchable_builtin(klass):
+    refs = gc.get_referents(klass.__dict__)
+    assert len(refs) == 1
+    return refs[0]
+
+
+def patch_builtins(klass, attr, value):
+    """Based on forbiddenfruit package:
+    https://github.com/clarete/forbiddenfruit/blob/master/forbiddenfruit/__init__.py#L421
+    ---
+    Patch a built-in `klass` with `attr` set to `value`
+
+    This function monkey-patches the built-in python object `attr` adding a new
+    attribute to it. You can add any kind of argument to the `class`.
+
+    It's possible to attach methods as class methods, just do the following:
+
+      >>> def myclassmethod(cls):
+      ...     return cls(1.5)
+      >>> curse(float, "myclassmethod", classmethod(myclassmethod))
+      >>> float.myclassmethod()
+      1.5
+
+    Methods will be automatically bound, so don't forget to add a self
+    parameter to them, like this:
+
+      >>> def hello(self):
+      ...     return self * 2
+      >>> curse(str, "hello", hello)
+      >>> "yo".hello()
+      "yoyo"
+    """
+    dikt = patchable_builtin(klass)
+
+    old_value = dikt.get(attr, None)
+    old_name = "_c_%s" % attr  # do not use .format here, it breaks py2.{5,6}
+
+    # Patch the thing
+    dikt[attr] = value
+
+    if old_value:
+        dikt[old_name] = old_value
+
+        try:
+            dikt[attr].__name__ = old_value.__name__
+        except (AttributeError, TypeError):  # py2.5 will raise `TypeError`
+            pass
+        try:
+            dikt[attr].__qualname__ = old_value.__qualname__
+        except AttributeError:
+            pass
+
+    ctypes.pythonapi.PyType_Modified(ctypes.py_object(klass))

ddtrace/appsec/iast/taint_sinks/_base.py

@@ -11,7 +11,6 @@
 from ddtrace.constants import IAST_CONTEXT_KEY
 from ddtrace.internal import _context
 from ddtrace.internal.logger import get_logger
-from ddtrace.vendor.wrapt import wrap_function_wrapper
 
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -77,10 +76,3 @@ def report(cls, evidence_value=""):
                     }
                 )
             _context.set_item(IAST_CONTEXT_KEY, report, span=span)
-
-
-def _wrap_function_wrapper_exception(module, name, wrapper):
-    try:
-        wrap_function_wrapper(module, name, wrapper)
-    except (ImportError, AttributeError):
-        log.debug("IAST patching. Module %s.%s not exists", module, name)

ddtrace/appsec/iast/taint_sinks/weak_hash.py

@@ -3,10 +3,10 @@
 from typing import TYPE_CHECKING
 
 from ddtrace.appsec.iast import oce
+from ddtrace.appsec.iast._patch import try_wrap_function_wrapper
 from ddtrace.appsec.iast.constants import EVIDENCE_ALGORITHM_TYPE
 from ddtrace.appsec.iast.constants import VULN_INSECURE_HASHING_TYPE
 from ddtrace.appsec.iast.taint_sinks._base import VulnerabilityBase
-from ddtrace.appsec.iast.taint_sinks._base import _wrap_function_wrapper_exception
 from ddtrace.internal.logger import get_logger
 
 
@@ -37,22 +37,22 @@ def patch():
     setattr(hashlib, "_datadog_patch", True)
 
     if sys.version_info >= (3, 0, 0):
-        _wrap_function_wrapper_exception("_hashlib", "HASH.digest", wrapped_digest_function)
-        _wrap_function_wrapper_exception("_hashlib", "HASH.hexdigest", wrapped_digest_function)
-        _wrap_function_wrapper_exception(("_%s" % MD5_DEF), "MD5Type.digest", wrapped_md5_function)
-        _wrap_function_wrapper_exception(("_%s" % MD5_DEF), "MD5Type.hexdigest", wrapped_md5_function)
-        _wrap_function_wrapper_exception(("_%s" % SHA1_DEF), "SHA1Type.digest", wrapped_sha1_function)
-        _wrap_function_wrapper_exception(("_%s" % SHA1_DEF), "SHA1Type.hexdigest", wrapped_sha1_function)
+        try_wrap_function_wrapper("_hashlib", "HASH.digest", wrapped_digest_function)
+        try_wrap_function_wrapper("_hashlib", "HASH.hexdigest", wrapped_digest_function)
+        try_wrap_function_wrapper(("_%s" % MD5_DEF), "MD5Type.digest", wrapped_md5_function)
+        try_wrap_function_wrapper(("_%s" % MD5_DEF), "MD5Type.hexdigest", wrapped_md5_function)
+        try_wrap_function_wrapper(("_%s" % SHA1_DEF), "SHA1Type.digest", wrapped_sha1_function)
+        try_wrap_function_wrapper(("_%s" % SHA1_DEF), "SHA1Type.hexdigest", wrapped_sha1_function)
     else:
-        _wrap_function_wrapper_exception("hashlib", MD5_DEF, wrapped_md5_function)
-        _wrap_function_wrapper_exception("hashlib", SHA1_DEF, wrapped_sha1_function)
-        _wrap_function_wrapper_exception("hashlib", "new", wrapped_new_function)
+        try_wrap_function_wrapper("hashlib", MD5_DEF, wrapped_md5_function)
+        try_wrap_function_wrapper("hashlib", SHA1_DEF, wrapped_sha1_function)
+        try_wrap_function_wrapper("hashlib", "new", wrapped_new_function)
 
     # pycryptodome methods
-    _wrap_function_wrapper_exception("Crypto.Hash.MD5", "MD5Hash.digest", wrapped_md5_function)
-    _wrap_function_wrapper_exception("Crypto.Hash.MD5", "MD5Hash.hexdigest", wrapped_md5_function)
-    _wrap_function_wrapper_exception("Crypto.Hash.SHA1", "SHA1Hash.digest", wrapped_sha1_function)
-    _wrap_function_wrapper_exception("Crypto.Hash.SHA1", "SHA1Hash.hexdigest", wrapped_sha1_function)
+    try_wrap_function_wrapper("Crypto.Hash.MD5", "MD5Hash.digest", wrapped_md5_function)
+    try_wrap_function_wrapper("Crypto.Hash.MD5", "MD5Hash.hexdigest", wrapped_md5_function)
+    try_wrap_function_wrapper("Crypto.Hash.SHA1", "SHA1Hash.digest", wrapped_sha1_function)
+    try_wrap_function_wrapper("Crypto.Hash.SHA1", "SHA1Hash.hexdigest", wrapped_sha1_function)
 
 
 @WeakHash.wrap

ddtrace/vendor/wrapt/wrappers.py

@@ -802,10 +802,8 @@ def lookup_attribute(parent, attribute):
                     return vars(cls)[attribute]
             else:
                 return getattr(parent, attribute)
-        elif hasattr(parent, attribute):
-            return getattr(parent, attribute)
         else:
-            return None
+            return getattr(parent, attribute)
 
     original = lookup_attribute(parent, attribute)
 
@@ -817,15 +815,7 @@ def lookup_attribute(parent, attribute):
 
 
 def apply_patch(parent, attribute, replacement):
-    try:
-        setattr(parent, attribute, replacement)
-    except (TypeError, AttributeError):
-        # It is a built-in/extension type
-        # CAVEAT: Global import raises an error, i.e, asynctest package raises:
-        # 'NoneType' object has no attribute '_spec_coroutines'
-        from forbiddenfruit import curse
-
-        curse(parent, attribute, replacement)
+    setattr(parent, attribute, replacement)
 
 
 def wrap_object(module, name, factory, args=(), kwargs={}):

releasenotes/notes/fix-forbiddenfruit-error-ffbb174ade1c6d57.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Remove ``forbiddenfruit`` as dependency and rollback ``wrapt`` changes where ``forbiddenfruit`` was called.
+    IAST: Patch builtins only when IAST is enabled.
+    

setup.py

@@ -262,7 +262,6 @@ def get_exts_for(name):
         "xmltodict>=0.12",
         "ipaddress; python_version<'3.7'",
         "envier",
-        "forbiddenfruit>=0.1.4",
     ]
     + bytecode,
     extras_require={