From eaa83d586587d1225b6dd821c4e4a1c7f1ebc2c1 Mon Sep 17 00:00:00 2001
From: bain <bain@bain.cz>
Date: Thu, 31 Oct 2024 22:36:11 +0100
Subject: [PATCH] add option to specify num of tokens in cmd; quota bypass list

---
 matrix-invitation-dealer/env.py  |  2 +
 matrix-invitation-dealer/main.py | 68 +++++++++++++++++++++-----------
 2 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/matrix-invitation-dealer/env.py b/matrix-invitation-dealer/env.py
index 17ec1d0..334d06d 100644
--- a/matrix-invitation-dealer/env.py
+++ b/matrix-invitation-dealer/env.py
@@ -67,3 +67,5 @@ INVITE_CODE_QUOTA_TIMESPAN = td_parse(icq_timespan)
 INVITE_CODE_EXPIRATION = env(
     td_parse, "INVITE_CODE_EXPIRATION", datetime.timedelta(days=7)
 )
+
+USERS_WITHOUT_QUOTA = env(lambda x: x.split(), "USERS_WITHOUT_QUOTA", [])
diff --git a/matrix-invitation-dealer/main.py b/matrix-invitation-dealer/main.py
index 3e4f5ea..c36bbeb 100644
--- a/matrix-invitation-dealer/main.py
+++ b/matrix-invitation-dealer/main.py
@@ -1,10 +1,12 @@
+import asyncio
 import datetime
+import logging
 import re
 import time
-import asyncio
+from collections import deque
 from typing import Deque, Optional
-import aiosqlite
 
+import aiosqlite
 from nio import (
     AsyncClient,
     InviteMemberEvent,
@@ -14,8 +16,6 @@ from nio import (
     RoomMessage,
     RoomMessageText,
 )
-from collections import deque
-import logging
 
 from . import env
 from .admin import SynapseAdmin
@@ -82,7 +82,10 @@ class Bot:
 
         user = event.sender
 
-        if type(event) is not RoomMessageText or event.body != "!new":
+        if (
+            type(event) is not RoomMessageText
+            or (command := event.body.split())[0] != "!new"
+        ):
             await self.send_hi_message(user, room)
             return
 
@@ -101,23 +104,38 @@ class Bot:
             )
             return
 
-        if await self.quota_exceeded(user):
-            await self.send_message(
-                room.room_id,
-                plain="Sorry, you can't create any more invites right now. Come back later.",
-            )
-            return
+        num = 1
+        try:
+            if len(command) > 1:
+                num = int(command[1])
+        except ValueError:
+            pass
+
+
 
-        token = await self.admin_api.create_token()
         async with self.db_lock:
-            await self.db.execute(
-                "INSERT INTO tokens (user, token) VALUES (?, ?);", (user, token)
+            if (
+                user not in env.USERS_WITHOUT_QUOTA
+                and await self.remaining_quota(user) - num < 0
+            ):
+                await self.send_message(
+                    room.room_id,
+                    plain="Sorry, you can't create any more invites right now. Come back later.",
+                )
+                return
+
+            tokens: list[str] = list(filter(lambda t: t, [await self.admin_api.create_token() for _ in range(num)]))  # type: ignore
+
+            await self.db.executemany(
+                "INSERT INTO tokens (user, token) VALUES (?, ?);",
+                ((user, token) for token in tokens),
             )
+
             await self.db.commit()
 
         await self.send_message(
             room.room_id,
-            formatted=f"<code>{token}</code>",
+            formatted="<code>{}</code>".format('\n'.join(tokens)),
         )
 
     async def send_hi_message(self, user_id: str, room: MatrixRoom):
@@ -158,15 +176,18 @@ class Bot:
             >= env.USER_REQUIRED_AGE
         )
 
-    async def quota_exceeded(self, user: str) -> bool:
+    async def remaining_quota(self, user: str) -> int:
         timespan = env.INVITE_CODE_QUOTA_TIMESPAN.total_seconds()
-        async with self.db_lock:
-            async with self.db.execute(
-                "SELECT count(token) AS amount FROM tokens WHERE unixepoch(CURRENT_TIMESTAMP)-unixepoch(created) < ? AND user = ?;",
-                (timespan, user),
-            ) as cursor:
-                res = await cursor.fetchone()
-                return res is not None and res[0] >= env.INVITE_CODE_QUOTA_AMOUNT
+
+        async with self.db.execute(
+            "SELECT count(token) AS amount FROM tokens WHERE unixepoch(CURRENT_TIMESTAMP)-unixepoch(created) < ? AND user = ?;",
+            (timespan, user),
+        ) as cursor:
+            res = await cursor.fetchone()
+            if res is None:
+                return 0
+            else:
+                return env.INVITE_CODE_QUOTA_AMOUNT - res[0]
 
     async def require_dm_partner(self, room: MatrixRoom) -> Optional[str]:
         """
@@ -246,7 +267,6 @@ class Bot:
                 # and then gracefully shutdown all connections
                 logger.info("Gracefully shutting down")
 
-
                 await self.db_lock.acquire()
 
                 try: