mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
refactor(auth): type the breach-screen DB dicts and flatten the session-key kwargs
Annotate screen_login_password_for_breach's update/where dicts with
prisma input TypedDicts and replace authenticate_user's conditional
dict splat with plain keyword arguments, clearing the LIT002 lines
this branch added in login_utils.py. No behavior change: an unflagged
login now passes allowed_routes=None and metadata={} explicitly, which
are the parameter defaults
This commit is contained in:
parent
7e8683aef0
commit
c6f724b56f
2 changed files with 22 additions and 27 deletions
|
|
@ -10,7 +10,7 @@ import secrets
|
|||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, cast
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast
|
||||
|
||||
import jwt
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -47,6 +47,9 @@ from litellm.repositories.user_repository import UserRepository
|
|||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import types as prisma_types
|
||||
|
||||
BREACH_RECHECK_INTERVAL: Final = timedelta(hours=24)
|
||||
PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change",)
|
||||
|
||||
|
|
@ -80,12 +83,15 @@ async def screen_login_password_for_breach(
|
|||
if not _breach_recheck_due(last_breach_check_at):
|
||||
return False
|
||||
breached: Final = await is_password_breached(password, general_settings, client)
|
||||
update_data: Final = {
|
||||
"last_breach_check_at": datetime.now(timezone.utc),
|
||||
**({"password_reset_required": True} if breached else {}),
|
||||
}
|
||||
checked_at: Final = datetime.now(timezone.utc)
|
||||
update_data: Final[prisma_types.LiteLLM_UserTableUpdateInput] = (
|
||||
{"last_breach_check_at": checked_at, "password_reset_required": True}
|
||||
if breached
|
||||
else {"last_breach_check_at": checked_at}
|
||||
)
|
||||
find_user: Final[prisma_types.LiteLLM_UserTableWhereInput] = {"user_id": user_id}
|
||||
try:
|
||||
await UserRepository(prisma_client).table.update(where={"user_id": user_id}, data=update_data)
|
||||
await UserRepository(prisma_client).table.update(where=find_user, data=update_data)
|
||||
except Exception as e: # noqa: BLE001 # a failed stamp must never surface into the login
|
||||
verbose_proxy_logger.warning("Login-time breach screening could not update user %s: %s", user_id, e)
|
||||
return breached
|
||||
|
|
@ -382,25 +388,14 @@ async def authenticate_user(
|
|||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_role,
|
||||
"duration": LITELLM_UI_SESSION_DURATION,
|
||||
"key_max_budget": litellm.max_ui_session_budget,
|
||||
"models": [],
|
||||
"aliases": {},
|
||||
"config": {},
|
||||
"spend": 0,
|
||||
"user_id": user_id,
|
||||
"team_id": "litellm-dashboard",
|
||||
**(
|
||||
{
|
||||
"allowed_routes": list(PASSWORD_RESET_ALLOWED_ROUTES),
|
||||
"metadata": {"password_reset_required": True},
|
||||
}
|
||||
if password_reset_required
|
||||
else {}
|
||||
),
|
||||
},
|
||||
user_role=user_role,
|
||||
duration=LITELLM_UI_SESSION_DURATION,
|
||||
key_max_budget=litellm.max_ui_session_budget,
|
||||
spend=0,
|
||||
user_id=user_id,
|
||||
team_id="litellm-dashboard",
|
||||
allowed_routes=list(PASSWORD_RESET_ALLOWED_ROUTES) if password_reset_required else None,
|
||||
metadata={"password_reset_required": True} if password_reset_required else {},
|
||||
)
|
||||
else:
|
||||
raise ProxyException(
|
||||
|
|
|
|||
|
|
@ -1032,8 +1032,8 @@ class TestPasswordResetRequiredSessionMinting:
|
|||
row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=None)
|
||||
result, key_kwargs = await self._login(_prisma_with_user(row))
|
||||
|
||||
assert "allowed_routes" not in key_kwargs
|
||||
assert "metadata" not in key_kwargs
|
||||
assert key_kwargs["allowed_routes"] is None
|
||||
assert not key_kwargs["metadata"]
|
||||
assert result.password_reset_required is False
|
||||
|
||||
async def _login_with_screen_result(self, mock_prisma_client, breached: bool) -> tuple[LoginResult, dict, dict]:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue