diff --git a/database/async_db.py b/database/async_db.py index bd2e62b..859a49f 100644 --- a/database/async_db.py +++ b/database/async_db.py @@ -188,13 +188,20 @@ class AsyncBotDB: ) # Методы для работы с черным списком - async def set_user_blacklist(self, user_id: int, user_name: str = None, - message_for_user: str = None, date_to_unban: int = None): + async def set_user_blacklist( + self, + user_id: int, + user_name: str = None, + message_for_user: str = None, + date_to_unban: int = None, + ban_author: Optional[int] = None, + ): """Добавляет пользователя в черный список.""" blacklist_user = BlacklistUser( user_id=user_id, message_for_user=message_for_user, - date_to_unban=date_to_unban + date_to_unban=date_to_unban, + ban_author=ban_author, ) await self.factory.blacklist.add_user(blacklist_user) diff --git a/database/models.py b/database/models.py index 67b11ed..12fe524 100644 --- a/database/models.py +++ b/database/models.py @@ -26,6 +26,7 @@ class BlacklistUser: message_for_user: Optional[str] = None date_to_unban: Optional[int] = None created_at: Optional[int] = None + ban_author: Optional[int] = None @dataclass diff --git a/database/repositories/blacklist_repository.py b/database/repositories/blacklist_repository.py index d02d6af..d66d514 100644 --- a/database/repositories/blacklist_repository.py +++ b/database/repositories/blacklist_repository.py @@ -14,7 +14,9 @@ class BlacklistRepository(DatabaseConnection): message_for_user TEXT, date_to_unban INTEGER, created_at INTEGER DEFAULT (strftime('%s', 'now')), - FOREIGN KEY (user_id) REFERENCES our_users (user_id) ON DELETE CASCADE + ban_author INTEGER, + FOREIGN KEY (user_id) REFERENCES our_users (user_id) ON DELETE CASCADE, + FOREIGN KEY (ban_author) REFERENCES our_users (user_id) ON DELETE SET NULL ) ''' await self._execute_query(query) @@ -23,10 +25,15 @@ class BlacklistRepository(DatabaseConnection): async def add_user(self, blacklist_user: BlacklistUser) -> None: """Добавляет пользователя в черный список.""" query = """ - INSERT INTO blacklist (user_id, message_for_user, date_to_unban) - VALUES (?, ?, ?) + INSERT INTO blacklist (user_id, message_for_user, date_to_unban, ban_author) + VALUES (?, ?, ?, ?) """ - params = (blacklist_user.user_id, blacklist_user.message_for_user, blacklist_user.date_to_unban) + params = ( + blacklist_user.user_id, + blacklist_user.message_for_user, + blacklist_user.date_to_unban, + blacklist_user.ban_author, + ) await self._execute_query(query, params) self.logger.info(f"Пользователь добавлен в черный список: user_id={blacklist_user.user_id}") @@ -52,7 +59,11 @@ class BlacklistRepository(DatabaseConnection): async def get_user(self, user_id: int) -> Optional[BlacklistUser]: """Возвращает информацию о пользователе в черном списке по user_id.""" - query = "SELECT user_id, message_for_user, date_to_unban, created_at FROM blacklist WHERE user_id = ?" + query = """ + SELECT user_id, message_for_user, date_to_unban, created_at, ban_author + FROM blacklist + WHERE user_id = ? + """ rows = await self._execute_query_with_result(query, (user_id,)) row = rows[0] if rows else None @@ -61,40 +72,54 @@ class BlacklistRepository(DatabaseConnection): user_id=row[0], message_for_user=row[1], date_to_unban=row[2], - created_at=row[3] + created_at=row[3], + ban_author=row[4] if len(row) > 4 else None, ) return None async def get_all_users(self, offset: int = 0, limit: int = 10) -> List[BlacklistUser]: """Возвращает список пользователей в черном списке.""" - query = "SELECT user_id, message_for_user, date_to_unban, created_at FROM blacklist LIMIT ?, ?" + query = """ + SELECT user_id, message_for_user, date_to_unban, created_at, ban_author + FROM blacklist + LIMIT ?, ? + """ rows = await self._execute_query_with_result(query, (offset, limit)) users = [] for row in rows: - users.append(BlacklistUser( - user_id=row[0], - message_for_user=row[1], - date_to_unban=row[2], - created_at=row[3] - )) + users.append( + BlacklistUser( + user_id=row[0], + message_for_user=row[1], + date_to_unban=row[2], + created_at=row[3], + ban_author=row[4] if len(row) > 4 else None, + ) + ) self.logger.info(f"Получен список пользователей в черном списке (offset={offset}, limit={limit}): {len(users)}") return users async def get_all_users_no_limit(self) -> List[BlacklistUser]: """Возвращает список всех пользователей в черном списке без лимитов.""" - query = "SELECT user_id, message_for_user, date_to_unban, created_at FROM blacklist" + query = """ + SELECT user_id, message_for_user, date_to_unban, created_at, ban_author + FROM blacklist + """ rows = await self._execute_query_with_result(query) users = [] for row in rows: - users.append(BlacklistUser( - user_id=row[0], - message_for_user=row[1], - date_to_unban=row[2], - created_at=row[3] - )) + users.append( + BlacklistUser( + user_id=row[0], + message_for_user=row[1], + date_to_unban=row[2], + created_at=row[3], + ban_author=row[4] if len(row) > 4 else None, + ) + ) self.logger.info(f"Получен список всех пользователей в черном списке: {len(users)}") return users diff --git a/helper_bot/handlers/admin/admin_handlers.py b/helper_bot/handlers/admin/admin_handlers.py index b9b229a..9d8d8f4 100644 --- a/helper_bot/handlers/admin/admin_handlers.py +++ b/helper_bot/handlers/admin/admin_handlers.py @@ -359,7 +359,8 @@ async def confirm_ban( user_id=user_data['target_user_id'], username=user_data['target_username'], reason=user_data['ban_reason'], - ban_days=user_data['ban_days'] + ban_days=user_data['ban_days'], + ban_author_id=message.from_user.id, ) safe_username = escape_html(user_data['target_username']) diff --git a/helper_bot/handlers/admin/services.py b/helper_bot/handlers/admin/services.py index 55fe57b..df3f3eb 100644 --- a/helper_bot/handlers/admin/services.py +++ b/helper_bot/handlers/admin/services.py @@ -117,7 +117,7 @@ class AdminService: @track_time("ban_user", "admin_service") @track_errors("admin_service", "ban_user") - async def ban_user(self, user_id: int, username: str, reason: str, ban_days: Optional[int]) -> None: + async def ban_user(self, user_id: int, username: str, reason: str, ban_days: Optional[int], ban_author_id: int) -> None: """Заблокировать пользователя""" try: # Проверяем, не заблокирован ли уже пользователь @@ -130,7 +130,7 @@ class AdminService: date_to_unban = add_days_to_date(ban_days) # Сохраняем в БД (username больше не передается, так как не используется в новой схеме) - await self.bot_db.set_user_blacklist(user_id, None, reason, date_to_unban) + await self.bot_db.set_user_blacklist(user_id, None, reason, date_to_unban, ban_author=ban_author_id) logger.info(f"Пользователь {user_id} ({username}) заблокирован. Причина: {reason}, срок: {ban_days} дней") diff --git a/helper_bot/handlers/callback/services.py b/helper_bot/handlers/callback/services.py index 991b9f5..c8ce08f 100644 --- a/helper_bot/handlers/callback/services.py +++ b/helper_bot/handlers/callback/services.py @@ -439,11 +439,14 @@ class BanService: current_date = datetime.now() date_to_unban = int((current_date + timedelta(days=7)).timestamp()) + ban_author_id = call.from_user.id + await self.db.set_user_blacklist( user_id=author_id, user_name=None, message_for_user="Спам", - date_to_unban=date_to_unban + date_to_unban=date_to_unban, + ban_author=ban_author_id, ) await self._get_bot(call.message).delete_message(chat_id=self.group_for_posts, message_id=call.message.message_id) diff --git a/helper_bot/utils/helper_func.py b/helper_bot/utils/helper_func.py index 43e930d..8e3dad7 100644 --- a/helper_bot/utils/helper_func.py +++ b/helper_bot/utils/helper_func.py @@ -143,17 +143,17 @@ def get_text_message(post_text: str, first_name: str, username: str = None, is_a # TODO: Уверен можно укоротить if is_anonymous is not None: if is_anonymous: - return f'Пост из ТГ:\n{safe_post_text}\n\nПост опубликован анонимно' + return f'{safe_post_text}\n\nПост опубликован анонимно' else: - return f'Пост из ТГ:\n{safe_post_text}\n\nАвтор поста: {author_info}' + return f'{safe_post_text}\n\nАвтор поста: {author_info}' else: # Legacy: определяем по тексту if "неанон" in post_text or "не анон" in post_text: - return f'Пост из ТГ:\n{safe_post_text}\n\nАвтор поста: {author_info}' + return f'{safe_post_text}\n\nАвтор поста: {author_info}' elif "анон" in post_text: - return f'Пост из ТГ:\n{safe_post_text}\n\nПост опубликован анонимно' + return f'{safe_post_text}\n\nПост опубликован анонимно' else: - return f'Пост из ТГ:\n{safe_post_text}\n\nАвтор поста: {author_info}' + return f'{safe_post_text}\n\nАвтор поста: {author_info}' @track_time("download_file", "helper_func") @track_errors("helper_func", "download_file") diff --git a/scripts/add_ban_author_column_to_blacklist.py b/scripts/add_ban_author_column_to_blacklist.py new file mode 100644 index 0000000..3513a77 --- /dev/null +++ b/scripts/add_ban_author_column_to_blacklist.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +""" +Скрипт миграции для добавления колонки ban_author в таблицу blacklist. +Колонка хранит user_id администратора, инициировавшего бан. +""" +import argparse +import asyncio +import os +import sys +from pathlib import Path + +import aiosqlite + +project_root = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(project_root)) + +from logs.custom_logger import logger # noqa: E402 + +DEFAULT_DB_PATH = "database/tg-bot-database.db" + + +def _column_exists(rows: list, name: str) -> bool: + """PRAGMA table_info returns (cid, name, type, notnull, dflt_value, pk).""" + for row in rows: + if row[1] == name: + return True + return False + + +async def main(db_path: str) -> None: + db_path = os.path.abspath(db_path) + if not os.path.exists(db_path): + logger.error("База данных не найдена: %s", db_path) + print(f"Ошибка: база данных не найдена: {db_path}") + return + + async with aiosqlite.connect(db_path) as conn: + await conn.execute("PRAGMA foreign_keys = ON") + + # Проверяем наличие колонки ban_author + cursor = await conn.execute("PRAGMA table_info(blacklist)") + rows = await cursor.fetchall() + await cursor.close() + + if not _column_exists(rows, "ban_author"): + logger.info("Добавление колонки ban_author в blacklist") + await conn.execute( + "ALTER TABLE blacklist " + "ADD COLUMN ban_author INTEGER REFERENCES our_users (user_id) ON DELETE SET NULL" + ) + await conn.commit() + print("Колонка ban_author добавлена в таблицу blacklist.") + else: + print("Колонка ban_author уже существует в таблице blacklist.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Добавление колонки ban_author в blacklist" + ) + parser.add_argument( + "--db", + default=os.environ.get("DB_PATH", DEFAULT_DB_PATH), + help="Путь к БД (или DB_PATH)", + ) + args = parser.parse_args() + asyncio.run(main(args.db)) + diff --git a/tests/test_blacklist_repository.py b/tests/test_blacklist_repository.py index 1d9c7b9..f4de1b7 100644 --- a/tests/test_blacklist_repository.py +++ b/tests/test_blacklist_repository.py @@ -37,7 +37,8 @@ class TestBlacklistRepository: user_id=12345, message_for_user="Нарушение правил", date_to_unban=int(time.time()) + 86400, # +1 день - created_at=int(time.time()) + created_at=int(time.time()), + ban_author=999, ) @pytest.fixture @@ -47,7 +48,8 @@ class TestBlacklistRepository: user_id=67890, message_for_user="Постоянный бан", date_to_unban=None, - created_at=int(time.time()) + created_at=int(time.time()), + ban_author=None, ) @pytest.mark.asyncio @@ -82,11 +84,11 @@ class TestBlacklistRepository: # Проверяем SQL запрос (учитываем форматирование) sql_query = call_args[0][0].replace('\n', ' ').replace(' ', ' ').replace(' ', ' ').strip() - expected_sql = "INSERT INTO blacklist (user_id, message_for_user, date_to_unban) VALUES (?, ?, ?)" + expected_sql = "INSERT INTO blacklist (user_id, message_for_user, date_to_unban, ban_author) VALUES (?, ?, ?, ?)" assert sql_query == expected_sql # Проверяем параметры - assert call_args[0][1] == (12345, "Нарушение правил", sample_blacklist_user.date_to_unban) + assert call_args[0][1] == (12345, "Нарушение правил", sample_blacklist_user.date_to_unban, 999) # Проверяем логирование blacklist_repository.logger.info.assert_called_once_with( @@ -99,7 +101,7 @@ class TestBlacklistRepository: await blacklist_repository.add_user(sample_blacklist_user_permanent) call_args = blacklist_repository._execute_query.call_args - assert call_args[0][1] == (67890, "Постоянный бан", None) + assert call_args[0][1] == (67890, "Постоянный бан", None, None) blacklist_repository.logger.info.assert_called_once_with( "Пользователь добавлен в черный список: user_id=67890" @@ -182,7 +184,7 @@ class TestBlacklistRepository: async def test_get_user_success(self, blacklist_repository): """Тест успешного получения пользователя по ID""" # Симулируем результат запроса - mock_row = (12345, "Нарушение правил", int(time.time()) + 86400, int(time.time())) + mock_row = (12345, "Нарушение правил", int(time.time()) + 86400, int(time.time()), 111) blacklist_repository._execute_query_with_result.return_value = [mock_row] result = await blacklist_repository.get_user(12345) @@ -193,12 +195,13 @@ class TestBlacklistRepository: assert result.message_for_user == "Нарушение правил" assert result.date_to_unban == mock_row[2] assert result.created_at == mock_row[3] + assert result.ban_author == mock_row[4] # Проверяем, что метод вызван с правильными параметрами blacklist_repository._execute_query_with_result.assert_called_once() call_args = blacklist_repository._execute_query_with_result.call_args - assert call_args[0][0] == "SELECT user_id, message_for_user, date_to_unban, created_at FROM blacklist WHERE user_id = ?" + assert "SELECT user_id, message_for_user, date_to_unban, created_at, ban_author" in call_args[0][0] assert call_args[0][1] == (12345,) @pytest.mark.asyncio diff --git a/tests/test_refactored_admin_handlers.py b/tests/test_refactored_admin_handlers.py index 8319e3e..10d0333 100644 --- a/tests/test_refactored_admin_handlers.py +++ b/tests/test_refactored_admin_handlers.py @@ -153,7 +153,7 @@ class TestAdminService: self.mock_db.set_user_blacklist = AsyncMock(return_value=None) # Act - await self.admin_service.ban_user(user_id, username, reason, ban_days) + await self.admin_service.ban_user(user_id, username, reason, ban_days, ban_author_id=999) # Assert self.mock_db.check_user_in_blacklist.assert_called_once_with(user_id) @@ -187,10 +187,10 @@ class TestAdminService: self.mock_db.set_user_blacklist = AsyncMock(return_value=None) # Act - await self.admin_service.ban_user(user_id, username, reason, ban_days) + await self.admin_service.ban_user(user_id, username, reason, ban_days, ban_author_id=999) # Assert - self.mock_db.set_user_blacklist.assert_called_once_with(user_id, None, reason, None) + self.mock_db.set_user_blacklist.assert_called_once_with(user_id, None, reason, None, ban_author=999) @pytest.mark.asyncio async def test_unban_user_success(self):