first commit

2026-07-25 00:37:09 +03:00
commit 3ca7f9f7cc
35 changed files with 2478 additions and 0 deletions
@@ -0,0 +1,94 @@
+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