refactor: promo and email bruteforce defence

db/dal/__init__.py

@@ -6,6 +6,7 @@ from . import panel_sync_dal
 from . import message_log_dal
 from . import user_billing_dal
 from . import ad_dal
+from . import security_dal
 
 __all__ = (
     "user_dal",
@@ -16,6 +17,7 @@ __all__ = (
     "message_log_dal",
     "user_billing_dal",
     "ad_dal",
+    "security_dal",
 )
 
 

db/dal/security_dal.py

@@ -0,0 +1,160 @@
+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"
+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)

db/migrator.py

@@ -202,6 +202,43 @@ def _migration_0005_add_email_auth_fields(connection: Connection) -> None:
         )
     )
 
+
+def _migration_0006_add_security_throttles(connection: Connection) -> None:
+    connection.execute(
+        text(
+            """
+            CREATE TABLE IF NOT EXISTS security_throttles (
+                throttle_id SERIAL PRIMARY KEY,
+                scope VARCHAR(64) NOT NULL,
+                identifier VARCHAR(512) NOT NULL,
+                failures INTEGER NOT NULL DEFAULT 0,
+                window_started_at TIMESTAMPTZ NULL,
+                locked_until TIMESTAMPTZ NULL,
+                last_attempt_at TIMESTAMPTZ NULL,
+                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                updated_at TIMESTAMPTZ NULL,
+                CONSTRAINT uq_security_throttles_scope_identifier UNIQUE (scope, identifier)
+            )
+            """
+        )
+    )
+    connection.execute(
+        text(
+            """
+            CREATE INDEX IF NOT EXISTS ix_security_throttles_scope
+            ON security_throttles (scope)
+            """
+        )
+    )
+    connection.execute(
+        text(
+            """
+            CREATE INDEX IF NOT EXISTS ix_security_throttles_locked_until
+            ON security_throttles (locked_until)
+            """
+        )
+    )
+
 MIGRATIONS: List[Migration] = [
     Migration(
         id="0001_add_channel_subscription_fields",
@@ -228,6 +265,11 @@ MIGRATIONS: List[Migration] = [
         description="Add email login identities and verification codes",
         upgrade=_migration_0005_add_email_auth_fields,
     ),
+    Migration(
+        id="0006_add_security_throttles",
+        description="Add generic lockout tracking for brute-force protection",
+        upgrade=_migration_0006_add_security_throttles,
+    ),
 ]
 
 

db/models.py

@@ -110,6 +110,24 @@ class EmailVerificationCode(Base):
     target_user = relationship("User")
 
 
+class SecurityThrottle(Base):
+    __tablename__ = "security_throttles"
+
+    throttle_id = Column(Integer, primary_key=True, autoincrement=True)
+    scope = Column(String(64), nullable=False, index=True)
+    identifier = Column(String(512), nullable=False, index=True)
+    failures = Column(Integer, nullable=False, default=0)
+    window_started_at = Column(DateTime(timezone=True), nullable=True)
+    locked_until = Column(DateTime(timezone=True), nullable=True, index=True)
+    last_attempt_at = Column(DateTime(timezone=True), nullable=True)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+    updated_at = Column(DateTime(timezone=True), onupdate=func.now(), nullable=True)
+
+    __table_args__ = (
+        UniqueConstraint("scope", "identifier", name="uq_security_throttles_scope_identifier"),
+    )
+
+
 class Payment(Base):
     __tablename__ = "payments"