162 lines
4.6 KiB
Python
162 lines
4.6 KiB
Python
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import case, delete, or_, select
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ..models import SecurityThrottle
|
|
|
|
EMAIL_CODE_VERIFY_SCOPE = "email_code_verify"
|
|
EMAIL_PASSWORD_LOGIN_SCOPE = "email_password_login"
|
|
PROMO_CODE_APPLY_SCOPE = "promo_code_apply"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ThrottleDecision:
|
|
locked: bool
|
|
retry_after: Optional[int] = None
|
|
|
|
|
|
def _utc_now(value: Optional[datetime] = None) -> datetime:
|
|
if value is None:
|
|
value = datetime.now(timezone.utc)
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=timezone.utc)
|
|
return value.astimezone(timezone.utc)
|
|
|
|
|
|
def _retry_after_seconds(locked_until: Optional[datetime], now: datetime) -> Optional[int]:
|
|
if not locked_until:
|
|
return None
|
|
locked_until = _utc_now(locked_until)
|
|
remaining = int((locked_until - now).total_seconds())
|
|
return max(1, remaining) if remaining > 0 else None
|
|
|
|
|
|
async def get_throttle_state(
|
|
session: AsyncSession,
|
|
*,
|
|
scope: str,
|
|
identifier: str,
|
|
) -> Optional[SecurityThrottle]:
|
|
stmt = (
|
|
select(SecurityThrottle)
|
|
.where(
|
|
SecurityThrottle.scope == scope,
|
|
SecurityThrottle.identifier == identifier,
|
|
)
|
|
.limit(1)
|
|
)
|
|
result = await session.execute(stmt)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def check_throttle(
|
|
session: AsyncSession,
|
|
*,
|
|
scope: str,
|
|
identifier: str,
|
|
now: Optional[datetime] = None,
|
|
) -> ThrottleDecision:
|
|
now = _utc_now(now)
|
|
row = await get_throttle_state(session, scope=scope, identifier=identifier)
|
|
if not row or not row.locked_until:
|
|
return ThrottleDecision(locked=False)
|
|
|
|
locked_until = _utc_now(row.locked_until)
|
|
if locked_until <= now:
|
|
return ThrottleDecision(locked=False)
|
|
|
|
return ThrottleDecision(
|
|
locked=True,
|
|
retry_after=_retry_after_seconds(locked_until, now),
|
|
)
|
|
|
|
|
|
async def record_throttle_failure(
|
|
session: AsyncSession,
|
|
*,
|
|
scope: str,
|
|
identifier: str,
|
|
max_failures: int,
|
|
window_seconds: int,
|
|
lock_seconds: int,
|
|
now: Optional[datetime] = None,
|
|
) -> ThrottleDecision:
|
|
now = _utc_now(now)
|
|
max_failures = max(1, int(max_failures))
|
|
window_seconds = max(1, int(window_seconds))
|
|
lock_seconds = max(1, int(lock_seconds))
|
|
window_cutoff = now - timedelta(seconds=window_seconds)
|
|
lock_until = now + timedelta(seconds=lock_seconds)
|
|
|
|
failure_count_expr = case(
|
|
(
|
|
or_(
|
|
SecurityThrottle.window_started_at.is_(None),
|
|
SecurityThrottle.window_started_at <= window_cutoff,
|
|
),
|
|
1,
|
|
),
|
|
else_=SecurityThrottle.failures + 1,
|
|
)
|
|
|
|
stmt = (
|
|
pg_insert(SecurityThrottle)
|
|
.values(
|
|
scope=scope,
|
|
identifier=identifier,
|
|
failures=1,
|
|
window_started_at=now,
|
|
last_attempt_at=now,
|
|
locked_until=lock_until if max_failures <= 1 else None,
|
|
)
|
|
.on_conflict_do_update(
|
|
index_elements=[SecurityThrottle.scope, SecurityThrottle.identifier],
|
|
set_={
|
|
"failures": failure_count_expr,
|
|
"window_started_at": case(
|
|
(
|
|
or_(
|
|
SecurityThrottle.window_started_at.is_(None),
|
|
SecurityThrottle.window_started_at <= window_cutoff,
|
|
),
|
|
now,
|
|
),
|
|
else_=SecurityThrottle.window_started_at,
|
|
),
|
|
"last_attempt_at": now,
|
|
"locked_until": case(
|
|
(failure_count_expr >= max_failures, lock_until),
|
|
else_=None,
|
|
),
|
|
},
|
|
)
|
|
.returning(SecurityThrottle.locked_until)
|
|
)
|
|
|
|
result = await session.execute(stmt)
|
|
locked_until = result.scalar_one_or_none()
|
|
locked_until = _utc_now(locked_until) if locked_until else None
|
|
if locked_until and locked_until > now:
|
|
return ThrottleDecision(
|
|
locked=True,
|
|
retry_after=_retry_after_seconds(locked_until, now),
|
|
)
|
|
return ThrottleDecision(locked=False)
|
|
|
|
|
|
async def clear_throttle_state(
|
|
session: AsyncSession,
|
|
*,
|
|
scope: str,
|
|
identifier: str,
|
|
) -> None:
|
|
stmt = delete(SecurityThrottle).where(
|
|
SecurityThrottle.scope == scope,
|
|
SecurityThrottle.identifier == identifier,
|
|
)
|
|
await session.execute(stmt)
|