|
4 | 4 | from typing import Any |
5 | 5 | from typing import Dict |
6 | 6 | from typing import Optional |
| 7 | +from typing import Union |
7 | 8 |
|
8 | 9 | import xmltodict |
9 | 10 |
|
10 | 11 | from ddtrace._trace.span import Span |
11 | 12 | from ddtrace.appsec._asm_request_context import _call_waf |
12 | 13 | from ddtrace.appsec._asm_request_context import _call_waf_first |
13 | 14 | from ddtrace.appsec._asm_request_context import get_blocked |
| 15 | +from ddtrace.appsec._asm_request_context import set_body_response |
14 | 16 | from ddtrace.appsec._constants import SPAN_DATA_NAMES |
15 | 17 | from ddtrace.appsec._http_utils import extract_cookies_from_headers |
16 | 18 | from ddtrace.appsec._http_utils import normalize_headers |
@@ -157,6 +159,14 @@ def _on_lambda_start_response( |
157 | 159 | _call_waf(("aws_lambda",)) |
158 | 160 |
|
159 | 161 |
|
| 162 | +def _on_lambda_parse_body( |
| 163 | + response_body: Optional[Union[str, Dict[str, Any]]], |
| 164 | +): |
| 165 | + if asm_config._api_security_feature_active: |
| 166 | + if response_body: |
| 167 | + set_body_response(response_body) |
| 168 | + |
| 169 | + |
160 | 170 | # ASGI |
161 | 171 |
|
162 | 172 |
|
@@ -408,6 +418,7 @@ def listen(): |
408 | 418 |
|
409 | 419 | core.on("aws_lambda.start_request", _on_lambda_start_request) |
410 | 420 | core.on("aws_lambda.start_response", _on_lambda_start_response) |
| 421 | + core.on("aws_lambda.parse_body", _on_lambda_parse_body) |
411 | 422 |
|
412 | 423 | core.on("grpc.server.response.message", _on_grpc_server_response) |
413 | 424 | core.on("grpc.server.data", _on_grpc_server_data) |
|
0 commit comments