mab update, !short command, last_access fix
- Updated mab to v0.5.0 - Added `!short` command to add named tokens - Fixed `last_access_at` field not being updated
This commit is contained in:
185
bot_callbacks.py
185
bot_callbacks.py
@@ -4,8 +4,7 @@ import html
|
||||
import time
|
||||
import traceback
|
||||
|
||||
from mab import MatrixBot, TextCommandFilter
|
||||
from nio import MatrixRoom, RoomMessageText
|
||||
from mab import * # type: ignore
|
||||
|
||||
import util
|
||||
from database import Database
|
||||
@@ -14,7 +13,6 @@ from datatypes import *
|
||||
#
|
||||
# PRIVATE
|
||||
#
|
||||
_bot: MatrixBot = None # type: ignore
|
||||
_db: Database = None # type: ignore
|
||||
|
||||
_COMMANDS = {
|
||||
@@ -27,6 +25,7 @@ _COMMANDS = {
|
||||
"ban": "Забанить IP адрес",
|
||||
"unban": "Разбанить IP адрес",
|
||||
"bans": "Получить список забаненных IP адресов",
|
||||
"short": "(Небезопасно) Создать токен-имя",
|
||||
}
|
||||
|
||||
def _generate_help_message() -> str:
|
||||
@@ -46,38 +45,35 @@ def _generate_help_message() -> str:
|
||||
#
|
||||
# COMMAND HANDLERS
|
||||
#
|
||||
async def _on_cmd_help(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_help(ctx: EventContext) -> None:
|
||||
"""!help"""
|
||||
await _bot.send_text_to_room(room.room_id, _generate_help_message(), is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, _generate_help_message())
|
||||
|
||||
async def _on_cmd_info(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_info(ctx: EventContext) -> None:
|
||||
"""!info"""
|
||||
# get room info (or add it)
|
||||
room_info = await _db.room_get(room.room_id)
|
||||
room_info = await _db.room_get(ctx.room.room_id)
|
||||
if room_info is None:
|
||||
room_info = await _db.room_create(room.room_id)
|
||||
room_info = await _db.room_create(ctx.room.room_id)
|
||||
# failure
|
||||
if room_info is None:
|
||||
await _bot.send_text_to_room(
|
||||
room.room_id,
|
||||
"<strong>Нет информации о комнате</strong>",
|
||||
is_html=True
|
||||
await ctx.bot.send_text(
|
||||
ctx.room, "<strong>Нет информации о комнате</strong>"
|
||||
)
|
||||
return
|
||||
# respond
|
||||
response = "<strong><i>Сведения о комнате</i></strong><br>"
|
||||
response += f"<strong>Канал:</strong> <code>{html.escape(room_info.code)}</code>"
|
||||
await _bot.send_text_to_room(room.room_id, response, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
async def _on_cmd_tokens(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_tokens(ctx: EventContext) -> None:
|
||||
"""!info"""
|
||||
# get all tokens and check if there are none
|
||||
tokens = await _db.token_get_all()
|
||||
if not tokens:
|
||||
await _bot.send_text_to_room(
|
||||
room.room_id,
|
||||
"<strong>Нет токенов, используйте <code>!auth</code></strong>",
|
||||
is_html=True
|
||||
await ctx.bot.send_text(
|
||||
ctx.room,
|
||||
"<strong>Нет токенов, используйте <code>!auth</code></strong>"
|
||||
)
|
||||
return
|
||||
# create the list of tokens
|
||||
@@ -92,15 +88,15 @@ async def _on_cmd_tokens(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
response += f"<li><strong>Последнее использование:</strong> <i>{util.date_to_text(token.last_access_at)}</i></li>"
|
||||
response += "</ul>"
|
||||
# respond
|
||||
await _bot.send_text_to_room(room.room_id, response, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
async def _on_cmd_auth(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_auth(ctx: EventContext) -> None:
|
||||
"""!auth"""
|
||||
# create new token
|
||||
token = await _db.token_create()
|
||||
# set the name if it is provided
|
||||
if event.command_args:
|
||||
token.name = " ".join(event.command_args)
|
||||
if ctx[CTX_CMD_ARGS]:
|
||||
token.name = " ".join(ctx[CTX_CMD_ARGS])
|
||||
await _db.token_set_name(token.code, token.name)
|
||||
# create the response
|
||||
response = "<strong>Создан новый токен</strong>"
|
||||
@@ -111,39 +107,39 @@ async def _on_cmd_auth(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
else:
|
||||
response += "указывается в запросе"
|
||||
# respond
|
||||
await _bot.send_text_to_room(room.room_id, response, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
async def _on_cmd_deauth(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_deauth(ctx: EventContext) -> None:
|
||||
"""!deauth"""
|
||||
# no token provided
|
||||
if len(event.command_args) != 1:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Укажите токен, который надо удалить (должен быть ровно один аргумент)</strong>", is_html=True)
|
||||
if len(ctx[CTX_CMD_ARGS]) != 1:
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Укажите токен, который надо удалить (должен быть ровно один аргумент)</strong>")
|
||||
return
|
||||
token = event.command_args[0]
|
||||
token = ctx[CTX_CMD_ARGS][0]
|
||||
# check if token does not exist
|
||||
if await _db.token_get(token) is None:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Токен не найден</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Токен не найден</strong>")
|
||||
return
|
||||
# remove the token
|
||||
await _db.token_delete(token)
|
||||
# respond
|
||||
response = f"<strong>Удалён токен <code>{token}</code></strong>"
|
||||
await _bot.send_text_to_room(room.room_id, response, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
async def _on_cmd_name(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_name(ctx: EventContext) -> None:
|
||||
"""!name"""
|
||||
# check arguments
|
||||
if len(event.command_args) < 1:
|
||||
if len(ctx[CTX_CMD_ARGS]) < 1:
|
||||
error = "<strong>Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя.</strong>"
|
||||
await _bot.send_text_to_room(room.room_id, error, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, error)
|
||||
return
|
||||
token = event.command_args[0]
|
||||
new_name = " ".join(event.command_args[1:])
|
||||
token = ctx[CTX_CMD_ARGS][0]
|
||||
new_name = " ".join(ctx[CTX_CMD_ARGS][1:])
|
||||
if not new_name.strip():
|
||||
new_name = None
|
||||
# check if token exists
|
||||
if await _db.token_get(token) is None:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Токен не существует</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Токен не существует</strong>")
|
||||
return
|
||||
# set new name
|
||||
await _db.token_set_name(token, new_name)
|
||||
@@ -153,51 +149,51 @@ async def _on_cmd_name(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
else:
|
||||
response = f"<strong>Удалено имя для токена <code>{token}</code></strong>"
|
||||
# respond
|
||||
await _bot.send_text_to_room(room.room_id, response, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
async def _on_cmd_ban(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_ban(ctx: EventContext) -> None:
|
||||
"""!ban"""
|
||||
# check arguments
|
||||
if len(event.command_args) < 2:
|
||||
if len(ctx[CTX_CMD_ARGS]) < 2:
|
||||
error = "<strong>Формат: <code>!ban <IP> <SECONDS> [REASON]</code></strong>"
|
||||
await _bot.send_text_to_room(room.room_id, error, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, error)
|
||||
return
|
||||
# get event.command_args
|
||||
try:
|
||||
ip = event.command_args[0]
|
||||
duration = float(event.command_args[1])
|
||||
reason = " ".join(event.command_args[2:]) if event.command_args[2:] else "Manual ban"
|
||||
ip = ctx[CTX_CMD_ARGS][0]
|
||||
duration = float(ctx[CTX_CMD_ARGS][1])
|
||||
reason = " ".join(ctx[CTX_CMD_ARGS][2:]) if ctx[CTX_CMD_ARGS][2:] else "Manual ban"
|
||||
except:
|
||||
error = "<strong>Возникла ошибка. Наверняка неправильно указаны секунды.</strong>"
|
||||
await _bot.send_text_to_room(room.room_id, error, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, error)
|
||||
return
|
||||
# ban
|
||||
try:
|
||||
await _db.ban_create(ip, time.time() + duration, reason)
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>IP адрес заблокирован</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>IP адрес заблокирован</strong>")
|
||||
except:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Возникла ошибка</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Возникла ошибка</strong>")
|
||||
|
||||
async def _on_cmd_unban(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_unban(ctx: EventContext) -> None:
|
||||
"""!unban"""
|
||||
# check arguments
|
||||
if len(event.command_args) != 1:
|
||||
if len(ctx[CTX_CMD_ARGS]) != 1:
|
||||
error = "<strong>Формат: <code>!ban <IP></code></strong>"
|
||||
await _bot.send_text_to_room(room.room_id, error, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, error)
|
||||
return
|
||||
# unban
|
||||
try:
|
||||
await _db.ban_delete(event.command_args[0])
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>IP адрес разблокирован (если он был заблокирован)</strong>", is_html=True)
|
||||
await _db.ban_delete(ctx[CTX_CMD_ARGS][0])
|
||||
await ctx.bot.send_text(ctx.room, "<strong>IP адрес разблокирован (если он был заблокирован)</strong>")
|
||||
except:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Возникла ошибка</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Возникла ошибка</strong>")
|
||||
|
||||
async def _on_cmd_bans(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def _on_cmd_bans(ctx: EventContext) -> None:
|
||||
"""!bans"""
|
||||
# list
|
||||
bans = _db.ban_get_all()
|
||||
if not bans:
|
||||
await _bot.send_text_to_room(room.room_id, "<strong>Нет заблокированных IP адресов</strong>", is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, "<strong>Нет заблокированных IP адресов</strong>")
|
||||
return
|
||||
# create the response
|
||||
result = "<strong>Список заблокированных IP</strong><br><ul>"
|
||||
@@ -207,7 +203,27 @@ async def _on_cmd_bans(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
result += f"<li><strong>Причина:</strong> <code>{html.escape(ban["reason"])}</code></li>"
|
||||
result += "</ul></li>"
|
||||
result += "</ul>"
|
||||
await _bot.send_text_to_room(room.room_id, result, is_html=True)
|
||||
await ctx.bot.send_text(ctx.room, result)
|
||||
|
||||
async def _on_cmd_short(ctx: EventContext) -> None:
|
||||
"""!auth"""
|
||||
# get token to use
|
||||
if len(ctx[CTX_CMD_ARGS]) != 1:
|
||||
error = "<strong>Формат: <code>!short <token></code></strong>"
|
||||
await ctx.bot.send_text(ctx.room, error)
|
||||
return
|
||||
# create new token
|
||||
try:
|
||||
token = await _db.token_create(ctx[CTX_CMD_ARGS][0])
|
||||
except Exception as e:
|
||||
await ctx.bot.send_text(ctx.room, f"Не удалось создать токен: {" ".join(e.args)}")
|
||||
return
|
||||
# create the response
|
||||
response = "<strong>Создан новый токен</strong>"
|
||||
response += f"<br><strong>Код:</strong> <code>{token.code}</code>"
|
||||
response += f"<br><strong>Имя сервиса:</strong> указывается в запросе"
|
||||
# respond
|
||||
await ctx.bot.send_text(ctx.room, response)
|
||||
|
||||
|
||||
#
|
||||
@@ -215,24 +231,45 @@ async def _on_cmd_bans(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
#
|
||||
def setup(bot: MatrixBot, db: Database) -> None:
|
||||
"""Setup the callbacks"""
|
||||
global _bot, _db
|
||||
global _db
|
||||
_db = db
|
||||
_bot = bot
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_help, TextCommandFilter(["help", "h", "?"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_info, TextCommandFilter(["info", "room", "i", "r"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_tokens, TextCommandFilter(["tokens", "t"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_auth, TextCommandFilter(["auth", "create", "a", "c"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_deauth, TextCommandFilter(["deauth", "delete", "d"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_name, TextCommandFilter(["name", "n"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_ban, TextCommandFilter(["ban", "b"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_unban, TextCommandFilter(["unban", "u"]))
|
||||
_bot.add_event_callback(
|
||||
_on_cmd_bans, TextCommandFilter(["bans", "l"]))
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["info", "room", "i", "r"]),
|
||||
_on_cmd_info
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["tokens", "t"]),
|
||||
_on_cmd_tokens
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["auth", "create", "a", "c"]),
|
||||
_on_cmd_auth
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["deauth", "delete", "d"]),
|
||||
_on_cmd_deauth
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["name", "n"]),
|
||||
_on_cmd_name
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["ban", "b"]),
|
||||
_on_cmd_ban
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["unban", "u"]),
|
||||
_on_cmd_unban
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["bans", "l"]),
|
||||
_on_cmd_bans
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & BodyCommandFilter(["short", "s"]),
|
||||
_on_cmd_short
|
||||
)
|
||||
bot.add_callback(
|
||||
~SenderIsBotFilter() & NewMessageFilter(),
|
||||
_on_cmd_help
|
||||
)
|
||||
21
database.py
21
database.py
@@ -252,11 +252,17 @@ class Database:
|
||||
return None
|
||||
|
||||
|
||||
async def token_create(self) -> ObjectToken:
|
||||
async def token_create(self, code: str | None = None) -> ObjectToken:
|
||||
"""Create a token."""
|
||||
if self._connection is None:
|
||||
raise RuntimeError("Not connected to the database")
|
||||
code = get_hash(str(time.time()).encode() + os.urandom(64))
|
||||
if code is None:
|
||||
code = get_hash(str(time.time()).encode() + os.urandom(64))
|
||||
else:
|
||||
if len(code) < 2 or len(code) > 16:
|
||||
raise RuntimeError("Token must be 2-16 symbols long")
|
||||
if not all(ord(c) < 128 and (c.islower() or c.isdigit()) for c in code):
|
||||
raise RuntimeError("Only digits and ASCII lowercase letters are allowed")
|
||||
create_time = time.time()
|
||||
statement = """
|
||||
INSERT INTO tokens (code, name, created_at, last_access_at)
|
||||
@@ -329,6 +335,17 @@ class Database:
|
||||
except:
|
||||
self._logger.error(traceback.format_exc())
|
||||
|
||||
async def token_set_last_access(self, code: str, timestamp: float) -> None:
|
||||
"""Set last_access_at for the token."""
|
||||
if self._connection is None:
|
||||
raise RuntimeError("Not connected to the database")
|
||||
try:
|
||||
statement = "UPDATE tokens SET last_access_at=? WHERE code=?"
|
||||
await self._connection.execute(statement, (timestamp, code))
|
||||
await self._connection.commit()
|
||||
except:
|
||||
self._logger.error(traceback.format_exc())
|
||||
|
||||
|
||||
async def ban_create(self, ip: str, expires_at: float, reason: str) -> None:
|
||||
"""Save information about banned IP address. Replaces existing IPs."""
|
||||
|
||||
2
main.py
2
main.py
@@ -55,7 +55,7 @@ async def main() -> None:
|
||||
if not await db.connect():
|
||||
util.log_error("Can't connect to the database!")
|
||||
return
|
||||
bot.start()
|
||||
await bot.start()
|
||||
await web_server.start()
|
||||
|
||||
# wait for Ctrl+C
|
||||
|
||||
@@ -34,7 +34,7 @@ idna==3.19
|
||||
Jinja2==3.1.6
|
||||
jsonschema==4.26.0
|
||||
jsonschema-specifications==2025.9.1
|
||||
mab @ git+https://git.tyukalov.su/nikita/mab@v0.1.0
|
||||
mab @ git+https://git.tyukalov.su/nikita/mab@v0.5.0
|
||||
markdown-it-py==4.2.0
|
||||
MarkupSafe==3.0.3
|
||||
matrix-nio==0.26.0
|
||||
@@ -42,6 +42,7 @@ mdurl==0.1.2
|
||||
multidict==6.7.1
|
||||
packaging==26.3
|
||||
peewee==3.19.0
|
||||
pillow==12.3.0
|
||||
propcache==0.5.2
|
||||
pycryptodome==3.23.0
|
||||
pydantic==2.13.5
|
||||
@@ -51,6 +52,7 @@ pydantic_core==2.46.5
|
||||
Pygments==2.21.0
|
||||
pyproject_hooks==1.2.0
|
||||
python-dotenv==1.2.3
|
||||
python-magic==0.4.27
|
||||
python-multipart==0.0.32
|
||||
python-socks==3.0.0
|
||||
PyYAML==6.0.3
|
||||
|
||||
@@ -84,6 +84,7 @@ class AuthMiddleware(BaseHTTPMiddleware):
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
content={"detail": "No such token"}
|
||||
)
|
||||
await database.token_set_last_access(token, time.time())
|
||||
request.state.token = token_data
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@@ -58,7 +58,7 @@ async def _get_notify(request: Request, params: ModelGetNotify = Depends()):
|
||||
text += html.escape(params.text)
|
||||
# try to send the notification
|
||||
try:
|
||||
notification_id = await matrix.send_text_to_room(
|
||||
notification_id = await matrix.send_text(
|
||||
room.matrix_id,
|
||||
text,
|
||||
is_html=True
|
||||
|
||||
Reference in New Issue
Block a user