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"{token}
",
+ formatted="{}
".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: