diff --git a/.gitignore b/.gitignore index 5b92d54..f062bfb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ .idea .venv -db.json +bot.db +config.json diff --git a/bot.py b/bot.py index 857e7fe..925ee98 100644 --- a/bot.py +++ b/bot.py @@ -7,25 +7,26 @@ from vkbottle.bot import Bot, Message from vkbottle_types.objects import MessagesGetConversationMembers import config -from config import db_load, db_save +from config import config_load +import database -db_load() -db_save() -bot = Bot(config.DB['api_token']) + +config_load() +bot = Bot(config.Config['api_token']) def posix_time(): return calendar.timegm(time.gmtime()) -async def get_user_info(user_id): +async def vk_get_user_info(user_id): info = await bot.api.users.get(user_ids=[user_id]) if len(info) > 0: return info[0] return None -def user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers): +def vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers): for member in chat_members.items: if member.member_id != user_id: continue @@ -33,51 +34,80 @@ def user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers): return False +def create_user_if_not_exists(chat_id: str, user_id: str): + res = database.DB.get_user(chat_id, user_id) + if res is None: + database.DB.add_user(chat_id, user_id) + + @bot.on.chat_message(text="!старт") -async def stats_handler(message: Message): +async def start_handler(message: Message): chat_members = await bot.api.messages.get_conversation_members(peer_id=message.peer_id, extended=False) - if not user_is_admin(message.from_id, chat_members): + if not vk_user_is_admin(message.from_id, chat_members): await message.answer('О таком меня может попросить только администратор беседы.') return - chat_id = str(message.peer_id) - if chat_id not in config.DB['chats']: - config.DB['chats'][chat_id] = {'users': {}} - for member in chat_members.items: + # Пропустить ботов if member.member_id < 0: continue - member_id = str(member.member_id) - if member_id not in config.DB['chats'][chat_id]['users']: - config.DB['chats'][chat_id]['users'][member_id] = { - 'last_message': 0, - 'messages_today': 0, - 'messages_month': 0 - } - - db_save() + create_user_if_not_exists(message.peer_id, member.member_id) await message.answer('Готова к работе!') +@bot.on.chat_message(text="!предупреждение") +async def warning_handler(message: Message): + chat_members = await bot.api.messages.get_conversation_members(peer_id=message.peer_id, extended=False) + + if not vk_user_is_admin(message.from_id, chat_members): + await message.answer('О таком меня может попросить только администратор беседы.') + return + + if message.reply_message is None: + await message.answer('Эту команду нужно вызывать в ответном сообщении.') + return + + chat_id = message.peer_id + user_id = message.reply_message.from_id + 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 = await vk_get_user_info(user_id) + if user_info is not None: + await message.answer('У участника {} {} {} предупреждений.'.format( + user_info.first_name, + user_info.last_name, + user['warnings']) + ) + + +@bot.on.chat_message(text="!предупреждения") +async def warnings_handler(message: Message): + top_users = database.DB.get_top_warnings(message.peer_id) + + response = '* Участники с предупреждениями:\n' + i = 1 + for user in top_users: + info = await vk_get_user_info(user['user_id']) + response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['warnings']) + i = i + 1 + + await message.answer(response) + + @bot.on.chat_message(text="!сегодня") async def stats_today_handler(message: Message): - chat_id = str(message.peer_id) + top_users = database.DB.get_top_messages_today(message.peer_id) - chat_users = [] - for user_id in config.DB['chats'][chat_id]['users'].keys(): - flat_user = {**{'id': user_id}, **config.DB['chats'][chat_id]['users'][user_id]} - chat_users.append(flat_user) - - top = sorted(chat_users, key=lambda item: item['messages_today'], reverse=True) response = '* Статистика за сегодня:\n' - i = 0 - while i < len(top): - if top[i]['messages_today'] == 0: - break - info = await get_user_info(top[i]['id']) - response += '{}. {} {} - {}\n'.format(i + 1, info.first_name, info.last_name, top[i]['messages_today']) + i = 1 + for user in top_users: + info = await vk_get_user_info(user['user_id']) + response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['messages_today']) i = i + 1 await message.answer(response) @@ -85,21 +115,13 @@ async def stats_today_handler(message: Message): @bot.on.chat_message(text="!месяц") async def stats_month_handler(message: Message): - chat_id = str(message.peer_id) + top_users = database.DB.get_top_messages_month(message.peer_id) - chat_users = [] - for user_id in config.DB['chats'][chat_id]['users'].keys(): - flat_user = {**{'id': user_id}, **config.DB['chats'][chat_id]['users'][user_id]} - chat_users.append(flat_user) - - top = sorted(chat_users, key=lambda item: item['messages_month'], reverse=True) - response = '\n* Статистика за месяц:\n' - i = 0 - while i < len(top): - if top[i]['messages_month'] == 0: - break - info = await get_user_info(top[i]['id']) - response += '{}. {} {} - {}\n'.format(i + 1, info.first_name, info.last_name, top[i]['messages_month']) + response = '* Статистика за месяц:\n' + i = 1 + for user in top_users: + info = await vk_get_user_info(user['user_id']) + response += '{}. {} {} - {}\n'.format(i, info.first_name, info.last_name, user['messages_month']) i = i + 1 await message.answer(response) @@ -107,60 +129,34 @@ async def stats_month_handler(message: Message): @bot.on.chat_message(text="!молчуны") async def silent_handler(message: Message): - chat_id = str(message.peer_id) - - chat_users = [] - for user_id in config.DB['chats'][chat_id]['users'].keys(): - flat_user = {**{'id': user_id}, - **{'last_message': config.DB['chats'][chat_id]['users'][user_id]['last_message']}} - chat_users.append(flat_user) - top = sorted(chat_users, key=lambda item: item['last_message']) - now = posix_time() + threshold = now - 14 * 24 * 3600 + top_users = database.DB.get_top_silent(message.peer_id, threshold) response = '* Молчуны чата (не писали больше 2 недель):\n' - for i in range(0, len(top)): - if top[i]['last_message'] == 0: - info = await get_user_info(top[i]['id']) - if info is None: - continue - response += '{}. {} {} - никогда\n'.format(i + 1, info.first_name, info.last_name) + i = 1 + for user in top_users: + info = await vk_get_user_info(user['user_id']) + if user['last_message'] == 0: + response += '{}. {} {} - никогда\n'.format(i, info.first_name, info.last_name) else: - days_silent = round((now - top[i]['last_message']) / 3600 / 24) - if days_silent < 14: - break - info = await get_user_info(top[i]['id']) - response += '{}. {} {} - {} дней\n'.format(i + 1, info.first_name, info.last_name, days_silent) + days_silent = round((now - user['last_message']) / 3600 / 24) + response += '{}. {} {} - {} дней\n'.format(i, info.first_name, info.last_name, days_silent) await message.answer(response) @bot.on.chat_message() async def any_message_handler(message: Message): + # Игнорировать ботов if message.from_id < 0: return - chat_id = str(message.peer_id) - user_id = str(message.from_id) - if chat_id not in config.DB['chats']: - config.DB['chats'][chat_id] = {'users': {}} - if user_id not in config.DB['chats'][chat_id]['users']: - config.DB['chats'][chat_id]['users'][user_id] = {} - - config.DB['chats'][chat_id]['users'][user_id]['last_message'] = posix_time() - config.DB['chats'][chat_id]['users'][user_id]['messages_today'] = ( - config.DB['chats'][chat_id]['users'][user_id].get('messages_today', 0) + 1) - config.DB['chats'][chat_id]['users'][user_id]['messages_month'] = ( - config.DB['chats'][chat_id]['users'][user_id].get('messages_month', 0) + 1) - - db_save() - - -# @bot.on.raw_event(GroupEventType.GROUP_JOIN, GroupTypes.GroupJoin) -# async def group_join_handler(event: GroupTypes.GroupJoin): -# chat_id = str(event.group_id) -# user_id = str(event.object.user_id) -# if user_id == bot.api.users.get() + chat_id = message.peer_id + user_id = message.from_id + create_user_if_not_exists(chat_id, user_id) + database.DB.user_set_last_message(chat_id, user_id, posix_time()) + database.DB.user_increment_messages(chat_id, user_id) async def wait_until(target_time: datetime.datetime): @@ -187,17 +183,12 @@ async def counters_reset_task(): await wait_until(target_datetime) print('Resetting daily counters...') - for chat_id in config.DB['chats']: - for user_id in config.DB['chats'][chat_id]['users']: - config.DB['chats'][chat_id]['users'][user_id]['messages_today'] = 0 + database.DB.reset_messages_today() if target_datetime.day == 1: print('Resetting monthly counters...') - for chat_id in config.DB['chats']: - for user_id in config.DB['chats'][chat_id]['users']: - config.DB['chats'][chat_id]['users'][user_id]['messages_month'] = 0 + database.DB.reset_messages_month() - db_save() target_datetime = target_datetime + datetime.timedelta(days=1) diff --git a/config.py b/config.py index 6be0cc5..d44836e 100644 --- a/config.py +++ b/config.py @@ -1,39 +1,23 @@ import json -# db = { -# chats: { -# CHAT_ID: { -# users: { -# USER_ID: { -# messages_today: 10 -# messages_in_month: 1045 -# last_message: 1755627442 -# } -# } -# } -# }, +# Config = { # api_token: "xxxxxxxx" # } -DB = {} +Config = {} -def db_load(): - global DB +def config_load(): + global Config try: - file = open('db.json', 'r') - DB = json.load(file) + file = open('config.json', 'r') + Config = json.load(file) file.close() except IOError: - DB = {'chats': {}, 'api_token': ''} + Config = {'api_token': ''} -def db_save(): - global DB - with open('db.json', 'w') as file: - json.dump(DB, file, indent=4) - - -def db_print(): - global DB - print(DB) +def config_save(): + global Config + with open('config.json', 'w') as file: + json.dump(Config, file, indent=4) diff --git a/database.py b/database.py new file mode 100644 index 0000000..5b457c4 --- /dev/null +++ b/database.py @@ -0,0 +1,99 @@ +import sqlite3 +from typing import List + + +class Database: + 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 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() + + def get_user(self, chat_id: int, user_id: int): + self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id)) + return self.cursor.fetchone() + + def get_users(self, chat_id: int): + self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,)) + return self.cursor.fetchall() + + def add_user(self, chat_id: int, user_id: int): + self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", (chat_id, user_id)) + self.conn.commit() + + def user_set_last_message(self, chat_id: int, user_id: int, last_message: int): + self.user_update(chat_id, user_id, last_message=last_message) + + def user_increment_messages(self, chat_id: int, user_id: int): + self.user_increment(chat_id, user_id, ['messages_today', 'messages_month']) + self.conn.commit() + + def user_increment_warnings(self, chat_id: int, user_id: int): + self.user_increment(chat_id, user_id, ['warnings']) + + 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)) + self.conn.commit() + + def user_update(self, chat_id: int, user_id: int, **kwargs): + # query = ("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) + + # " WHERE chat_id = ? AND user_id = ?") + # print(query) + self.cursor.execute("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) + + f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs) + self.conn.commit() + + def get_top_messages_today(self, chat_id: int): + self.cursor.execute(""" + SELECT user_id, messages_today FROM users + WHERE chat_id = ? AND messages_today > 0 + ORDER BY messages_today DESC + """, (chat_id,)) + return self.cursor.fetchall() + + def get_top_messages_month(self, chat_id: int): + self.cursor.execute(""" + SELECT user_id, messages_month FROM users + WHERE chat_id = ? AND messages_month > 0 + ORDER BY messages_month DESC + """, (chat_id,)) + return self.cursor.fetchall() + + def get_top_silent(self, chat_id: int, threshold_time: int): + self.cursor.execute(""" + SELECT user_id, last_message FROM users + WHERE chat_id = ? AND last_message < ? + ORDER BY last_message ASC + """, (chat_id, threshold_time)) + return self.cursor.fetchall() + + def get_top_warnings(self, chat_id: int): + self.cursor.execute(""" + SELECT user_id, warnings FROM users + WHERE chat_id = ? AND warnings > 0 + ORDER BY warnings DESC + """, (chat_id,)) + return self.cursor.fetchall() + + def reset_messages_today(self): + self.cursor.execute("UPDATE users SET messages_today = 0") + self.conn.commit() + + def reset_messages_month(self): + self.cursor.execute("UPDATE users SET messages_month = 0") + self.conn.commit() + + +DB = Database()