Merge 6d246da917 into e16d79d635

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(

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,
                 )

litellm/proxy/hooks/key_management_event_hooks.py

@@ -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)
 
@@ -523,7 +523,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.
 
@@ -537,6 +537,9 @@ 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] = response.key or response.token or ""
+
         ##########################
         # v2 integration for emails (enterprise)
         ##########################
@@ -553,16 +556,16 @@ class KeyManagementEventHooks:
             )
             if len(initialized_email_loggers) > 0:
                 event = SendKeyCreatedEmailEvent(
-                    virtual_key=response.get("key", ""),
+                    virtual_key=plaintext_key,
                     event="key_created",
                     event_group=Litellm_EntityType.KEY,
                     event_message="API Key Created",
-                    token=response.get("token", ""),
-                    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),
+                    token=plaintext_key,
+                    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):
@@ -583,12 +586,12 @@ class KeyManagementEventHooks:
                 event="key_created",
                 event_group=Litellm_EntityType.KEY,
                 event_message="API Key Created",
-                token=response.get("token", ""),
-                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),
+                token=plaintext_key,
+                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(

tests/unit/proxy/hooks/test_send_invite_email.py

@@ -1,16 +1,55 @@
-import pytest
-from unittest.mock import AsyncMock, patch, MagicMock
-from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
-from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
-from litellm.proxy._types import (
-    NewUserRequest,
-    NewUserResponse,
-    GenerateKeyRequest,
-    GenerateKeyResponse,
-    UserAPIKeyAuth,
-)
+import asyncio
+import html
 import sys
 from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+from litellm.proxy._types import (
+    GenerateKeyRequest,
+    GenerateKeyResponse,
+    NewUserRequest,
+    NewUserResponse,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
+from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
+
+_PLAINTEXT_KEY = "sk-plaintext-key-39315"
+_TOKEN_ID = "a1b2c3d4e5f6789012345678abcdef01"
+_MARKUP_KEY = 'sk-9876543210987654</pre><a href="https://attacker.example/login">Activate key</a><pre>'
+
+
+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:
+    return SimpleNamespace(
+        general_settings={"alerting": ["email"]},
+        proxy_logging_obj=SimpleNamespace(slack_alerting_instance=slack_alerting),
+        premium_user=False,
+        prisma_client=None,
+        CommonProxyErrors=SimpleNamespace(not_premium_user=SimpleNamespace(value="not premium")),
+        litellm_proxy_admin_name="admin-user",
+    )
 
 
 @pytest.mark.asyncio
@@ -23,9 +62,7 @@ async def test_v1_user_creation_no_email_when_send_invite_email_false():
     mock_proxy_logging_obj = MagicMock()
     mock_proxy_logging_obj.slack_alerting_instance = mock_slack_alerting
 
-    with patch(
-        "litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]
-    ):
+    with patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]):
         mock_proxy_server = SimpleNamespace(
             general_settings={"alerting": ["email"]},
             proxy_logging_obj=mock_proxy_logging_obj,
@@ -41,9 +78,7 @@ async def test_v1_user_creation_no_email_when_send_invite_email_false():
                 user_email="test@example.com",
                 key="sk-test-key",
             )
-            user_api_key_dict = UserAPIKeyAuth(
-                user_id="admin-user", api_key="admin-key"
-            )
+            user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
             await UserManagementEventHooks.async_send_user_invitation_email(
                 data=data,
                 response=response,
@@ -62,9 +97,7 @@ async def test_v1_user_creation_sends_email_when_send_invite_email_true():
     mock_proxy_logging_obj = MagicMock()
     mock_proxy_logging_obj.slack_alerting_instance = mock_slack_alerting
 
-    with patch(
-        "litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]
-    ):
+    with patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]):
         mock_proxy_server = SimpleNamespace(
             general_settings={"alerting": ["email"]},
             proxy_logging_obj=mock_proxy_logging_obj,
@@ -80,9 +113,7 @@ async def test_v1_user_creation_sends_email_when_send_invite_email_true():
                 user_email="test@example.com",
                 key="sk-test-key",
             )
-            user_api_key_dict = UserAPIKeyAuth(
-                user_id="admin-user", api_key="admin-key"
-            )
+            user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
             await UserManagementEventHooks.async_send_user_invitation_email(
                 data=data,
                 response=response,
@@ -225,9 +256,7 @@ async def test_v1_key_generation_sends_email_when_send_invite_email_true():
                 proxy_logging_obj=mock_proxy_logging_obj,
                 litellm_proxy_admin_name="admin-user",
             )
-            with patch.dict(
-                sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}
-            ):
+            with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
                 data = GenerateKeyRequest(
                     user_email="test@example.com",
                     send_invite_email=True,  # Should send key email
@@ -236,9 +265,7 @@ async def test_v1_key_generation_sends_email_when_send_invite_email_true():
                     user_email="test@example.com",
                     key="sk-test-key",
                 )
-                user_api_key_dict = UserAPIKeyAuth(
-                    user_id="admin-user", api_key="admin-key"
-                )
+                user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
                 await KeyManagementEventHooks.async_key_generated_hook(
                     data=data,
                     response=response,
@@ -271,9 +298,7 @@ async def test_v1_key_generation_no_email_when_send_invite_email_false():
                 proxy_logging_obj=mock_proxy_logging_obj,
                 litellm_proxy_admin_name="admin-user",
             )
-            with patch.dict(
-                sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}
-            ):
+            with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
                 data = GenerateKeyRequest(
                     user_email="test@example.com",
                     send_invite_email=False,  # Should NOT send key email
@@ -282,12 +307,130 @@ async def test_v1_key_generation_no_email_when_send_invite_email_false():
                     user_email="test@example.com",
                     key="sk-test-key",
                 )
-                user_api_key_dict = UserAPIKeyAuth(
-                    user_id="admin-user", api_key="admin-key"
-                )
+                user_api_key_dict = UserAPIKeyAuth(user_id="admin-user", api_key="admin-key")
                 await KeyManagementEventHooks.async_key_generated_hook(
                     data=data,
                     response=response,
                     user_api_key_dict=user_api_key_dict,
                 )
                 mock_send_key_created_email.assert_not_called()
+
+
+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)
+
+    sent_html: list[str] = []
+
+    async def capture_send_email(*, receiver_email, subject, html):
+        sent_html.append(html)
+
+    scheduled: list[asyncio.Task] = []
+    real_create_task = asyncio.create_task
+
+    def schedule_task(coro, *args, **kwargs):
+        task = real_create_task(coro, *args, **kwargs)
+        scheduled.append(task)
+        return task
+
+    slack_alerting = SlackAlerting(alerting=["email"])
+    with (
+        patch("litellm.logging_callback_manager.get_custom_loggers_for_type", return_value=[]),
+        patch("litellm.proxy.utils.send_email", capture_send_email),
+        patch(
+            "litellm.proxy.hooks.key_management_event_hooks.asyncio.create_task",
+            side_effect=schedule_task,
+        ),
+        patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(slack_alerting)}),
+    ):
+        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"
+    return sent_html[0]
+
+
+@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."""
+    sent_html = await _render_v0_key_created_email(monkeypatch, _PLAINTEXT_KEY)
+    assert f"Key: <pre>{_PLAINTEXT_KEY}</pre>" 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: <pre>{escaped_key}</pre>" in sent_html
+    assert f'api_key="{escaped_key}"' in sent_html
+    assert '<a href="https://attacker.example/login">' 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"):
+        monkeypatch.delenv(name, raising=False)
+    from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
+        BaseEmailLogger,
+    )
+
+    class RecordingEmailLogger(BaseEmailLogger):
+        def __init__(self):
+            super().__init__()
+            self.html_bodies: list[str] = []
+            self.events: list = []
+
+        async def send_key_created_email(self, send_key_created_email_event):
+            self.events.append(send_key_created_email_event)
+            event = send_key_created_email_event
+            if event.user_email is None:
+                event = event.model_copy(update={"user_email": "user@example.com"})
+            await super().send_key_created_email(send_key_created_email_event=event)
+
+        async def send_email(self, from_email, to_email, subject, html_body):
+            self.html_bodies.append(html_body)
+
+    recording_logger = RecordingEmailLogger()
+    with (
+        patch(
+            "litellm.logging_callback_manager.get_custom_loggers_for_type",
+            return_value=[recording_logger],
+        ),
+        patch.dict(sys.modules, {"litellm.proxy.proxy_server": _proxy_server_module(None)}),
+    ):
+        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"
+    sent_html = recording_logger.html_bodies[0]
+    assert f'<div class="key">{_PLAINTEXT_KEY}</div>' 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'<div class="key">{escaped_key}</div>' in sent_html
+    assert f'"{escaped_key}"' in sent_html
+    assert '<a href="https://attacker.example/login">' not in sent_html
+    assert _TOKEN_ID not in sent_html