95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from database import get_db
|
|
from remnawave import RemnawaveClient, RemnawaveError
|
|
|
|
logger = logging.getLogger("traffic_grants")
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
|
|
|
|
|
|
async def revert_expired_grants() -> int:
|
|
"""Снимает временный трафик после expires_at. Возвращает число успешных откатов."""
|
|
now = _now_iso()
|
|
db = await get_db()
|
|
try:
|
|
cursor = await db.execute(
|
|
"""
|
|
SELECT * FROM traffic_grants
|
|
WHERE status = 'active' AND expires_at <= ?
|
|
ORDER BY id ASC
|
|
""",
|
|
(now,),
|
|
)
|
|
grants = [dict(row) for row in await cursor.fetchall()]
|
|
finally:
|
|
await db.close()
|
|
|
|
if not grants:
|
|
return 0
|
|
|
|
client = RemnawaveClient()
|
|
done = 0
|
|
for grant in grants:
|
|
try:
|
|
result = await client.remove_traffic_bytes(
|
|
grant["user_uuid"], int(grant["bytes_added"])
|
|
)
|
|
db = await get_db()
|
|
try:
|
|
await db.execute(
|
|
"""
|
|
UPDATE traffic_grants
|
|
SET status = 'reverted', reverted_at = ?, last_error = ''
|
|
WHERE id = ?
|
|
""",
|
|
(_now_iso(), grant["id"]),
|
|
)
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
done += 1
|
|
logger.info(
|
|
"Reverted traffic grant #%s user=%s -%s bytes (limit %s → %s)",
|
|
grant["id"],
|
|
grant.get("username") or grant["user_uuid"],
|
|
grant["bytes_added"],
|
|
result.get("old_limit_bytes"),
|
|
result.get("new_limit_bytes"),
|
|
)
|
|
except RemnawaveError as exc:
|
|
db = await get_db()
|
|
try:
|
|
await db.execute(
|
|
"""
|
|
UPDATE traffic_grants
|
|
SET last_error = ?
|
|
WHERE id = ?
|
|
""",
|
|
(str(exc)[:500], grant["id"]),
|
|
)
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
logger.warning("Failed to revert grant #%s: %s", grant["id"], exc)
|
|
|
|
return done
|
|
|
|
|
|
async def traffic_grant_worker(stop_event: asyncio.Event) -> None:
|
|
while not stop_event.is_set():
|
|
try:
|
|
await revert_expired_grants()
|
|
except Exception:
|
|
logger.exception("traffic grant worker error")
|
|
try:
|
|
await asyncio.wait_for(stop_event.wait(), timeout=60)
|
|
except asyncio.TimeoutError:
|
|
pass
|