Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 24 additions & 6 deletions litellm/proxy/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import json
import os
import smtplib
import ssl
import sys
import threading
import time
Expand Down Expand Up @@ -5134,6 +5135,23 @@ async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient):
return


def _should_use_smtp_ssl(smtp_port: int) -> bool:
"""
Port 465 expects an immediate TLS handshake (implicit SSL), so a plain
smtplib.SMTP connection hangs waiting for an SMTP banner. Use SMTP_SSL
there, or when SMTP_USE_SSL is explicitly enabled.
"""
return os.getenv("SMTP_USE_SSL", "False") == "True" or smtp_port == 465


def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP:
if _should_use_smtp_ssl(smtp_port=smtp_port):
return smtplib.SMTP_SSL(
host=smtp_host, port=smtp_port, context=ssl.create_default_context()
)
return smtplib.SMTP(host=smtp_host, port=smtp_port)


async def send_email(
receiver_email: Optional[str] = None,
subject: Optional[str] = None,
Expand Down Expand Up @@ -5179,13 +5197,13 @@ async def send_email(
email_message.attach(MIMEText(html, "html"))

try:
# Establish a secure connection with the SMTP server
with smtplib.SMTP(
host=smtp_host,
port=smtp_port,
using_ssl = _should_use_smtp_ssl(smtp_port=smtp_port)
with _create_smtp_connection(
smtp_host=smtp_host,
smtp_port=smtp_port,
) as server:
if os.getenv("SMTP_TLS", "True") != "False":
server.starttls()
if not using_ssl and os.getenv("SMTP_TLS", "True") != "False":
server.starttls(context=ssl.create_default_context())

# Login to your email account only if smtp_username and smtp_password are provided
if smtp_username and smtp_password:
Expand Down
97 changes: 95 additions & 2 deletions tests/test_litellm/proxy/test_proxy_utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import datetime as real_datetime
import json
import os
import smtplib
import sys

import pytest
Expand All @@ -15,7 +15,7 @@
) # Adds the parent directory to the system path


from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch

from litellm.proxy.utils import get_custom_url, join_paths

Expand Down Expand Up @@ -368,3 +368,96 @@ async def test_logging_obj_without_anchor_is_noop(self):
await self._run(request_data)
assert "first_api_call_start_time" not in request_data
assert "litellm_logging_obj" not in request_data


class TestShouldUseSmtpSsl:
def test_port_465_uses_ssl(self, monkeypatch):
from litellm.proxy.utils import _should_use_smtp_ssl

monkeypatch.delenv("SMTP_USE_SSL", raising=False)
assert _should_use_smtp_ssl(smtp_port=465) is True

def test_smtp_use_ssl_env_var_forces_ssl_on_any_port(self, monkeypatch):
from litellm.proxy.utils import _should_use_smtp_ssl

monkeypatch.setenv("SMTP_USE_SSL", "True")
assert _should_use_smtp_ssl(smtp_port=2465) is True

def test_port_587_uses_plain_smtp(self, monkeypatch):
from litellm.proxy.utils import _should_use_smtp_ssl

monkeypatch.delenv("SMTP_USE_SSL", raising=False)
assert _should_use_smtp_ssl(smtp_port=587) is False
Comment thread
mubashir1osmani marked this conversation as resolved.


class TestCreateSmtpConnection:
def test_port_465_creates_smtp_ssl_with_verified_context(self, monkeypatch):
import ssl

from litellm.proxy.utils import _create_smtp_connection

monkeypatch.delenv("SMTP_USE_SSL", raising=False)
with (
patch("smtplib.SMTP_SSL") as mock_smtp_ssl,
patch("smtplib.SMTP") as mock_smtp,
):
result = _create_smtp_connection(
smtp_host="mail.example.com", smtp_port=465
)

mock_smtp.assert_not_called()
assert result is mock_smtp_ssl.return_value
_, kwargs = mock_smtp_ssl.call_args
assert kwargs["host"] == "mail.example.com"
assert kwargs["port"] == 465
context = kwargs["context"]
assert isinstance(context, ssl.SSLContext)
assert context.verify_mode == ssl.CERT_REQUIRED
assert context.check_hostname is True

def test_port_587_creates_plain_smtp(self, monkeypatch):
from litellm.proxy.utils import _create_smtp_connection

monkeypatch.delenv("SMTP_USE_SSL", raising=False)
with (
patch("smtplib.SMTP_SSL") as mock_smtp_ssl,
patch("smtplib.SMTP") as mock_smtp,
):
result = _create_smtp_connection(
smtp_host="mail.example.com", smtp_port=587
)

mock_smtp_ssl.assert_not_called()
assert result is mock_smtp.return_value
mock_smtp.assert_called_once_with(host="mail.example.com", port=587)


class TestSendEmailStartTls:
@pytest.mark.asyncio
async def test_starttls_uses_verified_context(self, monkeypatch):
import ssl

from litellm.proxy.utils import send_email

monkeypatch.setenv("SMTP_HOST", "mail.example.com")
monkeypatch.setenv("SMTP_PORT", "587")
monkeypatch.setenv("SMTP_SENDER_EMAIL", "sender@example.com")
monkeypatch.delenv("SMTP_TLS", raising=False)
monkeypatch.delenv("SMTP_USE_SSL", raising=False)

mock_server = MagicMock(spec=smtplib.SMTP)
with patch(
"litellm.proxy.utils._create_smtp_connection"
) as mock_create_connection:
mock_create_connection.return_value.__enter__.return_value = mock_server
await send_email(
receiver_email="receiver@example.com",
subject="test",
html="<p>test</p>",
)

_, kwargs = mock_server.starttls.call_args
context = kwargs["context"]
assert isinstance(context, ssl.SSLContext)
assert context.verify_mode == ssl.CERT_REQUIRED
assert context.check_hostname is True
8 changes: 5 additions & 3 deletions tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,7 @@ def __enter__(self) -> "_Conn":
def __exit__(self, *exc: Any) -> None:
return None

def starttls(self) -> None:
def starttls(self, **kwargs: Any) -> None:
self._starttls_called = True

def login(self, user: str, password: str) -> None:
Expand Down Expand Up @@ -378,10 +378,12 @@ def _factory(*args: Any, **kwargs: Any) -> _Conn:

@pytest.fixture
def in_memory_smtp(monkeypatch: pytest.MonkeyPatch) -> InMemorySMTP:
"""Patch ``smtplib.SMTP`` to capture sends in memory.
"""Patch ``smtplib.SMTP`` and ``smtplib.SMTP_SSL`` to capture sends in memory.

Override ``smtp.raise_on_send`` to test the SMTP error path.
"""
smtp = InMemorySMTP()
monkeypatch.setattr("smtplib.SMTP", smtp.server_factory())
factory = smtp.server_factory()
monkeypatch.setattr("smtplib.SMTP", factory)
monkeypatch.setattr("smtplib.SMTP_SSL", factory)
return smtp
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,12 @@
@pytest.fixture(autouse=True)
def _smtp_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("SMTP_HOST", "smtp.invalid")
monkeypatch.setenv("SMTP_PORT", "2525")
monkeypatch.setenv("SMTP_PORT", "587")
monkeypatch.setenv("SMTP_USERNAME", "u")
monkeypatch.setenv("SMTP_PASSWORD", "p")
monkeypatch.setenv("SMTP_SENDER_EMAIL", "from@invalid")
monkeypatch.setenv("SMTP_TLS", "True")
monkeypatch.setenv("SMTP_USE_SSL", "False")


@pytest.mark.asyncio
Expand Down Expand Up @@ -50,10 +51,10 @@ async def test_send_email_dispatches_via_smtp(in_memory_smtp: Any) -> None:


@pytest.mark.asyncio
async def test_send_email_skips_starttls_when_disabled(
async def test_send_email_starttls_uses_ssl(
in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("SMTP_TLS", "False")
monkeypatch.setenv("SMTP_USE_SSL", "True")
await send_email(
receiver_email="to@invalid",
subject="Hi",
Expand Down
Loading