|
5 | 5 |
|
6 | 6 | try: |
7 | 7 | from ddtrace.appsec.iast import oce |
| 8 | + from ddtrace.appsec.iast._patch_modules import patch_iast |
8 | 9 | from ddtrace.appsec.iast._taint_tracking import OriginType |
9 | 10 | from ddtrace.appsec.iast._taint_tracking import create_context |
10 | 11 | from ddtrace.appsec.iast._taint_tracking import is_pyobject_tainted |
11 | 12 | from ddtrace.appsec.iast._taint_tracking import setup as taint_tracking_setup |
12 | 13 | from ddtrace.appsec.iast._taint_tracking import taint_pyobject |
13 | 14 | from ddtrace.appsec.iast._taint_utils import LazyTaintDict |
| 15 | + from ddtrace.appsec.iast._taint_utils import LazyTaintList |
14 | 16 | from ddtrace.appsec.iast._taint_utils import check_tainted_args |
15 | 17 | except (ImportError, AttributeError): |
16 | 18 | pytest.skip("IAST not supported for this Python version", allow_module_level=True) |
17 | 19 |
|
18 | 20 |
|
19 | 21 | def setup(): |
| 22 | + patch_iast() |
20 | 23 | create_context() |
21 | 24 | taint_tracking_setup(bytes.join, bytearray.join) |
22 | 25 | oce._enabled = True |
@@ -237,3 +240,31 @@ def test_checked_tainted_args(): |
237 | 240 | assert check_tainted_args( |
238 | 241 | args=(tainted_arg, untainted_arg), kwargs=None, tracer=None, integration_name="psycopg", method=cursor.execute |
239 | 242 | ) |
| 243 | + |
| 244 | + |
| 245 | +def test_json_encode_dict(): |
| 246 | + import json |
| 247 | + |
| 248 | + tainted_dict = LazyTaintDict( |
| 249 | + { |
| 250 | + "tr_key_001": ["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}], |
| 251 | + "tr_key_002": {"tr_key_003": {"tr_key_004": "tr_val_005"}}, |
| 252 | + }, |
| 253 | + origins=(OriginType.PARAMETER, OriginType.PARAMETER), |
| 254 | + ) |
| 255 | + |
| 256 | + assert json.dumps(tainted_dict) == ( |
| 257 | + '{"tr_key_001": ["tr_val_001", "tr_val_002", "tr_val_003", ' |
| 258 | + '{"tr_key_005": "tr_val_004"}], "tr_key_002": {"tr_key_003": {"tr_key_004": "tr_val_005"}}}' |
| 259 | + ) |
| 260 | + |
| 261 | + |
| 262 | +def test_json_encode_list(): |
| 263 | + import json |
| 264 | + |
| 265 | + tainted_list = LazyTaintList( |
| 266 | + ["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}], |
| 267 | + origins=(OriginType.PARAMETER, OriginType.PARAMETER), |
| 268 | + ) |
| 269 | + |
| 270 | + assert json.dumps(tainted_list) == '["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}]' |
0 commit comments