saksham
1a4c80958f
FastAPI + Celery + Next.js + Postgres/Redis app with company monitoring, source collection, LLM-based change analysis, enrichment, and account security (Turnstile, escalating lockout, email verification).
65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.enums import ThrottleAction
|
|
from app.models.ip_ban import IpBan
|
|
from app.models.ip_throttle_state import IpThrottleState
|
|
|
|
|
|
class IpThrottleRepository:
|
|
def __init__(self, db: AsyncSession) -> None:
|
|
self.db = db
|
|
|
|
async def get_state(self, ip_address: str, action: ThrottleAction) -> IpThrottleState | None:
|
|
result = await self.db.execute(
|
|
select(IpThrottleState).where(
|
|
IpThrottleState.ip_address == ip_address, IpThrottleState.action == action
|
|
)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def get_or_create_state(self, ip_address: str, action: ThrottleAction) -> IpThrottleState:
|
|
state = await self.get_state(ip_address, action)
|
|
if state is not None:
|
|
return state
|
|
state = IpThrottleState(ip_address=ip_address, action=action)
|
|
self.db.add(state)
|
|
await self.db.flush()
|
|
return state
|
|
|
|
async def get_ban(self, ip_address: str) -> IpBan | None:
|
|
result = await self.db.execute(select(IpBan).where(IpBan.ip_address == ip_address))
|
|
return result.scalar_one_or_none()
|
|
|
|
async def create_ban(self, ip_address: str, reason: str, banned_at: datetime) -> IpBan:
|
|
ban = IpBan(ip_address=ip_address, reason=reason, banned_at=banned_at)
|
|
self.db.add(ban)
|
|
await self.db.flush()
|
|
return ban
|
|
|
|
async def list_bans(self) -> list[IpBan]:
|
|
result = await self.db.execute(select(IpBan).order_by(IpBan.banned_at.desc()))
|
|
return list(result.scalars().all())
|
|
|
|
async def clear_ban_and_state(self, ip_address: str) -> bool:
|
|
"""Full pardon for an admin-approved unban - removes the ban record
|
|
and resets every throttle-state row for this IP to a clean slate
|
|
(not just lifting the terminal ban while leaving them one offense
|
|
from another)."""
|
|
ban_result = await self.db.execute(delete(IpBan).where(IpBan.ip_address == ip_address))
|
|
result = await self.db.execute(
|
|
select(IpThrottleState).where(IpThrottleState.ip_address == ip_address)
|
|
)
|
|
states = result.scalars().all()
|
|
for state in states:
|
|
state.attempt_count = 0
|
|
state.next_allowed_at = None
|
|
state.timeout_until = None
|
|
state.offense_count = 0
|
|
await self.db.flush()
|
|
return ban_result.rowcount > 0 or len(states) > 0
|