feat: add trusted groups functionality for duty information
- Introduced a new `trusted_groups` table to store groups authorized to receive duty information. - Implemented functions to add, remove, and check trusted groups in the database. - Enhanced command handlers to manage trusted groups, including `/trust_group` and `/untrust_group` commands for admin users. - Updated internationalization messages to support new commands and group status notifications. - Added unit tests for trusted groups repository functions to ensure correct behavior and data integrity.
This commit is contained in:
@@ -13,6 +13,7 @@ from telegram.ext import ChatMemberHandler, CommandHandler, ContextTypes
|
||||
|
||||
from duty_teller.db.session import session_scope
|
||||
from duty_teller.i18n import get_lang, t
|
||||
from duty_teller.handlers.common import is_admin_async
|
||||
from duty_teller.services.group_duty_pin_service import (
|
||||
get_duty_message_text,
|
||||
get_message_id,
|
||||
@@ -21,6 +22,9 @@ from duty_teller.services.group_duty_pin_service import (
|
||||
save_pin,
|
||||
delete_pin,
|
||||
get_all_pin_chat_ids,
|
||||
is_group_trusted,
|
||||
trust_group,
|
||||
untrust_group,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -62,6 +66,37 @@ def _sync_get_message_id(chat_id: int) -> int | None:
|
||||
return get_message_id(session, chat_id)
|
||||
|
||||
|
||||
def _sync_is_trusted(chat_id: int) -> bool:
|
||||
"""Check if the group is trusted (sync wrapper for handlers)."""
|
||||
with session_scope(config.DATABASE_URL) as session:
|
||||
return is_group_trusted(session, chat_id)
|
||||
|
||||
|
||||
def _sync_trust_group(chat_id: int, added_by_user_id: int | None) -> bool:
|
||||
"""Add group to trusted list. Returns True if already trusted (no-op)."""
|
||||
with session_scope(config.DATABASE_URL) as session:
|
||||
if is_group_trusted(session, chat_id):
|
||||
return True
|
||||
trust_group(session, chat_id, added_by_user_id)
|
||||
return False
|
||||
|
||||
|
||||
def _sync_untrust_group(chat_id: int) -> tuple[bool, int | None]:
|
||||
"""Remove group from trusted list.
|
||||
|
||||
Returns:
|
||||
(was_trusted, message_id): was_trusted False if group was not in list;
|
||||
message_id of pinned message if any (for cleanup), else None.
|
||||
"""
|
||||
with session_scope(config.DATABASE_URL) as session:
|
||||
if not is_group_trusted(session, chat_id):
|
||||
return (False, None)
|
||||
message_id = get_message_id(session, chat_id)
|
||||
delete_pin(session, chat_id)
|
||||
untrust_group(session, chat_id)
|
||||
return (True, message_id)
|
||||
|
||||
|
||||
async def _schedule_next_update(
|
||||
application, chat_id: int, when_utc: datetime | None
|
||||
) -> None:
|
||||
@@ -102,17 +137,40 @@ async def _schedule_next_update(
|
||||
|
||||
async def _refresh_pin_for_chat(
|
||||
context: ContextTypes.DEFAULT_TYPE, chat_id: int
|
||||
) -> Literal["updated", "no_message", "failed"]:
|
||||
) -> Literal["updated", "no_message", "failed", "untrusted"]:
|
||||
"""Refresh pinned duty message: send new message, unpin old, pin new, save new message_id, delete old.
|
||||
|
||||
Uses single DB session for message_id, text, next_shift_end (consolidated).
|
||||
If the group is no longer trusted, removes pin record, job, and message; returns "untrusted".
|
||||
|
||||
Returns:
|
||||
"updated" if the message was sent, pinned and saved successfully;
|
||||
"no_message" if there is no pin record for this chat;
|
||||
"failed" if send_message or permissions failed.
|
||||
"failed" if send_message or permissions failed;
|
||||
"untrusted" if the group was removed from trusted list (pin record and message cleaned up).
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
trusted = await loop.run_in_executor(None, _sync_is_trusted, chat_id)
|
||||
if not trusted:
|
||||
old_message_id = await loop.run_in_executor(None, _sync_get_message_id, chat_id)
|
||||
await loop.run_in_executor(None, _sync_delete_pin, chat_id)
|
||||
name = f"{JOB_NAME_PREFIX}{chat_id}"
|
||||
if context.application.job_queue:
|
||||
for job in context.application.job_queue.get_jobs_by_name(name):
|
||||
job.schedule_removal()
|
||||
if old_message_id is not None:
|
||||
try:
|
||||
await context.bot.unpin_chat_message(chat_id=chat_id)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
try:
|
||||
await context.bot.delete_message(
|
||||
chat_id=chat_id, message_id=old_message_id
|
||||
)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
logger.info("Chat_id=%s no longer trusted, removed pin record and job", chat_id)
|
||||
return "untrusted"
|
||||
message_id, text, next_end = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: _sync_get_pin_refresh_data(chat_id, config.DEFAULT_LANGUAGE),
|
||||
@@ -186,6 +244,17 @@ async def my_chat_member_handler(
|
||||
ChatMemberStatus.BANNED,
|
||||
):
|
||||
loop = asyncio.get_running_loop()
|
||||
trusted = await loop.run_in_executor(None, _sync_is_trusted, chat_id)
|
||||
if not trusted:
|
||||
lang = get_lang(update.effective_user)
|
||||
try:
|
||||
await context.bot.send_message(
|
||||
chat_id=chat_id,
|
||||
text=t(lang, "group.not_trusted"),
|
||||
)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
return
|
||||
lang = get_lang(update.effective_user)
|
||||
text = await loop.run_in_executor(
|
||||
None, lambda: _get_duty_message_text_sync(lang)
|
||||
@@ -255,6 +324,10 @@ async def pin_duty_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> No
|
||||
return
|
||||
chat_id = chat.id
|
||||
loop = asyncio.get_running_loop()
|
||||
trusted = await loop.run_in_executor(None, _sync_is_trusted, chat_id)
|
||||
if not trusted:
|
||||
await update.message.reply_text(t(lang, "group.not_trusted"))
|
||||
return
|
||||
message_id = await loop.run_in_executor(None, _sync_get_message_id, chat_id)
|
||||
if message_id is None:
|
||||
text = await loop.run_in_executor(
|
||||
@@ -312,13 +385,107 @@ async def refresh_pin_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
|
||||
await update.message.reply_text(t(lang, "refresh_pin.group_only"))
|
||||
return
|
||||
chat_id = chat.id
|
||||
loop = asyncio.get_running_loop()
|
||||
trusted = await loop.run_in_executor(None, _sync_is_trusted, chat_id)
|
||||
if not trusted:
|
||||
await update.message.reply_text(t(lang, "group.not_trusted"))
|
||||
return
|
||||
result = await _refresh_pin_for_chat(context, chat_id)
|
||||
await update.message.reply_text(t(lang, f"refresh_pin.{result}"))
|
||||
|
||||
|
||||
async def trust_group_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle /trust_group: add current group to trusted list (admin only)."""
|
||||
if not update.message or not update.effective_chat or not update.effective_user:
|
||||
return
|
||||
chat = update.effective_chat
|
||||
lang = get_lang(update.effective_user)
|
||||
if chat.type not in ("group", "supergroup"):
|
||||
await update.message.reply_text(t(lang, "trust_group.group_only"))
|
||||
return
|
||||
if not await is_admin_async(update.effective_user.id):
|
||||
await update.message.reply_text(t(lang, "import.admin_only"))
|
||||
return
|
||||
chat_id = chat.id
|
||||
loop = asyncio.get_running_loop()
|
||||
already_trusted = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: _sync_trust_group(chat_id, update.effective_user.id if update.effective_user else None),
|
||||
)
|
||||
if already_trusted:
|
||||
await update.message.reply_text(t(lang, "trust_group.already_trusted"))
|
||||
return
|
||||
await update.message.reply_text(t(lang, "trust_group.added"))
|
||||
message_id = await loop.run_in_executor(None, _sync_get_message_id, chat_id)
|
||||
if message_id is None:
|
||||
text = await loop.run_in_executor(
|
||||
None, lambda: _get_duty_message_text_sync(lang)
|
||||
)
|
||||
try:
|
||||
msg = await context.bot.send_message(chat_id=chat_id, text=text)
|
||||
except (BadRequest, Forbidden) as e:
|
||||
logger.warning(
|
||||
"Failed to send duty message after trust_group chat_id=%s: %s",
|
||||
chat_id,
|
||||
e,
|
||||
)
|
||||
return
|
||||
try:
|
||||
await context.bot.pin_chat_message(
|
||||
chat_id=chat_id,
|
||||
message_id=msg.message_id,
|
||||
disable_notification=not config.DUTY_PIN_NOTIFY,
|
||||
)
|
||||
except (BadRequest, Forbidden) as e:
|
||||
logger.warning(
|
||||
"Failed to pin message after trust_group chat_id=%s: %s", chat_id, e
|
||||
)
|
||||
await loop.run_in_executor(None, _sync_save_pin, chat_id, msg.message_id)
|
||||
next_end = await loop.run_in_executor(None, _get_next_shift_end_sync)
|
||||
await _schedule_next_update(context.application, chat_id, next_end)
|
||||
|
||||
|
||||
async def untrust_group_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle /untrust_group: remove current group from trusted list (admin only)."""
|
||||
if not update.message or not update.effective_chat or not update.effective_user:
|
||||
return
|
||||
chat = update.effective_chat
|
||||
lang = get_lang(update.effective_user)
|
||||
if chat.type not in ("group", "supergroup"):
|
||||
await update.message.reply_text(t(lang, "untrust_group.group_only"))
|
||||
return
|
||||
if not await is_admin_async(update.effective_user.id):
|
||||
await update.message.reply_text(t(lang, "import.admin_only"))
|
||||
return
|
||||
chat_id = chat.id
|
||||
loop = asyncio.get_running_loop()
|
||||
was_trusted, message_id = await loop.run_in_executor(
|
||||
None, _sync_untrust_group, chat_id
|
||||
)
|
||||
if not was_trusted:
|
||||
await update.message.reply_text(t(lang, "untrust_group.not_trusted"))
|
||||
return
|
||||
name = f"{JOB_NAME_PREFIX}{chat_id}"
|
||||
if context.application.job_queue:
|
||||
for job in context.application.job_queue.get_jobs_by_name(name):
|
||||
job.schedule_removal()
|
||||
if message_id is not None:
|
||||
try:
|
||||
await context.bot.unpin_chat_message(chat_id=chat_id)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
try:
|
||||
await context.bot.delete_message(chat_id=chat_id, message_id=message_id)
|
||||
except (BadRequest, Forbidden):
|
||||
pass
|
||||
await update.message.reply_text(t(lang, "untrust_group.removed"))
|
||||
|
||||
|
||||
group_duty_pin_handler = ChatMemberHandler(
|
||||
my_chat_member_handler,
|
||||
ChatMemberHandler.MY_CHAT_MEMBER,
|
||||
)
|
||||
pin_duty_handler = CommandHandler("pin_duty", pin_duty_cmd)
|
||||
refresh_pin_handler = CommandHandler("refresh_pin", refresh_pin_cmd)
|
||||
trust_group_handler = CommandHandler("trust_group", trust_group_cmd)
|
||||
untrust_group_handler = CommandHandler("untrust_group", untrust_group_cmd)
|
||||
|
||||
Reference in New Issue
Block a user