pyJWT + urllib instrumentations

.cursor/BUGBOT.md

@@ -0,0 +1,5 @@
+# BUGBOT Notes
+
+## Instrumentation Guidelines
+
+- When adding a new instrumentation, the README must be updated to document the new instrumentation.

README.md

@@ -56,6 +56,8 @@ Tusk Drift currently supports the following packages and versions:
 | psycopg2 | all versions |
 | Redis | `>=4.0.0` |
 | Kinde | `>=2.0.1` |
+| PyJWT | all versions |
+| urllib.request | all versions |
 
 If you're using packages or versions not listed above, please create an issue with the package + version you'd like an instrumentation for.
 

drift/core/content_type_utils.py

@@ -17,7 +17,7 @@
     "application/vnd.api+json": DecodedType.JSON,
     # Plain Text (ALLOWED)
     "text/plain": DecodedType.PLAIN_TEXT,
-    # HTML (BLOCKED)
+    # HTML
     "text/html": DecodedType.HTML,
     "application/xhtml+xml": DecodedType.HTML,
     # CSS (BLOCKED)
@@ -111,9 +111,7 @@
     "application/binary": DecodedType.BINARY,
 }
 
-# Only JSON and plain text are acceptable (matches Node SDK)
-# All other content types will cause trace blocking
-ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT}
+ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT, DecodedType.HTML}
 
 
 def get_decoded_type(content_type: str | None) -> DecodedType | None:

drift/core/drift_sdk.py

@@ -406,6 +406,16 @@ def _init_auto_instrumentations(self) -> None:
         except ImportError:
             pass
 
+        try:
+            import urllib.request
+
+            from ..instrumentation.urllib import UrllibInstrumentation
+
+            _ = UrllibInstrumentation()
+            logger.debug("urllib instrumentation initialized")
+        except ImportError:
+            pass
+
         # Initialize PostgreSQL instrumentation before Django
         # Instrument BOTH psycopg2 and psycopg if available
         # This allows apps to use either or both
@@ -481,6 +491,15 @@ def _init_auto_instrumentations(self) -> None:
             except Exception as e:
                 logger.debug(f"Socket instrumentation initialization failed: {e}")
 
+            # PyJWT instrumentation for JWT verification bypass
+            try:
+                from ..instrumentation.pyjwt import PyJWTInstrumentation
+
+                _ = PyJWTInstrumentation(mode=self.mode)
+                logger.debug("PyJWT instrumentation registered (REPLAY mode)")
+            except Exception as e:
+                logger.debug(f"PyJWT instrumentation registration failed: {e}")
+
     def create_env_vars_snapshot(self) -> None:
         """Create a span capturing all environment variables.
 

drift/core/mock_utils.py

@@ -183,7 +183,9 @@ def find_mock_response_sync(
         mock_response = sdk.request_mock_sync(mock_request)
 
         if not mock_response or not mock_response.found:
-            logger.debug(f"No matching mock found for {trace_id} with input value: {input_value}")
+            logger.debug(
+                f"No matching mock found for {trace_id} with input value: {input_value}, input schema: {input_schema_merges}, input schema hash: {outbound_span.input_schema_hash}, input value hash: {outbound_span.input_value_hash}"
+            )
             return None
 
         logger.debug(f"Found mock response for {trace_id}")

drift/core/trace_blocking_manager.py

@@ -204,7 +204,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     """Check if a span should be blocked due to size or server error status.
 
     Blocks the trace if:
-    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 300)
+    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 400)
     2. The span exceeds the maximum size limit (1MB)
 
     This matches Node SDK behavior in TdSpanExporter.ts.
@@ -221,7 +221,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     span_name = span.name
     blocking_manager = TraceBlockingManager.get_instance()
 
-    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 300)
+    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 400)
     if span.kind == SpanKind.SERVER and span.status.code == StatusCode.ERROR:
         logger.debug(f"Blocking trace {trace_id} - server span '{span_name}' has error status")
         blocking_manager.block_trace(trace_id, reason="server_error")

drift/instrumentation/django/middleware.py

@@ -390,8 +390,8 @@ def dict_to_schema_merges(merges_dict):
         duration_seconds = duration_ns // 1_000_000_000
         duration_nanos = duration_ns % 1_000_000_000
 
-        # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-        if status_code >= 300:
+        # Match Node SDK: >= 400 is considered an error
+        if status_code >= 400:
             status = SpanStatus(code=StatusCode.ERROR, message=f"HTTP {status_code}")
         else:
             status = SpanStatus(code=StatusCode.OK, message="")

drift/instrumentation/fastapi/instrumentation.py

@@ -530,8 +530,8 @@ def _finalize_span(
     TuskDrift.get_instance()
 
     status_code = response_data.get("status_code", 200)
-    # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-    if status_code >= 300:
+    # Match Node SDK: >= 400 is considered an error
+    if status_code >= 400:
         span_info.span.set_status(Status(OTelStatusCode.ERROR, f"HTTP {status_code}"))
     else:
         span_info.span.set_status(Status(OTelStatusCode.OK))

drift/instrumentation/pyjwt/__init__.py

@@ -0,0 +1,5 @@
+"""PyJWT instrumentation for REPLAY mode."""
+
+from .instrumentation import PyJWTInstrumentation
+
+__all__ = ["PyJWTInstrumentation"]

drift/instrumentation/pyjwt/instrumentation.py

@@ -0,0 +1,98 @@
+"""PyJWT instrumentation for REPLAY mode.
+
+Patches PyJWT to disable all verification during test replay:
+1. _merge_options - returns all verification options as False
+2. _verify_signature - no-op (defense in depth)
+3. _validate_claims - no-op (defense in depth)
+
+Only active in REPLAY mode.
+"""
+
+from __future__ import annotations
+
+import logging
+from types import ModuleType
+
+from ...core.types import TuskDriftMode
+from ..base import InstrumentationBase
+
+logger = logging.getLogger(__name__)
+
+
+class PyJWTInstrumentation(InstrumentationBase):
+    """Patches PyJWT to disable verification in REPLAY mode."""
+
+    def __init__(self, mode: TuskDriftMode = TuskDriftMode.DISABLED, enabled: bool = True) -> None:
+        self.mode = mode
+        should_enable = enabled and mode == TuskDriftMode.REPLAY
+
+        super().__init__(
+            name="PyJWTInstrumentation",
+            module_name="jwt",
+            supported_versions="*",
+            enabled=should_enable,
+        )
+
+    def patch(self, module: ModuleType) -> None:
+        if self.mode != TuskDriftMode.REPLAY:
+            return
+
+        self._patch_merge_options()
+        self._patch_signature_verification()
+        self._patch_claim_validation()
+        logger.debug("[PyJWTInstrumentation] All patches applied")
+
+    def _patch_signature_verification(self) -> None:
+        """No-op signature verification."""
+        try:
+            from jwt import api_jws
+
+            def patched_verify_signature(self, *args, **kwargs):
+                logger.debug("[PyJWTInstrumentation] _verify_signature called - skipping verification")
+                return None
+
+            api_jws.PyJWS._verify_signature = patched_verify_signature
+            logger.debug("[PyJWTInstrumentation] Patched PyJWS._verify_signature")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _verify_signature: {e}")
+
+    def _patch_claim_validation(self) -> None:
+        """No-op claim validation."""
+        try:
+            from jwt import api_jwt
+
+            def patched_validate_claims(self, *args, **kwargs):
+                logger.debug("[PyJWTInstrumentation] _validate_claims called - skipping validation")
+                return None
+
+            api_jwt.PyJWT._validate_claims = patched_validate_claims
+            logger.debug("[PyJWTInstrumentation] Patched PyJWT._validate_claims")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _validate_claims: {e}")
+
+    def _patch_merge_options(self) -> None:
+        """Patch _merge_options to always return disabled verification options."""
+        try:
+            from jwt import api_jwt
+
+            disabled_options = {
+                "verify_signature": False,
+                "verify_exp": False,
+                "verify_nbf": False,
+                "verify_iat": False,
+                "verify_aud": False,
+                "verify_iss": False,
+                "verify_sub": False,
+                "verify_jti": False,
+                "require": [],
+                "strict_aud": False,
+            }
+
+            def patched_merge_options(self, options=None):
+                logger.debug("[PyJWTInstrumentation] _merge_options called - returning disabled options")
+                return disabled_options
+
+            api_jwt.PyJWT._merge_options = patched_merge_options
+            logger.debug("[PyJWTInstrumentation] Patched PyJWT._merge_options")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _merge_options: {e}")

drift/instrumentation/urllib/__init__.py

@@ -0,0 +1,5 @@
+"""urllib.request instrumentation module."""
+
+from .instrumentation import UrllibInstrumentation
+
+__all__ = ["UrllibInstrumentation"]

drift/instrumentation/urllib/e2e-tests/.tusk/config.yaml

@@ -0,0 +1,27 @@
+version: 1
+
+service:
+  id: "urllib-e2e-test-id"
+  name: "urllib-e2e-test"
+  port: 8000
+  start:
+    command: "python src/app.py"
+  readiness_check:
+    command: "curl -f http://localhost:8000/health"
+    timeout: 45s
+    interval: 5s
+
+tusk_api:
+  url: "http://localhost:8000"
+
+test_execution:
+  concurrent_limit: 10
+  batch_size: 10
+  timeout: 30s
+
+recording:
+  sampling_rate: 1.0
+  export_spans: false
+
+replay:
+  enable_telemetry: false

drift/instrumentation/urllib/e2e-tests/Dockerfile

@@ -0,0 +1,21 @@
+FROM python-e2e-base:latest
+
+# Copy SDK source for editable install
+COPY . /sdk
+
+# Copy test files
+COPY drift/instrumentation/urllib/e2e-tests /app
+
+WORKDIR /app
+
+# Install dependencies (requirements.txt uses -e /sdk for SDK)
+RUN pip install -q -r requirements.txt
+
+# Make entrypoint executable
+RUN chmod +x entrypoint.py
+
+# Create .tusk directories
+RUN mkdir -p /app/.tusk/traces /app/.tusk/logs
+
+# Run entrypoint
+ENTRYPOINT ["python", "entrypoint.py"]

drift/instrumentation/urllib/e2e-tests/docker-compose.yml

@@ -0,0 +1,19 @@
+services:
+  app:
+    build:
+      context: ../../../..
+      dockerfile: drift/instrumentation/urllib/e2e-tests/Dockerfile
+      args:
+        - TUSK_CLI_VERSION=${TUSK_CLI_VERSION:-latest}
+    environment:
+      - PORT=8000
+      - TUSK_ANALYTICS_DISABLED=1
+      - PYTHONUNBUFFERED=1
+    working_dir: /app
+    volumes:
+      # Mount SDK source for hot reload (no rebuild needed for SDK changes)
+      - ../../../..:/sdk
+      # Mount app source for development
+      - ./src:/app/src
+      # Mount .tusk folder to persist traces
+      - ./.tusk:/app/.tusk

drift/instrumentation/urllib/e2e-tests/entrypoint.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+"""
+E2E Test Entrypoint for Urllib Instrumentation
+
+This script orchestrates the full e2e test lifecycle:
+1. Setup: Install dependencies
+2. Record: Start app in RECORD mode, execute requests
+3. Test: Run Tusk CLI tests
+4. Teardown: Cleanup and return exit code
+"""
+
+import sys
+from pathlib import Path
+
+# Add SDK to path for imports
+sys.path.insert(0, "/sdk")
+
+from drift.instrumentation.e2e_common.base_runner import E2ETestRunnerBase
+
+
+class UrllibE2ETestRunner(E2ETestRunnerBase):
+    """E2E test runner for Urllib instrumentation."""
+
+    def __init__(self):
+        import os
+
+        port = int(os.getenv("PORT", "8000"))
+        super().__init__(app_port=port)
+
+
+if __name__ == "__main__":
+    runner = UrllibE2ETestRunner()
+    exit_code = runner.run()
+    sys.exit(exit_code)

drift/instrumentation/urllib/e2e-tests/requirements.txt

@@ -0,0 +1,2 @@
+-e /sdk
+Flask>=3.1.2