From a1d3061d37975d57abe2926aa1fe13bcc693d0f7 Mon Sep 17 00:00:00 2001 From: JY Tan Date: Fri, 6 Feb 2026 15:53:26 -0800 Subject: [PATCH] Commit --- drift/instrumentation/django/html_utils.py | 11 ++++ tests/unit/test_html_utils.py | 62 ++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/drift/instrumentation/django/html_utils.py b/drift/instrumentation/django/html_utils.py index 6d51346..47b045b 100644 --- a/drift/instrumentation/django/html_utils.py +++ b/drift/instrumentation/django/html_utils.py @@ -62,6 +62,17 @@ def normalize_csrf_in_body(body: bytes | None) -> bytes | None: flags=re.IGNORECASE, ) + # Pattern 3: JavaScript CSRF token assignment (e.g., DRF Swagger UI) + # request.headers["X-CSRFTOKEN"] = "ABC123..."; + # Also handles single quotes and X-CSRFToken variants + csrf_js_pattern = r'(headers\[["\']X-CSRFTOKEN["\']\]\s*=\s*["\'])[^"\']+(["\'])' + body_str = re.sub( + csrf_js_pattern, + rf"\g<1>{CSRF_PLACEHOLDER}\2", + body_str, + flags=re.IGNORECASE, + ) + return body_str.encode("utf-8") except Exception as e: diff --git a/tests/unit/test_html_utils.py b/tests/unit/test_html_utils.py index c6cea6a..3b97167 100644 --- a/tests/unit/test_html_utils.py +++ b/tests/unit/test_html_utils.py @@ -172,6 +172,68 @@ def test_handles_non_utf8_gracefully(self): result = normalize_csrf_in_body(html) assert result == html + def test_normalizes_js_csrf_header_assignment_double_quotes(self): + """JavaScript CSRF header assignment with double quotes should be normalized.""" + html = b"""""" + result = normalize_csrf_in_body(html) + assert result is not None + assert CSRF_PLACEHOLDER.encode() in result + assert b"7180AeP5yScQ72FIE7z8tCtrrQc32Sgaqpqgk8ktVXeBXw7a5qDgXoGgLiWF22Ew" not in result + + def test_normalizes_js_csrf_header_assignment_single_quotes(self): + """JavaScript CSRF header assignment with single quotes should be normalized.""" + html = b"""""" + result = normalize_csrf_in_body(html) + assert result is not None + assert CSRF_PLACEHOLDER.encode() in result + assert b"someToken123" not in result + + def test_normalizes_js_csrf_header_case_insensitive(self): + """JavaScript CSRF header matching should be case-insensitive (X-CSRFToken).""" + html = b"""""" + result = normalize_csrf_in_body(html) + assert result is not None + assert CSRF_PLACEHOLDER.encode() in result + assert b"myToken456" not in result + + def test_normalizes_js_csrf_header_preserves_surrounding_js(self): + """Surrounding JavaScript should be preserved when normalizing CSRF header.""" + html = b"""""" + result = normalize_csrf_in_body(html) + assert result is not None + assert b"requestInterceptor" in result + assert b"return request" in result + assert CSRF_PLACEHOLDER.encode() in result + assert b"tokenABC" not in result + + def test_normalizes_both_form_and_js_csrf_tokens(self): + """Both form hidden inputs and JS CSRF header assignments should be normalized.""" + html = b""" +
+ + """ + result = normalize_csrf_in_body(html) + assert result is not None + assert result.count(CSRF_PLACEHOLDER.encode()) == 2 + assert b"formToken123" not in result + assert b"jsToken456" not in result + class TestNormalizeHtmlBody: """Tests for normalize_html_body function."""