diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index b1e3d51bbc1..6f077650a92 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -6,15 +6,17 @@ to detect and block/mask sensitive content. """ import asyncio +import bisect import contextvars import itertools import json import os import re import time -from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Iterable, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime +from functools import partial from re import Pattern from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast @@ -79,7 +81,7 @@ from .competitor_intent import ( AirlineCompetitorIntentChecker, BaseCompetitorIntentChecker, ) -from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern +from .patterns import PATTERN_EXTRA_CONFIG, PATTERN_STREAMING_HOLDBACK_CHARS, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) @@ -223,10 +225,70 @@ async def _stream_then_end_of_stream_chunk( @dataclass(frozen=True, slots=True) class _StreamedScanPlan: context_chars: int + holdback_chars: int exception_phrases: tuple[str, ...] conditional_words: tuple[str, ...] +@dataclass(frozen=True, slots=True) +class _RedactionSpan: + start: int + end: int + pattern_index: int + pattern_name: str + + +def _merge_redaction_spans(spans: Iterable[_RedactionSpan]) -> tuple[_RedactionSpan, ...]: + ordered: Final = sorted(spans, key=lambda span: (span.start, -span.end)) + reach: Final = tuple(itertools.accumulate((span.end for span in ordered), max)) + first_indices: Final = tuple( + index for index, span in enumerate(ordered) if index == 0 or span.start >= reach[index - 1] + ) + last_indices: Final = (*(index - 1 for index in first_indices[1:]), len(ordered) - 1) + return tuple(replace(ordered[first], end=reach[last]) for first, last in zip(first_indices, last_indices)) + + +@dataclass(frozen=True, slots=True) +class _RedactedText: + spans: tuple[_RedactionSpan, ...] + tags: tuple[str, ...] + tag_starts: tuple[int, ...] + + def render(self, text: str) -> str: + previous_ends: Final = (0, *(span.end for span in self.spans)) + kept_parts: Final = ( + text[previous_end : span.start] + tag + for previous_end, span, tag in zip(previous_ends, self.spans, self.tags) + ) + return "".join((*kept_parts, text[previous_ends[-1] :])) + + def original_start(self, position: int) -> int: + index: Final = bisect.bisect_right(self.tag_starts, position) - 1 + if index < 0: + return position + tag_end: Final = self.tag_starts[index] + len(self.tags[index]) + if position < tag_end: + return self.spans[index].start + return position - tag_end + self.spans[index].end + + def original_end(self, position: int) -> int: + index: Final = bisect.bisect_left(self.tag_starts, position) - 1 + if index < 0: + return position + tag_end: Final = self.tag_starts[index] + len(self.tags[index]) + if position <= tag_end: + return self.spans[index].end + return position - tag_end + self.spans[index].end + + +def _redacted_text(spans: tuple[_RedactionSpan, ...], redaction_format: str) -> _RedactedText: + tags: Final = tuple(redaction_format.format(pattern_name=span.pattern_name.upper()) for span in spans) + shifts: Final = itertools.accumulate( + (len(tag) - (span.end - span.start) for span, tag in zip(spans, tags)), initial=0 + ) + return _RedactedText(spans, tags, tuple(span.start + shift for span, shift in zip(spans, shifts))) + + class CategoryFileData(TypedDict, total=False): category_name: str description: str @@ -987,21 +1049,6 @@ class ContentFilterGuardrail(CustomGuardrail): merged.append((start, end)) return merged - def _mask_spans(self, text: str, spans: list[tuple[int, int]], redaction: str) -> str: - """Apply masking for the provided spans using the given redaction tag.""" - - if not spans: - return text - - result_parts: Final[list[str]] = [] - previous_end = 0 - for start, end in spans: - result_parts.append(text[previous_end:start]) - result_parts.append(redaction) - previous_end = end - result_parts.append(text[previous_end:]) - return "".join(result_parts) - def _convert_word_number_sequence(self, sequence: str) -> str | None: """Convert a spelled-out digit sequence (e.g., 'One-Two') into digits.""" @@ -1392,10 +1439,8 @@ class ContentFilterGuardrail(CustomGuardrail): self, pattern_name: str, action: ContentFilterAction, - text: str, - spans: list[tuple[int, int]], detections: list[ContentFilterDetection] | None, - ) -> str: + ) -> None: """Handle regex pattern match detection and action.""" if detections is not None: pattern_detection: Final[PatternDetection] = { @@ -1413,11 +1458,49 @@ class ContentFilterGuardrail(CustomGuardrail): detail={"error": error_msg, "pattern": pattern_name}, ) elif action == ContentFilterAction.MASK: - redaction_tag: Final = self.pattern_redaction_format.format(pattern_name=pattern_name.upper()) - text = self._mask_spans(text, spans, redaction_tag) verbose_proxy_logger.info("Masked all %s matches in content", pattern_name) - return text + def _mask_step( + self, + text: str, + spans: tuple[_RedactionSpan, ...], + indexed_entry: tuple[int, CompiledPatternEntry], + ) -> tuple[_RedactionSpan, ...]: + index, pattern_entry = indexed_entry + if pattern_entry["action"] != ContentFilterAction.MASK: + return spans + original_matches: Final = self._find_pattern_spans(text, pattern_entry) + redacted_matches: Final = self._redacted_text_matches(text, spans, pattern_entry) if spans else () + return _merge_redaction_spans( + ( + *spans, + *( + _RedactionSpan(start, end, index, pattern_entry["pattern_name"]) + for start, end in (*original_matches, *redacted_matches) + ), + ) + ) + + def _redacted_text_matches( + self, text: str, spans: tuple[_RedactionSpan, ...], pattern_entry: CompiledPatternEntry + ) -> tuple[tuple[int, int], ...]: + redacted: Final = _redacted_text(spans, self.pattern_redaction_format) + return tuple( + (redacted.original_start(start), redacted.original_end(end)) + for start, end in self._find_pattern_spans(redacted.render(text), pattern_entry) + ) + + def _pattern_detected( + self, + text: str, + index: int, + pattern_entry: CompiledPatternEntry, + span_states: Sequence[tuple[_RedactionSpan, ...]], + ) -> bool: + if pattern_entry["action"] == ContentFilterAction.MASK: + return any(span.pattern_index == index for span in span_states[-1]) + text_after_earlier_masks: Final = _redacted_text(span_states[index], self.pattern_redaction_format).render(text) + return bool(self._find_pattern_spans(text_after_earlier_masks, pattern_entry)) def _handle_blocked_word_match( self, @@ -1506,13 +1589,14 @@ class ContentFilterGuardrail(CustomGuardrail): keyword, category_name, severity, action = category_keyword_match text = self._handle_category_keyword_match(keyword, category_name, severity, action, text, detections) - # Check regex patterns - process ALL patterns, not just first match - for pattern_entry in self.compiled_patterns: - spans = self._find_pattern_spans(text, pattern_entry) - if spans: - pattern_name = pattern_entry["pattern_name"] - action = pattern_entry["action"] - text = self._handle_pattern_match(pattern_name, action, text, spans, detections) + no_spans: Final[tuple[_RedactionSpan, ...]] = () + span_states: Final = tuple( + itertools.accumulate(enumerate(self.compiled_patterns), partial(self._mask_step, text), initial=no_spans) + ) + for index, pattern_entry in enumerate(self.compiled_patterns): + if self._pattern_detected(text, index, pattern_entry, span_states): + self._handle_pattern_match(pattern_entry["pattern_name"], pattern_entry["action"], detections) + text = _redacted_text(span_states[-1], self.pattern_redaction_format).render(text) # Check blocked words - iterate through ALL blocked words text_lower = text.lower() @@ -2071,7 +2155,9 @@ class ContentFilterGuardrail(CustomGuardrail): def _streamed_scan_plan(self) -> _StreamedScanPlan: """ Per-stream inputs for buffer trimming: the retained tail length (the default - context, widened to the longest configured keyword), the category exception + context, widened to the longest configured keyword), the unreleased tail length + (widened to the longest holdback a configured prebuilt pattern declares, since + provider keys run longer than the default), the category exception phrases, which suppress matches anywhere in the scanned text, and the conditional category words, which only match when paired inside one sentence. """ @@ -2081,6 +2167,13 @@ class ContentFilterGuardrail(CustomGuardrail): ) return _StreamedScanPlan( context_chars=max(CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS, longest_keyword), + holdback_chars=max( + ( + PATTERN_STREAMING_HOLDBACK_CHARS.get(entry["pattern_name"], CONTENT_FILTER_STREAMING_HOLDBACK_CHARS) + for entry in self.compiled_patterns + ), + default=CONTENT_FILTER_STREAMING_HOLDBACK_CHARS, + ), exception_phrases=tuple( phrase for category in self.loaded_categories.values() for phrase in category.exceptions ), @@ -2221,7 +2314,7 @@ class ContentFilterGuardrail(CustomGuardrail): safe_to_yield_len = max( previous_state.yielded_masked_text_len, - len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + len(masked_text) - (0 if is_final else plan.holdback_chars), ) released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] if isinstance(choice, TextChoices): diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json index db6dd36726c..dcb2dc5ce77 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json @@ -98,6 +98,28 @@ "category": "Credential Patterns", "description": "Detects generic API key patterns" }, + { + "name": "openai_api_key", + "display_name": "OpenAI API Key", + "pattern": "sk-(?:(?:proj|svcacct|admin|None)-[A-Za-z0-9_\\-]{20,200}T3BlbkFJ(?:(?!sk-(?:proj|svcacct|admin|None)-)[A-Za-z0-9_\\-]){20,200}|[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}(?![A-Za-z0-9_]))", + "streaming_holdback_chars": 419, + "category": "Credential Patterns", + "description": "Detects OpenAI API keys (sk-proj-..., sk-svcacct-..., sk-admin-..., sk-None-... and legacy sk-... keys, all carrying the T3BlbkFJ marker)" + }, + { + "name": "anthropic_api_key", + "display_name": "Anthropic API Key", + "pattern": "(?:(? str: + return "".join(await cls._collect_streamed_deltas(guardrail, chunks, metadata)) + + @staticmethod + async def _collect_streamed_deltas( + guardrail: ContentFilterGuardrail, + chunks: list[str], + metadata: dict[str, list[StandardLoggingGuardrailInformation]], + ) -> list[str]: from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices async def mock_stream(): @@ -951,7 +1042,7 @@ class TestContentFilterGuardrail: request_data={"messages": [], "model": "gpt-4o", "metadata": metadata}, ): yielded.append(chunk.choices[0].delta.content or "") - return "".join(yielded) + return yielded @pytest.mark.asyncio async def test_streaming_hook_scans_bounded_window_per_chunk(self): @@ -1068,6 +1159,53 @@ class TestContentFilterGuardrail: assert entry["guardrail_status"] == "guardrail_intervened" assert [d["keyword"] for d in entry["guardrail_response"]] == [phrase] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("pattern_name", "key"), + [ + ("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74), + ("openai_api_key", "sk-svcacct-" + "a" * 200 + "T3BlbkFJ" + "b" * 200), + ("anthropic_api_key", "sk-ant-api03-" + "a" * 93 + "AA"), + ("gemini_api_key", "AQ." + "a" * 55), + ], + ) + async def test_streaming_hook_masks_whole_provider_key_split_across_chunks(self, pattern_name: str, key: str): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-streaming-provider-key", + patterns=[ + ContentFilterPattern( + pattern_type="prebuilt", + pattern_name=pattern_name, + action=ContentFilterAction.MASK, + ) + ], + event_hook=GuardrailEventHooks.post_call, + ) + text: Final = f"Use {key} to call the API." + chunks: Final = [text[i : i + 4] for i in range(0, len(text), 4)] + + streamed: Final = await self._collect_streamed_text(guardrail, chunks, {}) + + assert streamed == f"Use [{pattern_name.upper()}_REDACTED] to call the API." + + @pytest.mark.asyncio + async def test_streaming_hook_keeps_default_holdback_without_provider_key_patterns(self): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-streaming-older-credential-patterns", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK) + for name in ("aws_access_key", "aws_secret_key", "github_token", "slack_token", "generic_api_key") + ], + event_hook=GuardrailEventHooks.post_call, + ) + text: Final = "plain words with nothing secret in them. " * 12 + chunks: Final = [text[i : i + 10] for i in range(0, len(text), 10)] + + deltas: Final = await self._collect_streamed_deltas(guardrail, chunks, {}) + + assert "".join(deltas) == text + assert len(next(delta for delta in reversed(deltas) if delta)) <= CONTENT_FILTER_STREAMING_HOLDBACK_CHARS + @pytest.mark.asyncio async def test_streaming_hook_blocks_keyword_longer_than_scan_context(self): """ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py index d702b9e0116..3c71e814fe1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py @@ -4,6 +4,8 @@ Tests for content filter pattern loading from JSON import json import os +import time +from typing import Final import pytest @@ -183,3 +185,169 @@ def test_eu_patterns_have_category(): assert ( pattern_name in eu_category_patterns ), f"Pattern {pattern_name} not in EU PII Patterns category" + + +def test_provider_key_patterns_are_offered_in_the_ui_as_credential_patterns() -> None: + ui_categories: Final = {pattern["name"]: pattern["category"] for pattern in get_pattern_metadata()} + assert [ui_categories.get(name) for name in ("openai_api_key", "anthropic_api_key", "gemini_api_key")] == [ + "Credential Patterns" + ] * 3 + + +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="openai-project"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-project-shortest"), + pytest.param("openai_api_key", "sk-svcacct-" + "a_b-" * 15 + "T3BlbkFJ" + "c-d_" * 15, id="openai-service-account"), + pytest.param("openai_api_key", "sk-admin-" + "a" * 58 + "T3BlbkFJ" + "b" * 58, id="openai-admin"), + pytest.param("openai_api_key", "sk-None-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-user-key"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19 + "-", id="openai-trailing-hyphen"), + pytest.param( + "anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOP", id="anthropic-versioned" + ), + pytest.param( + "anthropic_api_key", "sk-ant-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU", id="anthropic-unversioned" + ), + pytest.param("anthropic_api_key", "sk-ant-" + "a" * 20, id="anthropic-shortest"), + pytest.param("anthropic_api_key", "sk-ant-api03-" + "a_b-" * 23 + "AA", id="anthropic-urlsafe"), + pytest.param( + "anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyz0123456789-", id="anthropic-trailing-hyphen" + ), + pytest.param("gemini_api_key", "AIzaSyD-9tSrke72I6_AQp2h8HI4-uE6QNKQV0", id="gemini-legacy"), + pytest.param("gemini_api_key", "AIza" + "a" * 33, id="gemini-legacy-shortest"), + pytest.param("gemini_api_key", "AIza" + "a" * 36, id="gemini-legacy-longest"), + pytest.param("gemini_api_key", "AIza" + "a" * 34 + "-", id="gemini-legacy-trailing-hyphen"), + pytest.param("gemini_api_key", "AQ.Ab8RN6I-7xK9pL2mQ0vW5yZ1uX3tS4rE8qD7cB6aN5mP4oI3uY2", id="gemini-auth-key"), + pytest.param("gemini_api_key", "AQ." + "a" * 45, id="gemini-auth-key-shortest"), + pytest.param("gemini_api_key", "AQ." + "a" * 55, id="gemini-auth-key-longest"), + pytest.param("gemini_api_key", "AQ." + "a" * 49 + "-", id="gemini-auth-key-trailing-hyphen"), + ], +) +def test_credential_pattern_matches_whole_key(pattern_name: str, key: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f"I keep getting 401 with {key}, is it expired?") + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize("separator", ["\\n", "\\t", "\\u003c", "Bearer%20", "?key%3D", "?key%253D", "GOOGLE_KEY_"]) +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai"), + pytest.param("anthropic_api_key", "sk-ant-api03-" + "a" * 20, id="anthropic"), + pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"), + pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"), + ], +) +def test_credential_pattern_matches_key_after_escaped_separator(pattern_name: str, key: str, separator: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f'{{"log": "retrying with{separator}{key} failed"}}') + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize( + ("pattern_name", "text"), + [ + pytest.param("openai_api_key", "sk-short", id="openai-too-short"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 19, id="openai-project-too-short"), + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 21, id="openai-legacy-too-long"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 48, id="openai-project-without-marker"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 19 + "T3BlbkFJ" + "b" * 20, id="openai-project-marker-too-early"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19, id="openai-project-marker-too-late"), + pytest.param("openai_api_key", "SK-ADMIN-dashboard-redesign-phase-two", id="openai-prefixed-slug"), + pytest.param("openai_api_key", "desk-proj-management-dashboard-v2", id="openai-inside-longer-word"), + pytest.param( + "openai_api_key", "sk-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRST", id="openai-legacy-without-marker" + ), + pytest.param("openai_api_key", "sk-0123456789abcdef0123456789abcdef", id="openai-other-vendor-sk-key"), + pytest.param("openai_api_key", "sk-Xq3vR9tLm2Pa8KdW5yZgQw", id="openai-litellm-virtual-key"), + pytest.param( + "openai_api_key", + "sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789AbCd-_x", + id="openai-litellm-virtual-key-raised-length", + ), + pytest.param("openai_api_key", "not-a-key", id="openai-no-prefix"), + pytest.param( + "openai_api_key", "sk-ant-api03-thisisanthropicnotopenaiabc123", id="openai-ignores-anthropic-key" + ), + pytest.param("openai_api_key", "ordinary sentence without any api key in it", id="openai-plain-text"), + pytest.param("anthropic_api_key", "sk-ant-" + "a" * 19, id="anthropic-too-short"), + pytest.param("anthropic_api_key", "flask-ant-design-system-components", id="anthropic-inside-longer-word"), + pytest.param( + "anthropic_api_key", + "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEF1234567890", + id="anthropic-ignores-openai-project-key", + ), + pytest.param( + "anthropic_api_key", + "sk-admin-abcdefghijklmnopqrstuvwxyzABCDEF123456", + id="anthropic-ignores-openai-admin-key", + ), + pytest.param("anthropic_api_key", "sk-ant", id="anthropic-prefix-only"), + pytest.param("anthropic_api_key", "plain text with nothing sensitive", id="anthropic-plain-text"), + pytest.param("gemini_api_key", "AIza", id="gemini-prefix-only"), + pytest.param("gemini_api_key", "AIza" + "a" * 32, id="gemini-legacy-too-short"), + pytest.param("gemini_api_key", "AIza" + "a" * 37, id="gemini-legacy-too-long"), + pytest.param("gemini_api_key", "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB", id="gemini-wrong-prefix"), + pytest.param("gemini_api_key", "AQ.", id="gemini-auth-prefix-only"), + pytest.param("gemini_api_key", "AQ." + "a" * 44, id="gemini-auth-key-too-short"), + pytest.param("gemini_api_key", "AQ." + "a" * 56, id="gemini-auth-key-too-long"), + pytest.param("gemini_api_key", "AQ-" + "a" * 50, id="gemini-auth-key-wrong-separator"), + pytest.param("gemini_api_key", "x" + "AIza" + "a" * 35, id="gemini-inside-longer-token"), + pytest.param("gemini_api_key", "The AQ. was reported at 42 ppm today.", id="gemini-aq-acronym"), + pytest.param("gemini_api_key", "Is the question AQ. Maybe it is BQ.", id="gemini-aq-sentence-end"), + pytest.param("gemini_api_key", "plain text with no google key", id="gemini-plain-text"), + pytest.param( + "gemini_api_key", + "https://example.com/blog/aizawa-takeshi-interview-2024-tokyo-01", + id="gemini-lowercase-legacy-prefix", + ), + pytest.param( + "gemini_api_key", + "ssh aq.internal-observability-metrics-collector-eu-west-01", + id="gemini-lowercase-auth-prefix", + ), + ], +) +def test_credential_pattern_ignores_non_keys(pattern_name: str, text: str) -> None: + assert get_compiled_pattern(pattern_name).search(text) is None + + +@pytest.mark.parametrize("suffix", ["请检查", "é", "ü der Schlüssel"]) +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"), + pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"), + pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"), + ], +) +def test_credential_pattern_matches_key_followed_by_non_ascii_letters(pattern_name: str, key: str, suffix: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f"我的密钥是{key}{suffix}") + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize( + "leading", + [ + pytest.param("sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="another-key"), + pytest.param("sk-proj-" + "x" * 20 + "T3BlbkFJ" + "y" * 192, id="key-shaped-filler"), + pytest.param("sk-proj-" + "x" * 100 + "T3BlbkFJ" + "y" * 100, id="long-key-shaped-filler"), + ], +) +def test_openai_key_pattern_masks_a_key_glued_to_preceding_key_shaped_text(leading: str) -> None: + key: Final = "sk-proj-" + "c" * 74 + "T3BlbkFJ" + "d" * 74 + assert get_compiled_pattern("openai_api_key").sub("[KEY]", leading + key) == "[KEY][KEY]" + + +@pytest.mark.parametrize("prefix", ["sk-proj-", "sk-svcacct-", "sk-admin-", "sk-None-"]) +def test_openai_key_pattern_scans_repeated_prefixes_without_stalling(prefix: str) -> None: + prompt: Final = prefix * (128_000 // len(prefix)) + started: Final = time.perf_counter() + match: Final = get_compiled_pattern("openai_api_key").search(prompt) + elapsed: Final = time.perf_counter() - started + assert match is None + assert elapsed < 1.0