diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index c648f7faea0..778a6e8e18b 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -149,10 +149,11 @@ class BaseEmailLogger(CustomLogger): ) if include_api_key is None: include_api_key = True # Default to True if not set - key_token_display = ( + key_token_display: Final = html.escape( send_key_created_email_event.virtual_key if include_api_key - else "[Key hidden for security - retrieve from dashboard]" + else "[Key hidden for security - retrieve from dashboard]", + quote=True, ) email_html_content = KEY_CREATED_EMAIL_TEMPLATE.format( diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94a3564e5cb..430f2fe03f3 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -2,6 +2,7 @@ # Class for sending Slack Alerts # import asyncio import datetime +import html import os import random import time @@ -1317,11 +1318,13 @@ Model Info: ) if webhook_event.event == "key_created": + # Custom keys can contain markup. Escape at render time; the event keeps the raw key. + displayed_key: Final[str] = html.escape(key_token, quote=True) if isinstance(key_token, str) else "" email_html_content = KEY_CREATED_EMAIL_TEMPLATE.format( email_logo_url=email_logo_url, recipient_email=recipient_email, key_budget=key_budget, - key_token=key_token, + key_token=displayed_key, base_url=base_url, email_support_contact=email_support_contact, ) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 519002f935a..ed87093fc4a 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -2,7 +2,7 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter @@ -82,7 +82,7 @@ class KeyManagementEventHooks: # Send email notification - non-blocking, independent operation if data.send_invite_email is True: try: - await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) + await KeyManagementEventHooks._send_key_created_email(response) except Exception as e: verbose_proxy_logger.warning("Failed to send key created email: %s", e) @@ -517,7 +517,7 @@ class KeyManagementEventHooks: return False @staticmethod - async def _send_key_created_email(response: dict): + async def _send_key_created_email(response: GenerateKeyResponse): """ Send key created email if email sending is enabled. @@ -532,7 +532,7 @@ class KeyManagementEventHooks: from litellm.proxy.proxy_server import general_settings, proxy_logging_obj # generate_key_fn remaps token to the hashed token_id. The plaintext sk- key stays on key. - plaintext_key: Final[str] = cast(str, response.get("key") or response.get("token") or "") + plaintext_key: Final[str] = response.key or response.token or "" ########################## # v2 integration for emails (enterprise) @@ -555,11 +555,11 @@ class KeyManagementEventHooks: event_group=Litellm_EntityType.KEY, event_message="API Key Created", token=plaintext_key, - spend=response.get("spend", 0.0), - max_budget=response.get("max_budget", 0.0), - user_id=response.get("user_id", None), - team_id=response.get("team_id", "Default Team"), - key_alias=response.get("key_alias", None), + spend=0.0 if response.spend is None else response.spend, + max_budget=0.0 if response.max_budget is None else response.max_budget, + user_id=response.user_id, + team_id=response.team_id if response.team_id is not None else "Default Team", + key_alias=response.key_alias, ) for email_logger in initialized_email_loggers: if isinstance(email_logger, BaseEmailLogger): @@ -581,11 +581,11 @@ class KeyManagementEventHooks: event_group=Litellm_EntityType.KEY, event_message="API Key Created", token=plaintext_key, - spend=response.get("spend", 0.0), - max_budget=response.get("max_budget", 0.0), - user_id=response.get("user_id", None), - team_id=response.get("team_id", "Default Team"), - key_alias=response.get("key_alias", None), + spend=0.0 if response.spend is None else response.spend, + max_budget=0.0 if response.max_budget is None else response.max_budget, + user_id=response.user_id, + team_id=response.team_id if response.team_id is not None else "Default Team", + key_alias=response.key_alias, ) # If user configured email alerting - send an Email letting their end-user know the key was created asyncio.create_task( diff --git a/tests/unit/proxy/hooks/test_send_invite_email.py b/tests/unit/proxy/hooks/test_send_invite_email.py index b1edee39d24..fec7c63352d 100644 --- a/tests/unit/proxy/hooks/test_send_invite_email.py +++ b/tests/unit/proxy/hooks/test_send_invite_email.py @@ -1,4 +1,5 @@ import asyncio +import html import sys from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -18,20 +19,26 @@ from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventH _PLAINTEXT_KEY = "sk-plaintext-key-39315" _TOKEN_ID = "a1b2c3d4e5f6789012345678abcdef01" +_MARKUP_KEY = 'sk-1234567890123456Activate key
'
-def _key_created_email_response() -> dict:
- """Response after generate_key_fn remaps token onto the hashed token_id."""
- return {
- "key": _PLAINTEXT_KEY,
- "token": _TOKEN_ID,
- "token_id": _TOKEN_ID,
- "user_id": "user-123",
- "team_id": "team-1",
- "key_alias": "ci-key",
- "spend": 0.0,
- "max_budget": 10.0,
- }
+def _key_created_email_response(key: str = _PLAINTEXT_KEY) -> GenerateKeyResponse:
+ """Response after generate_key_fn remaps token onto the hashed token_id.
+
+ Constructing the model with both fields copies token onto key, so the hash is
+ assigned afterwards, matching generate_key_fn.
+ """
+ response = GenerateKeyResponse(
+ key=key,
+ token_id=_TOKEN_ID,
+ user_id="user-123",
+ team_id="team-1",
+ key_alias="ci-key",
+ spend=0.0,
+ max_budget=10.0,
+ )
+ response.token = _TOKEN_ID
+ return response
def _proxy_server_module(slack_alerting: object) -> SimpleNamespace:
@@ -309,9 +316,7 @@ async def test_v1_key_generation_no_email_when_send_invite_email_false():
mock_send_key_created_email.assert_not_called()
-@pytest.mark.asyncio
-async def test_v0_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
- """The v0 invite email prints webhook token, which must be the sk- key."""
+async def _render_v0_key_created_email(monkeypatch: pytest.MonkeyPatch, key: str) -> str:
for name in ("EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", "SMTP_SENDER_LOGO"):
monkeypatch.delenv(name, raising=False)
@@ -338,18 +343,34 @@ async def test_v0_key_created_email_contains_plaintext_key_not_token_id(monkeypa
),
patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(slack_alerting)}),
):
- await KeyManagementEventHooks._send_key_created_email(_key_created_email_response())
+ await KeyManagementEventHooks._send_key_created_email(_key_created_email_response(key))
assert scheduled, "key created email was not scheduled"
await asyncio.gather(*scheduled)
assert sent_html, "key created email was not sent"
- assert f"Key: {_PLAINTEXT_KEY}" in sent_html[0]
- assert _TOKEN_ID not in sent_html[0]
+ return sent_html[0]
@pytest.mark.asyncio
-async def test_enterprise_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
- """The enterprise invite email prints virtual_key, which must be the sk- key."""
+async def test_v0_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
+ """The v0 invite email prints webhook token, which must be the sk- key."""
+ sent_html = await _render_v0_key_created_email(monkeypatch, _PLAINTEXT_KEY)
+ assert f"Key: {_PLAINTEXT_KEY}" in sent_html
+ assert _TOKEN_ID not in sent_html
+
+
+@pytest.mark.asyncio
+async def test_v0_key_created_email_escapes_markup_in_custom_key(monkeypatch):
+ """A custom key is text in the legacy HTML email, not markup."""
+ sent_html = await _render_v0_key_created_email(monkeypatch, _MARKUP_KEY)
+ escaped_key = html.escape(_MARKUP_KEY, quote=True)
+ assert f"Key: {escaped_key}" in sent_html
+ assert f'api_key="{escaped_key}"' in sent_html
+ assert '' not in sent_html
+ assert _TOKEN_ID not in sent_html
+
+
+async def _render_enterprise_key_created_email(monkeypatch: pytest.MonkeyPatch, key: str):
pytest.importorskip("litellm_enterprise")
monkeypatch.setenv("EMAIL_INCLUDE_API_KEY", "true")
for name in ("EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", "SMTP_SENDER_LOGO", "EMAIL_SIGNATURE"):
@@ -382,12 +403,34 @@ async def test_enterprise_key_created_email_contains_plaintext_key_not_token_id(
),
patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(None)}),
):
- await KeyManagementEventHooks._send_key_created_email(_key_created_email_response())
+ await KeyManagementEventHooks._send_key_created_email(_key_created_email_response(key))
+ return recording_logger
+
+
+@pytest.mark.asyncio
+async def test_enterprise_key_created_email_contains_plaintext_key_not_token_id(monkeypatch):
+ """The enterprise invite email prints virtual_key, which must be the sk- key."""
+ recording_logger = await _render_enterprise_key_created_email(monkeypatch, _PLAINTEXT_KEY)
assert recording_logger.events, "enterprise key created email was not sent"
assert recording_logger.events[0].virtual_key == _PLAINTEXT_KEY
assert recording_logger.events[0].token == _PLAINTEXT_KEY
assert recording_logger.html_bodies, "enterprise key created email was not sent"
- html = recording_logger.html_bodies[0]
- assert f'{_PLAINTEXT_KEY}' in html
- assert _TOKEN_ID not in html
+ sent_html = recording_logger.html_bodies[0]
+ assert f'{_PLAINTEXT_KEY}' in sent_html
+ assert _TOKEN_ID not in sent_html
+
+
+@pytest.mark.asyncio
+async def test_enterprise_key_created_email_escapes_markup_in_custom_key(monkeypatch):
+ """The stored event keeps the raw key, and the enterprise HTML escapes it."""
+ recording_logger = await _render_enterprise_key_created_email(monkeypatch, _MARKUP_KEY)
+ escaped_key = html.escape(_MARKUP_KEY, quote=True)
+
+ assert recording_logger.events[0].virtual_key == _MARKUP_KEY
+ assert recording_logger.events[0].token == _MARKUP_KEY
+ sent_html = recording_logger.html_bodies[0]
+ assert f'{escaped_key}' in sent_html
+ assert f'"{escaped_key}"' in sent_html
+ assert '' not in sent_html
+ assert _TOKEN_ID not in sent_html