diff --git a/database.py b/database.py index d5ce2b6..36a767c 100644 --- a/database.py +++ b/database.py @@ -2,37 +2,12 @@ import sqlite3 from typing import List -class Database: +class BasicDatabase: def __init__(self): self.conn = sqlite3.connect('bot.db') self.conn.row_factory = sqlite3.Row self.cursor = self.conn.cursor() - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS chats ( - "id" INTEGER, - "active" INTEGER NOT NULL DEFAULT 0, - "rules" TEXT, - "greeting_join" TEXT, - "greeting_rejoin" TEXT, - "birthday_message" TEXT, - PRIMARY KEY("id")) - """) - - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS users ( - "chat_id" INTEGER, - "user_id" INTEGER, - "last_message" INTEGER NOT NULL DEFAULT 0, - "messages_today" INTEGER NOT NULL DEFAULT 0, - "messages_month" INTEGER NOT NULL DEFAULT 0, - "warnings" INTEGER NOT NULL DEFAULT 0, - "happy_birthday" INTEGER NOT NULL DEFAULT 1, - PRIMARY KEY("chat_id","user_id")) - """) - - self.conn.commit() - def get_chats(self): self.cursor.execute("SELECT * FROM chats") return self.cursor.fetchall() @@ -76,9 +51,6 @@ class Database: def user_increment_warnings(self, chat_id: int, user_id: int): self.user_increment(chat_id, user_id, ['warnings']) - def user_toggle_happy_birthday(self, chat_id: int, user_id: int, happy_birthday: int): - self.user_update(chat_id, user_id, happy_birthday=happy_birthday) - def user_increment(self, chat_id: int, user_id: int, fields: List[str]): self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) + " WHERE chat_id = ? AND user_id = ?", (chat_id, user_id)) @@ -143,21 +115,16 @@ class Database: self.cursor.execute("UPDATE users SET messages_month = 0") self.conn.commit() + def create_chat_if_not_exists(self, chat_id: int): + chat = self.get_chat(chat_id) + if chat is None: + self.add_chat(chat_id) + chat = self.get_chat(chat_id) + return chat -DB = Database() - - -def create_chat_if_not_exists(chat_id: int): - chat = DB.get_chat(chat_id) - if chat is None: - DB.add_chat(chat_id) - chat = DB.get_chat(chat_id) - return chat - - -def create_user_if_not_exists(chat_id: int, user_id: int): - user = DB.get_user(chat_id, user_id) - if user is None: - DB.add_user(chat_id, user_id) - user = DB.get_user(chat_id, user_id) - return user + def create_user_if_not_exists(self, chat_id: int, user_id: int): + user = self.get_user(chat_id, user_id) + if user is None: + self.add_user(chat_id, user_id) + user = self.get_user(chat_id, user_id) + return user diff --git a/handlers/__init__.py b/handlers/__init__.py deleted file mode 100644 index fab1bbb..0000000 --- a/handlers/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -import handlers.user -import handlers.admin -import handlers.action -import handlers.default - - -__all__ = [] diff --git a/labeler.py b/labeler.py deleted file mode 100644 index c1b145a..0000000 --- a/labeler.py +++ /dev/null @@ -1,4 +0,0 @@ -from vkbottle.bot import BotLabeler - - -labeler = BotLabeler() diff --git a/tg/handlers/__init__.py b/tg/handlers/__init__.py new file mode 100644 index 0000000..07379d2 --- /dev/null +++ b/tg/handlers/__init__.py @@ -0,0 +1,9 @@ +from aiogram import Router + +from . import user, admin, action, default + +router = Router() +router.include_router(user.router) +router.include_router(admin.router) +router.include_router(action.router) +router.include_router(default.router) diff --git a/tg/handlers/action.py b/tg/handlers/action.py new file mode 100644 index 0000000..c747bb0 --- /dev/null +++ b/tg/handlers/action.py @@ -0,0 +1,40 @@ +from aiogram import Router, F +from aiogram.enums import ContentType, ParseMode +from aiogram.types import Message + +from messages import * +import tg.tg_database as database + +router = Router() + + +@router.message(F.content_type == ContentType.NEW_CHAT_MEMBERS) +async def user_join_handler(message: Message): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + for member in message.new_chat_members: + if member.is_bot: + continue + + response = chat['greeting_join'] or MESSAGE_DEFAULT_GREETING_JOIN + response = response.format(name=f'[{member.first_name}](tg://user?id={member.id})') + await message.answer(response, parse_mode=ParseMode.MARKDOWN) + + +@router.message(F.content_type == ContentType.LEFT_CHAT_MEMBER) +async def user_join_handler(message: Message): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + member = message.left_chat_member + if member.is_bot: + return + + database.DB.delete_user(chat_id, member.id) diff --git a/tg/handlers/admin.py b/tg/handlers/admin.py new file mode 100644 index 0000000..c470a21 --- /dev/null +++ b/tg/handlers/admin.py @@ -0,0 +1,103 @@ +from aiogram import Bot, Router, F +from aiogram.types import Message +from aiogram.utils.formatting import Bold + +import utils +from messages import * +import tg.tg_database as database + +router = Router() + + +async def tg_user_is_admin(chat_id: int, user_id: int, bot: Bot): + for admin in await bot.get_chat_administrators(chat_id=chat_id): + if admin.user.id == user_id: + return True + return False + + +@router.message(F.text == "!старт") +async def start_handler(message: Message, bot: Bot): + chat_id = message.chat.id + database.DB.create_chat_if_not_exists(chat_id) + + if not await tg_user_is_admin(chat_id, message.from_user.id, bot): + await message.answer(MESSAGE_PERMISSION_DENIED) + return + + database.DB.chat_update(chat_id, active=1) + await message.answer('Готова к работе!') + + +@router.message(F.text == "!правила") +async def rules_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + if message.reply_to_message is None: + if chat['rules'] is not None: + response = Bold('Правила чата') + '\n' + chat['rules'] + await message.answer(**response.as_kwargs()) + else: + await message.answer(MESSAGE_DEFAULT_RULES) + else: + if not await tg_user_is_admin(chat_id, message.from_user.id, bot): + await message.answer(MESSAGE_PERMISSION_DENIED) + return + + database.DB.chat_update(chat_id, rules=message.reply_to_message.text) + await message.answer('Правила чата изменены.') + + +@router.message(F.text == "!приветствие") +async def set_greeting_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + if not await tg_user_is_admin(chat_id, message.from_user.id, bot): + await message.answer(MESSAGE_PERMISSION_DENIED) + return + + if message.reply_to_message is None: + await message.answer(MESSAGE_NEED_REPLY) + return + + database.DB.chat_update(chat_id, greeting_join=message.reply_to_message.text) + await message.answer('Приветствие изменено.') + + +@router.message(F.text == "!предупреждение") +async def warning_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + if not await tg_user_is_admin(chat_id, message.from_user.id, bot): + await message.answer(MESSAGE_PERMISSION_DENIED) + return + + if message.reply_to_message is None: + await message.answer(MESSAGE_NEED_REPLY) + return + + user_id = message.reply_to_message.from_user.id + database.DB.create_user_if_not_exists(chat_id, user_id) + + database.DB.user_increment_warnings(chat_id, user_id) + user = database.DB.get_user(chat_id, user_id) + + user_info = message.reply_to_message.from_user + await message.answer('У {} {} {} {}.'.format( + user_info.first_name, + user_info.last_name, + user['warnings'], + utils.make_word_agree_with_number(user['warnings'], 'предупреждение')) + ) diff --git a/tg/handlers/default.py b/tg/handlers/default.py new file mode 100644 index 0000000..756cef3 --- /dev/null +++ b/tg/handlers/default.py @@ -0,0 +1,44 @@ +from aiogram import Router, F +from aiogram.types import Message +from aiogram.enums.content_type import ContentType + +import utils +import tg.tg_database as database + +router = Router() + +ACCEPTED_CONTENT_TYPES: list[ContentType] = [ + ContentType.TEXT, + ContentType.ANIMATION, + ContentType.AUDIO, + ContentType.DOCUMENT, + ContentType.PAID_MEDIA, + ContentType.PHOTO, + ContentType.STICKER, + ContentType.STORY, + ContentType.VIDEO, + ContentType.VIDEO_NOTE, + ContentType.VOICE, + ContentType.CHECKLIST, + ContentType.CONTACT, + ContentType.POLL, + ContentType.VENUE, + ContentType.LOCATION +] + + +@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES)) +async def any_message_handler(message: Message): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + return + + # Игнорировать ботов + if message.from_user.is_bot: + return + + user_id = message.from_user.id + database.DB.create_user_if_not_exists(chat_id, user_id) + database.DB.user_set_last_message(chat_id, user_id, utils.posix_time()) + database.DB.user_increment_messages(chat_id, user_id) diff --git a/tg/handlers/user.py b/tg/handlers/user.py new file mode 100644 index 0000000..1dd7c03 --- /dev/null +++ b/tg/handlers/user.py @@ -0,0 +1,129 @@ +from typing import List, Any + +from aiogram import Bot, Router, F +from aiogram.types import Message +from aiogram.utils.formatting import Bold, Italic + +from messages import * +import tg.tg_database as database + +router = Router() + + +# top_users - массив sqlite3.Row с двумя столбцами: user_id и value +async def format_rating(chat_id: int, top_users: List[Any], bot: Bot) -> str: + result = '' + i = 1 + for user in top_users: + info = await bot.get_chat_member(chat_id, user['user_id']) + result += '{}. {} {} - {}\n'.format(i, info.user.first_name, info.user.last_name, user['value']) + i = i + 1 + + return result + + +def calculate_total_messages(top_users: List[Any]) -> int: + total = 0 + for user in top_users: + total += user['value'] + return total + + +@router.message(F.text == '!помощь') +async def help_handler(message: Message): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + response = Bold('Команды для всех') + '\n' + response += '!правила - вывести правила\n' + response += '!сегодня - статистика сообщений за сегодня\n' + response += '!месяц - статистика сообщений за месяц\n' + response += '!молчуны - список молчунов\n' + response += '!предупреждения - список участников с предупреждениями\n' + response += '\n' + response += Bold('Команды для администраторов') + '\n' + response += '!старт - начать работу в чате\n' + response += '!правила* - изменить правила\n' + response += '!приветствие* - изменить приветствие новичков\n' + response += '!предупреждение* - выдать предупреждение участнику\n' + response += '\n' + response += Italic('Команды с пометкой * нужно вызывать в ответном сообщении.') + + await message.answer(**response.as_kwargs()) + + +@router.message(F.text == "!сегодня") +async def stats_today_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + top_users = database.DB.get_top_messages_today(chat_id) + if len(top_users) == 0: + await message.answer('Сегодня еще никто не писал.') + return + + response = Bold('Статистика за сегодня') + '\n' + response += 'Всего сообщений - {}\n'.format(calculate_total_messages(top_users)) + response += await format_rating(chat_id, top_users, bot) + await message.answer(**response.as_kwargs()) + + +@router.message(F.text == "!месяц") +async def stats_month_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + top_users = database.DB.get_top_messages_month(chat_id) + if len(top_users) == 0: + await message.answer('В этом месяце еще никто не писал.') + return + + response = Bold('Статистика за месяц') + '\n' + response += 'Всего сообщений - {}\n'.format(calculate_total_messages(top_users)) + response += await format_rating(chat_id, top_users, bot) + await message.answer(**response.as_kwargs()) + + +@router.message(F.text == "!молчуны") +async def silent_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + top_users = database.DB.get_top_messages_today(chat_id) + if len(top_users) == 0: + await message.answer('Молчунов нет. Все молодцы!') + return + + response = Bold('Молчуны чата') + ' (не писали более 14 дней)\n' + response += await format_rating(chat_id, top_users, bot) + await message.answer(**response.as_kwargs()) + + +@router.message(F.text == "!предупреждения") +async def warnings_handler(message: Message, bot: Bot): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + top_users = database.DB.get_top_warnings(chat_id) + if len(top_users) == 0: + await message.answer('Пока все спокойно. Продолжайте в том же духе.') + return + + response = Bold('Участники с предупреждениями') + '\n' + response += await format_rating(chat_id, top_users, bot) + await message.answer(**response.as_kwargs()) diff --git a/tg/tasks.py b/tg/tasks.py new file mode 100644 index 0000000..02617f4 --- /dev/null +++ b/tg/tasks.py @@ -0,0 +1,67 @@ +import datetime +import traceback + +from asyncio import sleep +from aiogram import Bot +from aiogram.utils.formatting import Bold + +from messages import * +import tg.tg_database as database +from tg.handlers.user import format_rating + + +async def reset_counters(reset_month: bool, bot: Bot): + database.DB.reset_messages_today() + print('Дневные счетчики сброшены.') + + if reset_month: + for chat in database.DB.get_chats(): + if chat['active'] == 0: + continue + + top_users = database.DB.get_top_messages_month(chat['id']) + if len(top_users) == 0: + continue + + message = Bold('Итоговая статистика за прошедший месяц') + '\n' + message += await format_rating(chat['id'], top_users, bot) + await bot.send_message(chat_id=chat['id'], **message.as_kwargs()) + + database.DB.reset_messages_month() + print('Месячные счетчики сброшены.') + + +async def wait_until(target_time: datetime.datetime): + now = datetime.datetime.now(target_time.tzinfo) + if now >= target_time: + return + + delay_seconds = (target_time - now).total_seconds() + await sleep(delay_seconds) + + +async def daily_maintenance_task(bot: Bot): + tz = datetime.timezone(datetime.timedelta(hours=3), name="MSK") + + target_time = datetime.time(6, 0, 0, tzinfo=tz) + now = datetime.datetime.now(tz) + if now.hour > target_time.hour or now.hour == target_time.hour and now.minute > target_time.minute: + target_date = now.date() + datetime.timedelta(days=1) + else: + target_date = now.date() + target_datetime = datetime.datetime.combine(target_date, target_time) + + while True: + await wait_until(target_datetime) + + try: + await reset_counters(target_datetime.day == 1, bot) + except Exception: + print(traceback.format_exc()) + + target_datetime = target_datetime + datetime.timedelta(days=1) + + +async def startup_task(bot: Bot): + me = await bot.get_me() + print(f"Бот '{me.full_name}' (id={me.id}) запущен.") diff --git a/tg/tg_bot.py b/tg/tg_bot.py new file mode 100644 index 0000000..bef74cc --- /dev/null +++ b/tg/tg_bot.py @@ -0,0 +1,20 @@ +import asyncio + +from aiogram import Bot, Dispatcher + +import config +import handlers +import tasks + + +async def main() -> None: + bot = Bot(token=config.Config['api_token']) + + dp = Dispatcher() + dp.include_router(handlers.router) + dp.startup.register(tasks.startup_task) + asyncio.create_task(tasks.daily_maintenance_task(bot)) + await dp.start_polling(bot) + + +asyncio.run(main()) diff --git a/tg/tg_database.py b/tg/tg_database.py new file mode 100644 index 0000000..02cb972 --- /dev/null +++ b/tg/tg_database.py @@ -0,0 +1,31 @@ +import database + + +class TgDatabase(database.BasicDatabase): + def __init__(self): + super().__init__() + + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS chats ( + "id" INTEGER, + "active" INTEGER NOT NULL DEFAULT 0, + "rules" TEXT, + "greeting_join" TEXT, + PRIMARY KEY("id")) + """) + + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS users ( + "chat_id" INTEGER, + "user_id" INTEGER, + "last_message" INTEGER NOT NULL DEFAULT 0, + "messages_today" INTEGER NOT NULL DEFAULT 0, + "messages_month" INTEGER NOT NULL DEFAULT 0, + "warnings" INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY("chat_id","user_id")) + """) + + self.conn.commit() + + +DB = TgDatabase() diff --git a/vk/handlers/__init__.py b/vk/handlers/__init__.py new file mode 100644 index 0000000..d7f529d --- /dev/null +++ b/vk/handlers/__init__.py @@ -0,0 +1,9 @@ +from vkbottle.framework.labeler import BotLabeler + +from . import user, admin, action, default + +labeler = BotLabeler() +labeler.load(user.labeler) +labeler.load(admin.labeler) +labeler.load(action.labeler) +labeler.load(default.labeler) diff --git a/handlers/action.py b/vk/handlers/action.py similarity index 84% rename from handlers/action.py rename to vk/handlers/action.py index 4510add..fb55e84 100644 --- a/handlers/action.py +++ b/vk/handlers/action.py @@ -1,15 +1,16 @@ from vkbottle.bot import Message +from vkbottle.framework.labeler import BotLabeler -import database from messages import * +import vk.vk_database as database -from labeler import labeler +labeler = BotLabeler() @labeler.chat_message(action=['chat_invite_user', 'chat_invite_user_by_link']) async def user_join_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: return @@ -37,7 +38,7 @@ async def user_join_handler(message: Message): @labeler.chat_message(action=['chat_kick_user']) async def user_leave_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: return diff --git a/handlers/admin.py b/vk/handlers/admin.py similarity index 90% rename from handlers/admin.py rename to vk/handlers/admin.py index 56ae675..aa16d76 100644 --- a/handlers/admin.py +++ b/vk/handlers/admin.py @@ -1,12 +1,13 @@ from vkbottle import bold, VKAPIError from vkbottle.bot import Message +from vkbottle.framework.labeler import BotLabeler from vkbottle_types.codegen.objects import MessagesGetConversationMembers -import database from messages import * import utils +import vk.vk_database as database -from labeler import labeler +labeler = BotLabeler() def vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers): @@ -20,7 +21,7 @@ def vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers) @labeler.chat_message(text="!старт") async def start_handler(message: Message): chat_id = message.peer_id - database.create_chat_if_not_exists(chat_id) + database.DB.create_chat_if_not_exists(chat_id) chat_members = await message.ctx_api.messages.get_conversation_members(peer_id=chat_id, extended=False) @@ -34,7 +35,7 @@ async def start_handler(message: Message): # Пропустить ботов if member.member_id < 0: continue - database.create_user_if_not_exists(chat_id, member.member_id) + database.DB.create_user_if_not_exists(chat_id, member.member_id) await message.answer('Готова к работе!') @@ -42,7 +43,7 @@ async def start_handler(message: Message): @labeler.chat_message(text="!правила") async def rules_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -65,7 +66,7 @@ async def rules_handler(message: Message): @labeler.chat_message(text="!приветствие") async def set_greeting_join_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -86,7 +87,7 @@ async def set_greeting_join_handler(message: Message): @labeler.chat_message(text="!возвращение") async def set_greeting_rejoin_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -107,7 +108,7 @@ async def set_greeting_rejoin_handler(message: Message): @labeler.chat_message(text="!деньрождения") async def set_birthday_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -128,7 +129,7 @@ async def set_birthday_handler(message: Message): @labeler.chat_message(text="!предупреждение") async def warning_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -144,7 +145,7 @@ async def warning_handler(message: Message): return user_id = message.reply_message.from_id - database.create_user_if_not_exists(chat_id, user_id) + database.DB.create_user_if_not_exists(chat_id, user_id) database.DB.user_increment_warnings(chat_id, user_id) user = database.DB.get_user(chat_id, user_id) @@ -162,7 +163,7 @@ async def warning_handler(message: Message): @labeler.chat_message(text="!исключить") async def ban_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return diff --git a/handlers/default.py b/vk/handlers/default.py similarity index 74% rename from handlers/default.py rename to vk/handlers/default.py index 584978f..121e94e 100644 --- a/handlers/default.py +++ b/vk/handlers/default.py @@ -1,15 +1,17 @@ from vkbottle.bot import Message +from vkbottle.framework.labeler import BotLabeler -import database import utils -from labeler import labeler +import vk.vk_database as database + +labeler = BotLabeler() # Обычные сообщения (не команды и не действия) @labeler.chat_message() async def any_message_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: return @@ -22,6 +24,6 @@ async def any_message_handler(message: Message): return user_id = message.from_id - database.create_user_if_not_exists(chat_id, user_id) + database.DB.create_user_if_not_exists(chat_id, user_id) database.DB.user_set_last_message(chat_id, user_id, utils.posix_time()) database.DB.user_increment_messages(chat_id, user_id) diff --git a/handlers/user.py b/vk/handlers/user.py similarity index 92% rename from handlers/user.py rename to vk/handlers/user.py index 03d9c17..c0be0c2 100644 --- a/handlers/user.py +++ b/vk/handlers/user.py @@ -2,11 +2,12 @@ from typing import List, Any from vkbottle import bold, italic, API from vkbottle.bot import Message +from vkbottle.framework.labeler import BotLabeler -import database from messages import * +import vk.vk_database as database -from labeler import labeler +labeler = BotLabeler() # top_users - массив sqlite3.Row с двумя столбцами: user_id и value @@ -37,7 +38,7 @@ def calculate_total_messages(top_users: List[Any]) -> int: @labeler.chat_message(text="!помощь") async def rules_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -67,7 +68,7 @@ async def rules_handler(message: Message): @labeler.chat_message(text="!сегодня") async def stats_today_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -86,7 +87,7 @@ async def stats_today_handler(message: Message): @labeler.chat_message(text="!месяц") async def stats_month_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -105,7 +106,7 @@ async def stats_month_handler(message: Message): @labeler.chat_message(text="!молчуны") async def silent_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -123,7 +124,7 @@ async def silent_handler(message: Message): @labeler.chat_message(text="!предупреждения") async def warnings_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return @@ -141,13 +142,13 @@ async def warnings_handler(message: Message): @labeler.chat_message(text="!поздравление") async def no_birthday_handler(message: Message): chat_id = message.peer_id - chat = database.create_chat_if_not_exists(chat_id) + chat = database.DB.create_chat_if_not_exists(chat_id) if chat['active'] == 0: await message.answer(MESSAGE_CHAT_NOT_ACTIVE) return user_id = message.from_id - user = database.create_user_if_not_exists(chat_id, user_id) + user = database.DB.create_user_if_not_exists(chat_id, user_id) happy_birthday = 1 if user['happy_birthday'] == 0 else 0 database.DB.user_toggle_happy_birthday(chat_id, user_id, happy_birthday) diff --git a/tasks.py b/vk/tasks.py similarity index 95% rename from tasks.py rename to vk/tasks.py index 7452c55..97f68d0 100644 --- a/tasks.py +++ b/vk/tasks.py @@ -4,9 +4,9 @@ import traceback from asyncio import sleep from vkbottle import API, bold -import database -from handlers.user import format_rating from messages import * +import vk.vk_database as database +from vk.handlers.user import format_rating async def reset_counters(reset_month: bool, api: API): @@ -46,7 +46,7 @@ async def check_birthdays(api: API): user_id = item.member_id if user_id < 0: continue - user = database.create_user_if_not_exists(chat_id, user_id) + user = database.DB.create_user_if_not_exists(chat_id, user_id) for profile in members.profiles: if profile.id == user_id: diff --git a/bot.py b/vk/vk_bot.py similarity index 59% rename from bot.py rename to vk/vk_bot.py index cbdcf23..ff8f9e5 100644 --- a/bot.py +++ b/vk/vk_bot.py @@ -2,13 +2,10 @@ import config import handlers import tasks -from labeler import labeler - -from vkbottle.bot import Bot +from vkbottle.bot import Bot as VkBot -assert handlers -bot = Bot(config.Config['api_token'], labeler=labeler) +bot = VkBot(config.Config['api_token'], labeler=handlers.labeler) bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api)) bot.loop_wrapper.add_task(tasks.daily_maintenance_task(bot.api)) bot.run_forever() diff --git a/vk/vk_database.py b/vk/vk_database.py new file mode 100644 index 0000000..7c3f383 --- /dev/null +++ b/vk/vk_database.py @@ -0,0 +1,37 @@ +import database + + +class VkDatabase(database.BasicDatabase): + def __init__(self): + super().__init__() + + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS chats ( + "id" INTEGER, + "active" INTEGER NOT NULL DEFAULT 0, + "rules" TEXT, + "greeting_join" TEXT, + "greeting_rejoin" TEXT, + "birthday_message" TEXT, + PRIMARY KEY("id")) + """) + + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS users ( + "chat_id" INTEGER, + "user_id" INTEGER, + "last_message" INTEGER NOT NULL DEFAULT 0, + "messages_today" INTEGER NOT NULL DEFAULT 0, + "messages_month" INTEGER NOT NULL DEFAULT 0, + "warnings" INTEGER NOT NULL DEFAULT 0, + "happy_birthday" INTEGER NOT NULL DEFAULT 1, + PRIMARY KEY("chat_id","user_id")) + """) + + self.conn.commit() + + def user_toggle_happy_birthday(self, chat_id: int, user_id: int, happy_birthday: int): + self.user_update(chat_id, user_id, happy_birthday=happy_birthday) + + +DB = VkDatabase()