""" Telemetry module for ComfyUI-Doctor. Provides opt-in anonymous usage data collection with local-only storage (Phase 1-3). All data is stored locally; no network upload in Phase 1-3. Features: - TelemetryStore: Thread-safe event buffer with atomic file writes - RateLimiter: Token bucket algorithm (60 events/minute) - Event validation: Allowlist-based category/action/label validation - PII detection: Blacklist scan for sensitive patterns """ import json import os import re import threading import time import tempfile import uuid from dataclasses import dataclass, field, asdict from datetime import timedelta from typing import Optional, List, Dict, Set, Tuple, Any try: try: from .services.doctor_paths import get_doctor_data_dir as _get_canonical_doctor_data_dir except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from services.doctor_paths import get_doctor_data_dir as _get_canonical_doctor_data_dir except Exception: _get_canonical_doctor_data_dir = None try: from .services.time_utils import parse_utc_timestamp, utc_filename_timestamp, utc_isoformat, utc_now except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from services.time_utils import parse_utc_timestamp, utc_filename_timestamp, utc_isoformat, utc_now try: from .terminal_output import emit_doctor_log except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from terminal_output import emit_doctor_log # ═══════════════════════════════════════════════════════════════════════════ # CONSTANTS # ═══════════════════════════════════════════════════════════════════════════ SCHEMA_VERSION = "1.0" MAX_EVENTS = 500 MAX_LABEL_LENGTH = 64 MAX_VALUE = 10000 EVENT_TTL_DAYS = 7 # Allowlist for category → actions ALLOWED_EVENTS: Dict[str, List[str]] = { "feature": ["tab_switch"], "analysis": ["pattern_matched", "llm_called"], "resolution": ["marked"], "session": ["start", "end"], } # Allowlist for labels per category/action ALLOWED_LABELS: Dict[str, Dict[str, List[str]]] = { "feature": { "tab_switch": ["chat", "stats", "settings"], }, "analysis": { "pattern_matched": [], # Dynamic: loaded from PatternLoader "llm_called": [ "openai", "deepseek", "anthropic", "ollama", "lmstudio", "gemini", "groq", "openrouter", "xai", "custom" ], }, "resolution": { "marked": ["resolved", "unresolved", "ignored"], }, "session": { "start": [], "end": [], }, } # PII detection patterns (blacklist) PII_PATTERNS = [ re.compile(r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}'), # Email re.compile(r'[A-Za-z]:\\Users\\[^\\]+'), # Windows user path re.compile(r'/home/[^/]+'), # Linux home re.compile(r'/Users/[^/]+'), # macOS home re.compile(r'sk-[a-zA-Z0-9]{20,}'), # API key pattern ] # ═══════════════════════════════════════════════════════════════════════════ # DATA CLASSES # ═══════════════════════════════════════════════════════════════════════════ @dataclass class TelemetryEvent: """Represents a single telemetry event.""" schema_version: str event_id: str timestamp: str category: str action: str label: Optional[str] = None value: Optional[int] = None def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for JSON serialization.""" result = { "schema_version": self.schema_version, "event_id": self.event_id, "timestamp": self.timestamp, "category": self.category, "action": self.action, } if self.label is not None: result["label"] = self.label if self.value is not None: result["value"] = self.value return result @classmethod def from_dict(cls, data: Dict[str, Any]) -> "TelemetryEvent": """Create from dictionary.""" return cls( schema_version=data.get("schema_version", SCHEMA_VERSION), event_id=data.get("event_id", ""), timestamp=data.get("timestamp", ""), category=data.get("category", ""), action=data.get("action", ""), label=data.get("label"), value=data.get("value"), ) # ═══════════════════════════════════════════════════════════════════════════ # RATE LIMITER # ═══════════════════════════════════════════════════════════════════════════ # R7: Import shared RateLimiter from rate_limiter.py (avoid code duplication) try: from .rate_limiter import RateLimiter except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from rate_limiter import RateLimiter # ═══════════════════════════════════════════════════════════════════════════ # PATTERN ALLOWLIST # ═══════════════════════════════════════════════════════════════════════════ _pattern_cache: Optional[Set[str]] = None _pattern_cache_time: float = 0 PATTERN_CACHE_TTL = 300 # 5 minutes def get_pattern_allowlist() -> Set[str]: """Get the current set of valid pattern IDs from PatternLoader.""" global _pattern_cache, _pattern_cache_time now = time.time() if _pattern_cache is not None and (now - _pattern_cache_time) < PATTERN_CACHE_TTL: return _pattern_cache try: try: from .pattern_loader import get_pattern_loader except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from pattern_loader import get_pattern_loader # IMPORTANT: read ids from the analyzer's shared loader (same pattern directories). The old # `PatternLoader()` call lacked its required directories; the swallowed TypeError left the # allowlist empty and mapped every label to __unknown__. _pattern_cache = set(get_pattern_loader().get_pattern_ids()) _pattern_cache_time = now return _pattern_cache except Exception: # If PatternLoader fails, return empty set (all patterns → __unknown__) return set() def validate_pattern_label(label: str) -> str: """Validate pattern ID, return __unknown__ if not in allowlist.""" if not label: return "__unknown__" allowlist = get_pattern_allowlist() if label in allowlist: return label return "__unknown__" # ═══════════════════════════════════════════════════════════════════════════ # VALIDATION # ═══════════════════════════════════════════════════════════════════════════ def contains_pii(text: str) -> bool: """Check if text contains PII patterns.""" if not text: return False for pattern in PII_PATTERNS: if pattern.search(text): return True return False def validate_event(data: Dict[str, Any]) -> Tuple[bool, str, Optional[TelemetryEvent]]: """ Validate event data against schema and allowlists. Returns: (is_valid, error_message, event_or_none) """ # 1. Required fields category = data.get("category") action = data.get("action") if not category or not action: return False, "Missing required field: category or action", None # 2. Category allowlist if category not in ALLOWED_EVENTS: return False, f"Invalid category: {category}", None # 3. Action allowlist if action not in ALLOWED_EVENTS[category]: return False, f"Invalid action: {action} for category {category}", None # 4. Label validation label = data.get("label") if label is not None: # Check length if len(label) > MAX_LABEL_LENGTH: return False, f"Label too long: {len(label)} > {MAX_LABEL_LENGTH}", None # Check PII if contains_pii(label): return False, "Label contains PII pattern", None # Check allowlist allowed_labels = ALLOWED_LABELS.get(category, {}).get(action, []) if allowed_labels: # Static allowlist exists if label not in allowed_labels: return False, f"Invalid label: {label}", None elif category == "analysis" and action == "pattern_matched": # Dynamic pattern validation label = validate_pattern_label(label) # 5. Value validation value = data.get("value") if value is not None: if not isinstance(value, int): return False, "Value must be integer", None if value < 0 or value > MAX_VALUE: return False, f"Value out of range: {value}", None # Create event event = TelemetryEvent( schema_version=SCHEMA_VERSION, event_id=str(uuid.uuid4()), timestamp=utc_isoformat(), category=category, action=action, label=label, value=value, ) return True, "", event # ═══════════════════════════════════════════════════════════════════════════ # TELEMETRY STORE # ═══════════════════════════════════════════════════════════════════════════ def get_doctor_data_dir() -> str: """ Get the Doctor data directory (cross-platform). R18: Prefer the shared canonical resolver (Desktop/portable safe). Fallback: OS temp dir (never write under Desktop resources). """ if _get_canonical_doctor_data_dir: try: return _get_canonical_doctor_data_dir() except Exception: pass data_dir = os.path.join(tempfile.gettempdir(), "ComfyUI-Doctor") try: os.makedirs(data_dir, exist_ok=True) except Exception: pass return data_dir def _legacy_telemetry_path() -> str: extension_root = os.path.dirname(os.path.abspath(__file__)) return os.path.join(extension_root, "data", "telemetry.json") def _migrate_legacy_telemetry_file(target_path: str) -> None: """ One-time best-effort migration from legacy `/data/telemetry.json` to the canonical data dir path. """ legacy_path = _legacy_telemetry_path() if legacy_path == target_path: return if not os.path.exists(legacy_path) or os.path.exists(target_path): return try: os.makedirs(os.path.dirname(target_path), exist_ok=True) with open(legacy_path, "rb") as src: payload = src.read() with open(target_path, "wb") as dst: dst.write(payload) ts = utc_filename_timestamp() os.rename(legacy_path, f"{legacy_path}.migrated-{ts}") except Exception: # Migration failures must not break startup or telemetry. pass def get_telemetry_path() -> str: """Get telemetry file path.""" target_path = os.path.join(get_doctor_data_dir(), "telemetry.json") _migrate_legacy_telemetry_file(target_path) return target_path class TelemetryStore: """ Thread-safe telemetry event storage with atomic file writes. Features: - Max 500 events buffer - 7-day TTL auto-purge - Atomic write via temp file + rename - Single-writer lock """ def __init__(self, filepath: Optional[str] = None, enabled: bool = False): """ Initialize telemetry store. Args: filepath: Path to telemetry JSON file (default: auto-detect) enabled: Whether telemetry is enabled """ self._filepath = filepath or get_telemetry_path() self._enabled = enabled self._lock = threading.Lock() self._buffer: List[TelemetryEvent] = [] self._loaded = False self._rate_limiter = RateLimiter(max_per_minute=60) @property def enabled(self) -> bool: """Check if telemetry is enabled.""" return self._enabled @enabled.setter def enabled(self, value: bool) -> None: """Set telemetry enabled state.""" self._enabled = value @property def filepath(self) -> str: """Get the telemetry file path.""" return self._filepath def _load(self) -> None: """Load events from file (internal, called with lock held).""" if self._loaded: return try: if os.path.exists(self._filepath): with open(self._filepath, "r", encoding="utf-8") as f: data = json.load(f) if isinstance(data, list): self._buffer = [ TelemetryEvent.from_dict(e) for e in data if isinstance(e, dict) ] except (json.JSONDecodeError, OSError) as e: emit_doctor_log(f"Telemetry load warning: {e}", "WARNING") self._buffer = [] self._loaded = True def _save_atomic(self) -> None: """Save events atomically via temp file + rename.""" try: # Ensure directory exists dir_path = os.path.dirname(self._filepath) if dir_path and not os.path.exists(dir_path): os.makedirs(dir_path, exist_ok=True) # Write to temp file temp_path = self._filepath + ".tmp" with open(temp_path, "w", encoding="utf-8") as f: json.dump( [e.to_dict() for e in self._buffer], f, ensure_ascii=False, indent=2, ) # Atomic rename os.replace(temp_path, self._filepath) except OSError as e: emit_doctor_log(f"Telemetry save warning: {e}", "WARNING") def _purge_old_events(self) -> None: """Remove events older than TTL.""" cutoff = utc_now() - timedelta(days=EVENT_TTL_DAYS) self._buffer = [ event for event in self._buffer if (parse_utc_timestamp(event.timestamp) or cutoff) >= cutoff ] def _enforce_limit(self) -> None: """Enforce max events limit (remove oldest).""" if len(self._buffer) > MAX_EVENTS: self._buffer = self._buffer[-MAX_EVENTS:] def track(self, data: Dict[str, Any]) -> Tuple[bool, str]: """ Record a telemetry event. Args: data: Event data (category, action, label?, value?) Returns: (success, message) """ # Check if enabled if not self._enabled: return False, "Telemetry disabled" # Rate limit if not self._rate_limiter.allow(): return False, "Rate limited" # Validate is_valid, error_msg, event = validate_event(data) if not is_valid or event is None: return False, error_msg # Store with self._lock: self._load() self._buffer.append(event) self._purge_old_events() self._enforce_limit() self._save_atomic() return True, "Event recorded" def get_buffer(self) -> List[Dict[str, Any]]: """Get all buffered events as dictionaries.""" with self._lock: self._load() return [e.to_dict() for e in self._buffer] def get_stats(self) -> Dict[str, Any]: """Get buffer statistics.""" with self._lock: self._load() if not self._buffer: return { "count": 0, "oldest": None, "newest": None, } return { "count": len(self._buffer), "oldest": self._buffer[0].timestamp if self._buffer else None, "newest": self._buffer[-1].timestamp if self._buffer else None, } def clear(self) -> None: """Clear all buffered events.""" with self._lock: self._buffer = [] self._save_atomic() def export_json(self) -> str: """Export buffer as JSON string.""" with self._lock: self._load() return json.dumps( [e.to_dict() for e in self._buffer], ensure_ascii=False, indent=2, ) def __len__(self) -> int: """Return buffer size.""" with self._lock: self._load() return len(self._buffer) # ═══════════════════════════════════════════════════════════════════════════ # SINGLETON INSTANCE # ═══════════════════════════════════════════════════════════════════════════ # Global telemetry store instance (disabled by default) _telemetry_store: Optional[TelemetryStore] = None def get_telemetry_store() -> TelemetryStore: """Get the global telemetry store instance.""" global _telemetry_store if _telemetry_store is None: _telemetry_store = TelemetryStore(enabled=False) return _telemetry_store def track_event(category: str, action: str, label: Optional[str] = None, value: Optional[int] = None) -> bool: """ Convenience function to track a telemetry event. Args: category: Event category (feature, analysis, resolution, session) action: Event action label: Optional label value: Optional numeric value Returns: True if event was recorded, False otherwise """ store = get_telemetry_store() success, _ = store.track({ "category": category, "action": action, "label": label, "value": value, }) return success