vk_chat_bot/bot.py

import asyncio
import calendar
import datetime
import time

from vkbottle.bot import Bot, Message
from vkbottle_types.objects import MessagesGetConversationMembers

import config
from config import config_load
import database


config_load()
bot = Bot(config.Config['api_token'])


def posix_time():
    return calendar.timegm(time.gmtime())


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 vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers):
    for member in chat_members.items:
        if member.member_id != user_id:
            continue
        return member.is_admin
    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 start_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

    for member in chat_members.items:
        # Пропустить ботов
        if member.member_id < 0:
            continue
        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):
    top_users = database.DB.get_top_messages_today(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['messages_today'])
        i = i + 1

    await message.answer(response)


@bot.on.chat_message(text="!месяц")
async def stats_month_handler(message: Message):
    top_users = database.DB.get_top_messages_month(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['messages_month'])
        i = i + 1

    await message.answer(response)


@bot.on.chat_message(text="!молчуны")
async def silent_handler(message: Message):
    now = posix_time()
    threshold = now - 14 * 24 * 3600
    top_users = database.DB.get_top_silent(message.peer_id, threshold)

    response = '* Молчуны чата (не писали больше 2 недель):\n'
    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 - 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 = 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):
    now = datetime.datetime.now(target_time.tzinfo)
    if now >= target_time:
        return

    delay_seconds = (target_time - now).total_seconds()
    await asyncio.sleep(delay_seconds)


async def counters_reset_task():
    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)

        print('Resetting daily counters...')
        database.DB.reset_messages_today()

        if target_datetime.day == 1:
            print('Resetting monthly counters...')
            database.DB.reset_messages_month()

        target_datetime = target_datetime + datetime.timedelta(days=1)


async def startup_task():
    print("Bot started.")


bot.loop_wrapper.on_startup.append(startup_task())
bot.loop_wrapper.add_task(counters_reset_task())
bot.run_forever()