PacketStation/server/packetstation/auth/ratelimit.py

"""Sign-in rate limits, checked before any LDAP bind and before any PIN hash check.

Only failed attempts count, and never those refused as `rate_limited`. Per account: when an account has
SIGNIN_ACCOUNT_LIMIT failures inside one window, it is refused until the limit-hitting failure plus the
cooldown. Per client address: SIGNIN_IP_LIMIT failures in the last window refuse it until enough of them
fall out of the window. Times come from the database clock unless the caller passes `now`.
"""

from __future__ import annotations

import math
from collections.abc import Sequence
from datetime import datetime, timedelta
from typing import Any

from ..settings import Settings

RATE_LIMITED = "rate_limited"
KINDS = ("password", "pin", "unlock", "signature", "pin_set")


def account_block_until(
    failures: Sequence[datetime], *, limit: int, window: timedelta, cooldown: timedelta
) -> datetime | None:
    """The latest "refused until" time implied by a list of failure times, or None."""
    times = sorted(failures)
    until: datetime | None = None
    for i in range(limit - 1, len(times)):
        if times[i] - times[i - limit + 1] < window:
            candidate = times[i] + cooldown
            if until is None or candidate > until:
                until = candidate
    return until


def ip_block_until(failures: Sequence[datetime], *, now: datetime, limit: int, window: timedelta) -> datetime | None:
    """When the address may try again, or None. Only failures inside (now - window, now] count."""
    recent = sorted(t for t in failures if now - window < t <= now)
    if len(recent) < limit:
        return None
    # Allowed again once the count drops below the limit: the (n - limit + 1)th oldest leaves the window
    return recent[len(recent) - limit] + window


def retry_seconds(until: datetime | None, now: datetime) -> int | None:
    if until is None or until <= now:
        return None
    return max(1, math.ceil((until - now).total_seconds()))


def db_now(conn: Any) -> datetime:
    row = conn.execute("SELECT now()").fetchone()
    now: datetime = row[0]
    return now


def check(
    conn: Any,
    settings: Settings,
    *,
    accounts: Sequence[str],
    ip: str | None,
    now: datetime | None = None,
) -> int | None:
    """Seconds until another attempt is allowed, or None when allowed. Failures of all `accounts` count together."""
    now = now or db_now(conn)
    window = timedelta(minutes=settings.signin_account_window_minutes)
    cooldown = timedelta(minutes=settings.signin_account_cooldown_minutes)
    waits: list[int] = []
    keys = sorted({a for a in accounts if a})
    if keys:
        since = now - window - cooldown
        rows = conn.execute(
            "SELECT at FROM signin_attempts WHERE account = ANY(%s) AND NOT ok"
            " AND reason IS DISTINCT FROM %s AND at > %s AND at <= %s",
            (keys, RATE_LIMITED, since, now),
        ).fetchall()
        until = account_block_until(
            [r[0] for r in rows], limit=settings.signin_account_limit, window=window, cooldown=cooldown
        )
        if (w := retry_seconds(until, now)) is not None:
            waits.append(w)
    if ip:
        ip_window = timedelta(minutes=settings.signin_ip_window_minutes)
        rows = conn.execute(
            "SELECT at FROM signin_attempts WHERE ip = %s::inet AND NOT ok"
            " AND reason IS DISTINCT FROM %s AND at > %s AND at <= %s",
            (ip, RATE_LIMITED, now - ip_window, now),
        ).fetchall()
        until = ip_block_until([r[0] for r in rows], now=now, limit=settings.signin_ip_limit, window=ip_window)
        if (w := retry_seconds(until, now)) is not None:
            waits.append(w)
    return max(waits) if waits else None


def record(
    conn: Any,
    *,
    account: str,
    ip: str | None,
    kind: str,
    ok: bool,
    reason: str | None = None,
    user_id: Any = None,
    session_id: Any = None,
    at: datetime | None = None,
) -> None:
    """Append one signin_attempts row. Never holds a password or PIN."""
    if kind not in KINDS:
        raise ValueError(f"bad attempt kind {kind!r}")
    conn.execute(
        "INSERT INTO signin_attempts (at, account, ip, kind, ok, reason, user_id, session_id)"
        " VALUES (coalesce(%s, now()), %s, %s::inet, %s, %s, %s, %s, %s)",
        (at, account, ip, kind, ok, reason, user_id, session_id),
    )