stemtrace 0.1.0__py3-none-any.whl

stemtrace/library/bootsteps.py

@@ -0,0 +1,169 @@
+"""Celery bootsteps for capturing task lifecycle events.
+
+Bootsteps hook into Celery's worker consumer pipeline to capture
+events that don't have dedicated signals (like RECEIVED).
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from celery import bootsteps
+
+from stemtrace.core.events import TaskEvent, TaskState
+
+if TYPE_CHECKING:
+    from celery.worker.consumer import Consumer
+
+logger = logging.getLogger(__name__)
+
+# Import here to avoid circular imports
+_publish_event: Any = None
+
+
+def _set_publisher(publisher: Any) -> None:
+    """Set the event publisher function (called from signals module)."""
+    global _publish_event
+    _publish_event = publisher
+
+
+# Celery's bootsteps module is not fully typed
+class ReceivedEventStep(bootsteps.ConsumerStep):  # type: ignore[misc]
+    """Bootstep that emits RECEIVED events when tasks are received by worker.
+
+    This hooks into the consumer's task message handling to capture the moment
+    a task message is received from the broker, before it's executed.
+    """
+
+    requires: ClassVar[set[str]] = {"celery.worker.consumer.tasks:Tasks"}
+
+    def __init__(
+        self,
+        consumer: Consumer,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the bootstep."""
+        super().__init__(consumer, **kwargs)
+        self._original_strategy: dict[str, Any] = {}
+
+    def start(self, consumer: Consumer) -> None:
+        """Wrap task strategies to emit RECEIVED events."""
+        logger.debug("ReceivedEventStep: wrapping task strategies")
+
+        # Wrap each task's strategy to emit RECEIVED before execution
+        for name, strategy in list(consumer.strategies.items()):
+            if name not in self._original_strategy:
+                self._original_strategy[name] = strategy
+                consumer.strategies[name] = self._wrap_strategy(name, strategy)
+
+    def stop(self, consumer: Consumer) -> None:
+        """Restore original strategies."""
+        for name, strategy in self._original_strategy.items():
+            if name in consumer.strategies:
+                consumer.strategies[name] = strategy
+        self._original_strategy.clear()
+
+    def _wrap_strategy(self, task_name: str, strategy: Any) -> Any:
+        """Wrap a task strategy to emit RECEIVED before execution."""
+
+        def wrapped_strategy(
+            message: Any,
+            body: Any,
+            ack: Any,
+            reject: Any,
+            callbacks: Any,
+            **kwargs: Any,
+        ) -> Any:
+            """Emit RECEIVED event before delegating to original strategy."""
+            # Emit RECEIVED event
+            try:
+                self._emit_received(task_name, body, message)
+            except Exception:
+                logger.debug("Failed to emit RECEIVED event", exc_info=True)
+
+            # Call original strategy
+            return strategy(message, body, ack, reject, callbacks, **kwargs)
+
+        return wrapped_strategy
+
+    def _emit_received(self, task_name: str, body: Any, message: Any) -> None:
+        """Emit a RECEIVED event for the task.
+
+        Skips retries since they already have RETRY + STARTED events.
+        """
+        if _publish_event is None:
+            return
+
+        # Extract task ID and retries from body (can be tuple or dict format)
+        task_id: str | None = None
+        parent_id: str | None = None
+        root_id: str | None = None
+        group_id: str | None = None
+        chord_id: str | None = None
+        retries: int = 0
+
+        if isinstance(body, list | tuple) and len(body) >= 3:
+            # Old format: (args, kwargs, embed)
+            embed = body[2] if len(body) > 2 else {}
+            if isinstance(embed, dict):
+                task_id = embed.get("id")
+                parent_id = embed.get("parent_id")
+                root_id = embed.get("root_id")
+                group_id = embed.get("group")
+                chord_id = embed.get("chord")
+                retries = embed.get("retries", 0) or 0
+        elif isinstance(body, dict):
+            # New format: {'task': ..., 'id': ..., ...}
+            task_id = body.get("id")
+            parent_id = body.get("parent_id")
+            root_id = body.get("root_id")
+            group_id = body.get("group")
+            chord_id = body.get("chord")
+            retries = body.get("retries", 0) or 0
+
+        # Fallback to message headers/properties
+        if hasattr(message, "headers"):
+            headers = message.headers or {}
+            if task_id is None:
+                task_id = headers.get("id")
+            parent_id = parent_id or headers.get("parent_id")
+            root_id = root_id or headers.get("root_id")
+            group_id = group_id or headers.get("group")
+            chord_id = chord_id or headers.get("chord")
+            if retries == 0:
+                retries = headers.get("retries", 0) or 0
+
+        if task_id is None:
+            logger.debug("Could not extract task_id for RECEIVED event")
+            return
+
+        # Skip RECEIVED for retries - they already have RETRY events
+        if retries > 0:
+            logger.debug("Skipping RECEIVED for retry %d of task %s", retries, task_id)
+            return
+
+        _publish_event(
+            TaskEvent(
+                task_id=task_id,
+                name=task_name,
+                state=TaskState.RECEIVED,
+                timestamp=datetime.now(timezone.utc),
+                parent_id=parent_id,
+                root_id=root_id,
+                group_id=group_id,
+                chord_id=chord_id,
+            )
+        )
+        logger.debug("Emitted RECEIVED for task %s", task_id)
+
+
+def register_bootsteps(app: Any) -> None:
+    """Register stemtrace bootsteps with the Celery app.
+
+    Args:
+        app: The Celery application instance.
+    """
+    app.steps["consumer"].add(ReceivedEventStep)
+    logger.debug("Registered ReceivedEventStep bootstep")

stemtrace/library/config.py

@@ -0,0 +1,55 @@
+"""Configuration for stemtrace library."""
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class StemtraceConfig(BaseModel):
+    """Frozen configuration for stemtrace initialization.
+
+    Attributes:
+        transport_url: Broker URL for event transport.
+        prefix: Redis key prefix for stemtrace data.
+        ttl: Time-to-live for stored events in seconds.
+        capture_args: Whether to capture task args/kwargs.
+        capture_result: Whether to capture task return values.
+        max_data_size: Maximum size in bytes for serialized data.
+        scrub_sensitive_data: Whether to scrub sensitive keys.
+        additional_sensitive_keys: Extra keys to treat as sensitive.
+        safe_keys: Keys to never scrub (overrides sensitive).
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    transport_url: str
+    prefix: str = "stemtrace"
+    ttl: int = 86400
+
+    # Data capture options (all enabled by default)
+    capture_args: bool = True
+    capture_result: bool = True
+    max_data_size: int = 10240  # 10KB
+
+    # Scrubbing options (enabled by default)
+    scrub_sensitive_data: bool = True
+    additional_sensitive_keys: frozenset[str] = Field(default_factory=frozenset)
+    safe_keys: frozenset[str] = Field(default_factory=frozenset)
+
+
+_config: StemtraceConfig | None = None
+
+
+def get_config() -> StemtraceConfig | None:
+    """Get the active configuration, or None if not initialized."""
+    return _config
+
+
+def set_config(config: StemtraceConfig) -> None:
+    """Set the active configuration."""
+    global _config
+    _config = config
+
+
+def _reset_config() -> None:
+    """Reset configuration. For testing only."""
+    global _config
+    _config = None

stemtrace/library/scrubbing.py

@@ -0,0 +1,195 @@
+"""Sensitive data scrubbing for task arguments and results.
+
+Implements Sentry-style data scrubbing with a configurable denylist of
+sensitive key patterns. Values matching sensitive keys are replaced with
+a filtered placeholder.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+FILTERED = "[Filtered]"
+
+# Default sensitive key patterns (case-insensitive, partial match)
+DEFAULT_SENSITIVE_KEYS: frozenset[str] = frozenset(
+    {
+        # Authentication
+        "password",
+        "passwd",
+        "pwd",
+        "pass",
+        "secret",
+        "api_key",
+        "apikey",
+        "api_secret",
+        "token",
+        "access_token",
+        "refresh_token",
+        "auth_token",
+        "bearer",
+        "authorization",
+        "auth",
+        "credentials",
+        "private_key",
+        "private",
+        # Financial/PII
+        "credit_card",
+        "card_number",
+        "cvv",
+        "ssn",
+        "social_security",
+        # Session
+        "cookie",
+        "session",
+        "sessionid",
+        "csrf",
+        "csrf_token",
+    }
+)
+
+
+def _is_sensitive_key(
+    key: str,
+    sensitive_keys: frozenset[str],
+    safe_keys: frozenset[str] | None = None,
+) -> bool:
+    """Check if a key matches any sensitive pattern.
+
+    Uses case-insensitive partial matching:
+    - 'user_password' matches 'password'
+    - 'API_KEY' matches 'api_key'
+
+    Args:
+        key: The key name to check.
+        sensitive_keys: Set of sensitive key patterns.
+        safe_keys: Optional set of keys to never scrub.
+
+    Returns:
+        True if the key should be scrubbed.
+    """
+    if safe_keys and key.lower() in {k.lower() for k in safe_keys}:
+        return False
+
+    key_lower = key.lower()
+    return any(pattern in key_lower for pattern in sensitive_keys)
+
+
+def scrub_dict(
+    data: dict[str, Any],
+    sensitive_keys: frozenset[str] = DEFAULT_SENSITIVE_KEYS,
+    additional_keys: frozenset[str] | None = None,
+    safe_keys: frozenset[str] | None = None,
+) -> dict[str, Any]:
+    """Recursively scrub sensitive keys from a dictionary.
+
+    Args:
+        data: Dictionary to scrub.
+        sensitive_keys: Base set of sensitive key patterns.
+        additional_keys: Extra keys to treat as sensitive.
+        safe_keys: Keys to never scrub (overrides sensitive).
+
+    Returns:
+        New dictionary with sensitive values replaced by FILTERED.
+
+    Example:
+        >>> scrub_dict({"password": "secret123", "name": "Alice"})
+        {'password': '[Filtered]', 'name': 'Alice'}
+    """
+    all_sensitive = (
+        sensitive_keys | additional_keys if additional_keys else sensitive_keys
+    )
+
+    result: dict[str, Any] = {}
+    for key, value in data.items():
+        if _is_sensitive_key(key, all_sensitive, safe_keys):
+            result[key] = FILTERED
+        else:
+            result[key] = _scrub_value(value, all_sensitive, safe_keys)
+
+    return result
+
+
+def _scrub_value(
+    value: Any,
+    sensitive_keys: frozenset[str],
+    safe_keys: frozenset[str] | None = None,
+) -> Any:
+    """Recursively scrub a value (handles nested dicts and lists)."""
+    if isinstance(value, dict):
+        return scrub_dict(value, sensitive_keys, safe_keys=safe_keys)
+    elif isinstance(value, list):
+        return [_scrub_value(item, sensitive_keys, safe_keys) for item in value]
+    elif isinstance(value, tuple):
+        return tuple(_scrub_value(item, sensitive_keys, safe_keys) for item in value)
+    else:
+        return value
+
+
+def scrub_args(
+    args: tuple[Any, ...],
+    sensitive_keys: frozenset[str] = DEFAULT_SENSITIVE_KEYS,
+    additional_keys: frozenset[str] | None = None,
+    safe_keys: frozenset[str] | None = None,
+) -> list[Any]:
+    """Scrub positional arguments (recursively handles nested structures).
+
+    Args:
+        args: Tuple of positional arguments.
+        sensitive_keys: Base set of sensitive key patterns.
+        additional_keys: Extra keys to treat as sensitive.
+        safe_keys: Keys to never scrub.
+
+    Returns:
+        List of scrubbed arguments.
+    """
+    all_sensitive = (
+        sensitive_keys | additional_keys if additional_keys else sensitive_keys
+    )
+    return [_scrub_value(arg, all_sensitive, safe_keys) for arg in args]
+
+
+def safe_serialize(
+    value: Any,
+    max_size: int = 10240,
+    sensitive_keys: frozenset[str] = DEFAULT_SENSITIVE_KEYS,
+    additional_keys: frozenset[str] | None = None,
+    safe_keys: frozenset[str] | None = None,
+) -> Any:
+    """Serialize a value to JSON-compatible format with size limit and scrubbing.
+
+    Args:
+        value: Value to serialize.
+        max_size: Maximum size in bytes before truncation.
+        sensitive_keys: Base set of sensitive key patterns.
+        additional_keys: Extra keys to treat as sensitive.
+        safe_keys: Keys to never scrub.
+
+    Returns:
+        JSON-serializable value, or string representation if serialization fails.
+    """
+    all_sensitive = (
+        sensitive_keys | additional_keys if additional_keys else sensitive_keys
+    )
+
+    # First scrub the value
+    scrubbed = _scrub_value(value, all_sensitive, safe_keys)
+
+    # Try to serialize to check size and ensure it's JSON-compatible
+    try:
+        serialized = json.dumps(scrubbed, default=str)
+        if len(serialized) > max_size:
+            return f"[Truncated: {len(serialized)} bytes > {max_size} max]"
+        # Return the original scrubbed value (not the JSON string)
+        return scrubbed
+    except (TypeError, ValueError) as e:
+        logger.debug("Failed to serialize value: %s", e)
+        # Fall back to string representation
+        str_repr = str(scrubbed)
+        if len(str_repr) > max_size:
+            return f"[Truncated: {len(str_repr)} bytes > {max_size} max]"
+        return str_repr