150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
"""Durable throttling state for operator sign-in attempts."""
|
|
|
|
import hashlib
|
|
import ipaddress
|
|
import math
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
class LoginAttemptStoreError(RuntimeError):
|
|
"""Raised when sign-in throttling state cannot be accessed safely."""
|
|
|
|
|
|
def client_source(peer_host: str, forwarded_for: str, trusted_proxy_cidrs: str) -> str:
|
|
"""Resolve a canonical client IP without trusting arbitrary forwarding headers."""
|
|
try:
|
|
peer = ipaddress.ip_address(peer_host)
|
|
trusted = [
|
|
ipaddress.ip_network(value.strip())
|
|
for value in trusted_proxy_cidrs.split(",")
|
|
if value.strip()
|
|
]
|
|
except ValueError:
|
|
return peer_host
|
|
if not any(peer in network for network in trusted):
|
|
return str(peer)
|
|
try:
|
|
forwarded = [
|
|
ipaddress.ip_address(value.strip())
|
|
for value in forwarded_for.split(",")
|
|
if value.strip()
|
|
]
|
|
except ValueError:
|
|
return str(peer)
|
|
for address in reversed(forwarded):
|
|
if not any(address in network for network in trusted):
|
|
return str(address)
|
|
return str(forwarded[0]) if forwarded else str(peer)
|
|
|
|
|
|
class LoginAttemptStore:
|
|
def __init__(
|
|
self,
|
|
path: str | Path,
|
|
*,
|
|
clock: Callable[[], float],
|
|
max_failures: int,
|
|
window_seconds: int,
|
|
max_entries: int = 10_000,
|
|
lock_timeout_seconds: float = 0.1,
|
|
) -> None:
|
|
self.path = Path(path)
|
|
self.clock = clock
|
|
self.max_failures = max(1, max_failures)
|
|
self.window_seconds = max(1, window_seconds)
|
|
self.max_entries = max(1, max_entries)
|
|
self.lock_timeout_seconds = lock_timeout_seconds
|
|
|
|
@staticmethod
|
|
def _digest(source: str) -> str:
|
|
return hashlib.sha256(source.encode()).hexdigest()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
try:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(self.path, timeout=self.lock_timeout_seconds)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS login_attempts (
|
|
source_hash TEXT PRIMARY KEY,
|
|
failures INTEGER NOT NULL,
|
|
window_started_at REAL NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
return connection
|
|
except (OSError, sqlite3.Error) as exc:
|
|
raise LoginAttemptStoreError(
|
|
"Sign-in throttling is temporarily unavailable"
|
|
) from exc
|
|
|
|
def retry_after(self, source: str) -> int:
|
|
now = self.clock()
|
|
try:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"SELECT failures, window_started_at FROM login_attempts WHERE source_hash = ?",
|
|
(self._digest(source),),
|
|
).fetchone()
|
|
except (OSError, sqlite3.Error) as exc:
|
|
raise LoginAttemptStoreError(
|
|
"Sign-in throttling is temporarily unavailable"
|
|
) from exc
|
|
if row is None or row[0] < self.max_failures:
|
|
return 0
|
|
remaining = row[1] + self.window_seconds - now
|
|
return max(0, math.ceil(remaining))
|
|
|
|
def record_failure(self, source: str) -> None:
|
|
now = self.clock()
|
|
source_hash = self._digest(source)
|
|
try:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
connection.execute(
|
|
"DELETE FROM login_attempts WHERE window_started_at + ? <= ?",
|
|
(self.window_seconds, now),
|
|
)
|
|
row = connection.execute(
|
|
"SELECT failures, window_started_at FROM login_attempts WHERE source_hash = ?",
|
|
(source_hash,),
|
|
).fetchone()
|
|
if row is None or row[1] + self.window_seconds <= now:
|
|
connection.execute(
|
|
"INSERT OR REPLACE INTO login_attempts VALUES (?, 1, ?)",
|
|
(source_hash, now),
|
|
)
|
|
else:
|
|
connection.execute(
|
|
"UPDATE login_attempts SET failures = failures + 1 WHERE source_hash = ?",
|
|
(source_hash,),
|
|
)
|
|
connection.execute(
|
|
"""
|
|
DELETE FROM login_attempts
|
|
WHERE source_hash NOT IN (
|
|
SELECT source_hash FROM login_attempts
|
|
ORDER BY window_started_at DESC, rowid DESC LIMIT ?
|
|
)
|
|
""",
|
|
(self.max_entries,),
|
|
)
|
|
except (OSError, sqlite3.Error) as exc:
|
|
raise LoginAttemptStoreError(
|
|
"Sign-in throttling is temporarily unavailable"
|
|
) from exc
|
|
|
|
def clear(self, source: str) -> None:
|
|
try:
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"DELETE FROM login_attempts WHERE source_hash = ?",
|
|
(self._digest(source),),
|
|
)
|
|
except (OSError, sqlite3.Error) as exc:
|
|
raise LoginAttemptStoreError(
|
|
"Sign-in throttling is temporarily unavailable"
|
|
) from exc
|