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