wip

Author: smola

ddtrace/contrib/internal/langchain/patch.py

@@ -1199,7 +1199,6 @@ def unpatch():
 
 
 def iast_propagate_prompt_template_format(func, instance, args, kwargs):
-    print(f"iast_propagate_prompt_template_format({kwargs})")
     result = func(*args, **kwargs)
     try:
         if not asm_config.is_iast_request_enabled:
@@ -1215,12 +1214,11 @@ def iast_propagate_prompt_template_format(func, instance, args, kwargs):
                 source = ranges[0].source
                 break
         if source:
-            print("TAINTED")
             return taint_pyobject(result, source.name, source.value, source.origin)
 
     except Exception as e:
+        from ddtrace.appsec._iast._metrics import _set_iast_error_metric
         _set_iast_error_metric("IAST propagation error. langchain iast_propagate_prompt_template_format. {}".format(e))
-        print(f"IAST propagation error. langchain iast_propagate_prompt_template_format. {e}")
     return result
 
 

tests/appsec/integrations/langchain_tests/test_iast_langchain.py

@@ -1,13 +1,15 @@
+from langchain.agents import AgentType
+from langchain.agents import initialize_agent
+from langchain_community.tools.shell.tool import ShellTool
+from langchain_core.language_models.fake import FakeListLLM
+
 from ddtrace.appsec._iast.constants import VULN_CMDI
 from tests.appsec.iast.conftest import iast_span_defaults  # noqa: F401
 from tests.appsec.iast.iast_utils import _iast_patched_module
 from tests.appsec.iast.iast_utils import get_line_and_hash
 from tests.appsec.iast.taint_sinks.conftest import _get_span_report
 from tests.utils import override_env
-from langchain.agents import AgentType
-from langchain.agents import initialize_agent
-from langchain_community.tools.shell.tool import ShellTool
-from langchain_core.language_models.fake import FakeListLLM
+
 
 with override_env({"DD_IAST_ENABLED": "True"}):
     from ddtrace.appsec._iast._taint_tracking import OriginType