diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index 106dd438..dd5adf37 100644 --- a/comfy_cli/tracking.py +++ b/comfy_cli/tracking.py @@ -1,4 +1,5 @@ import functools +import json import logging as logginglib import sys import uuid @@ -16,9 +17,55 @@ MIXPANEL_TOKEN = "93aeab8962b622d431ac19800ccc9f67" mp = Mixpanel(MIXPANEL_TOKEN) if MIXPANEL_TOKEN else None -# Kwargs whose values must never reach tracking system. -# The key is kept (with a redacted marker) so we can still see whether the option was supplied. -SENSITIVE_TRACKING_KEYS = frozenset({"api_key"}) +# --------------------------------------------------------------------------- +# Sensitive-kwarg redaction and kwarg filtering for track_command() +# +# Historical context (BE-992): the original SENSITIVE_TRACKING_KEYS only +# matched "api_key" by exact name, missing kwargs such as +# set_civitai_api_token and set_hf_api_token. Separately, the kwarg filter +# dropped "ctx" and "context" but not "_ctx" (underscore-prefixed Click +# Context). These two gaps *cancelled each other out*: the un-redacted +# credential payload always included the non-JSON-serializable _ctx object, +# so mp.track() raised a TypeError before the event reached Mixpanel. +# +# Fixing the underscore-prefix filter without also broadening the credential +# allowlist would silently activate a CWE-200 credential leak. Both filters +# are therefore fixed atomically here. +# --------------------------------------------------------------------------- + +_SENSITIVE_SUFFIXES = ("_token", "_api_key", "_secret", "_password") +_SENSITIVE_EXACT = frozenset({"api_key", "token", "password", "secret"}) + + +def _is_sensitive(name: str) -> bool: + """Return True if *name* looks like it holds a credential or secret. + + Matched case-insensitively so future kwargs added with non-PEP-8 casing + (``API_KEY``, ``apiKey``, ``AuthToken``, etc.) don't bypass redaction. + """ + lower = name.lower() + return lower in _SENSITIVE_EXACT or lower.endswith(_SENSITIVE_SUFFIXES) + + +def _is_trackable(name: str, value: object) -> bool: + """Return True if the kwarg (name, value) pair is safe to include in a tracking event. + + Excluded: + * ``ctx`` / ``context`` — Click context objects passed by Typer. + * Any name starting with ``_`` — covers ``_ctx`` and other private/internal params. + * Values that are not JSON-serializable — defensive guard so a single + unserializable kwarg doesn't silently swallow the entire event. + """ + if name in ("ctx", "context"): + return False + if name.startswith("_"): + return False + try: + json.dumps(value) + except (TypeError, ValueError, OverflowError, RecursionError): + return False + return True + # Generate a unique tracing ID per command. config_manager = ConfigManager() @@ -79,12 +126,10 @@ def decorator(func): def wrapper(*args, **kwargs): command_name = f"{sub_command}:{func.__name__}" if sub_command is not None else func.__name__ - # Copy kwargs to avoid mutating original dictionary - # Remove context and ctx from the dictionary as they are not needed for tracking and not serializable. filtered_kwargs = { - k: ("" if v is not None else None) if k in SENSITIVE_TRACKING_KEYS else v + k: ("" if v is not None else None) if _is_sensitive(k) else v for k, v in kwargs.items() - if k != "ctx" and k != "context" + if _is_trackable(k, v) } logging.debug(f"Tracking command: {command_name} with arguments: {filtered_kwargs}") diff --git a/tests/comfy_cli/test_tracking.py b/tests/comfy_cli/test_tracking.py index a2e619eb..c634d9c8 100644 --- a/tests/comfy_cli/test_tracking.py +++ b/tests/comfy_cli/test_tracking.py @@ -84,9 +84,6 @@ def some_cmd(workflow, api_key=None): assert "sk-supersecret" not in str(props) def test_api_key_none_stays_none(self, tracking_module): - # When the user didn't pass --api-key (or set $COMFY_API_KEY), we still - # want to be able to see in the analytics that it was absent — not a - # "" sentinel that would imply they did pass one. tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") @tracking_module.track_command() @@ -98,6 +95,84 @@ def some_cmd(workflow, api_key=None): _, kwargs = tracking_module.mp.track.call_args assert kwargs["properties"]["api_key"] is None + def test_set_civitai_api_token_is_redacted(self, tracking_module): + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command("model") + def download(url, set_civitai_api_token=None, set_hf_api_token=None): + return None + + download(url="https://example.com", set_civitai_api_token="civ-real-token") + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert props["set_civitai_api_token"] == "" + assert "civ-real-token" not in str(props) + + def test_set_hf_api_token_is_redacted(self, tracking_module): + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command("model") + def download(url, set_civitai_api_token=None, set_hf_api_token=None): + return None + + download(url="https://example.com", set_hf_api_token="hf_real-token") + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert props["set_hf_api_token"] == "" + assert "hf_real-token" not in str(props) + + def test_bare_token_kwarg_is_redacted(self, tracking_module): + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command() + def some_cmd(workflow, token=None): + return None + + some_cmd(workflow="wf.json", token="my-secret-token") + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert props["token"] == "" + assert "my-secret-token" not in str(props) + + def test_underscore_ctx_is_excluded(self, tracking_module): + import click + + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command("model") + def download(_ctx, url, set_civitai_api_token=None): + return None + + ctx = click.Context(click.Command("download")) + download(_ctx=ctx, url="https://example.com") + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert "_ctx" not in props + assert props["url"] == "https://example.com" + + def test_non_serializable_value_is_excluded(self, tracking_module): + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command() + def some_cmd(workflow, callback=None): + return None + + some_cmd(workflow="wf.json", callback=lambda x: x) + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert "callback" not in props + assert props["workflow"] == "wf.json" + class TestInitTrackingRoundTrip: """End-to-end: init_tracking() writes the string "False"/"True", and track_event honors it.