DispelEvilBot/common/repositories/admin_login_tokens_repo.py

74 lines
2.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Одноразовые токены входа в админку через Telegram deep-link."""
from __future__ import annotations
import hashlib
import secrets
from common.db.pool import get_pool
#: Время жизни токена.
TTL_MINUTES = 5
def _hash(token: str) -> str:
return hashlib.sha256(token.encode()).hexdigest()
async def create_token() -> str:
"""Создать токен, вернуть его в открытом виде (в БД — только хеш)."""
token = secrets.token_urlsafe(32)
async with get_pool().connection() as conn:
await conn.execute(
"""
INSERT INTO bot.admin_login_tokens (token_hash, expires_at)
VALUES (%s, now() + make_interval(mins => %s))
""",
(_hash(token), TTL_MINUTES),
)
await conn.execute("DELETE FROM bot.admin_login_tokens WHERE expires_at < now() - interval '1 day'")
return token
async def confirm_token(token: str, telegram_id: int) -> bool:
"""Отметить токен подтверждённым ботом. False — если истёк/уже подтверждён."""
async with get_pool().connection() as conn:
cur = await conn.execute(
"""
UPDATE bot.admin_login_tokens
SET confirmed_at = now(), telegram_id = %s
WHERE token_hash = %s
AND confirmed_at IS NULL
AND used_at IS NULL
AND expires_at > now()
""",
(telegram_id, _hash(token)),
)
return bool(cur.rowcount)
async def consume_token(token: str) -> int | None:
"""Забрать подтверждённый токен один раз. Возвращает Telegram-ID или None."""
async with get_pool().connection() as conn:
cur = await conn.execute(
"""
SELECT telegram_id
FROM bot.admin_login_tokens
WHERE token_hash = %s
AND confirmed_at IS NOT NULL
AND used_at IS NULL
AND expires_at > now()
FOR UPDATE
""",
(_hash(token),),
)
row = await cur.fetchone()
if row is None or row["telegram_id"] is None:
return None
await conn.execute(
"UPDATE bot.admin_login_tokens SET used_at = now() WHERE token_hash = %s",
(_hash(token),),
)
return int(row["telegram_id"])