TG Bot: переписал на aiogram
This commit is contained in:
Alexey Polyakov
parent
e5607adb9b
commit
0f5c06f80c
|
|
@ -1,4 +1,4 @@
|
|||
pyTelegramBotAPI
|
||||
aiogram
|
||||
aiomysql
|
||||
msgpack
|
||||
lz4
|
||||
|
|
|
|||
|
|
@ -2,15 +2,22 @@ import logging
|
|||
import random
|
||||
import json
|
||||
import time
|
||||
from telebot.async_telebot import AsyncTeleBot
|
||||
from textwrap import dedent
|
||||
|
||||
from aiogram import Bot, Dispatcher, Router
|
||||
from aiogram.filters import Command
|
||||
from aiogram.types import Message
|
||||
|
||||
from common.static import Static
|
||||
from common.sql_queries import SQLQueries
|
||||
from common.tools import Tools
|
||||
|
||||
class TelegramBot:
|
||||
def __init__(self, token, enabled, db_pool, whitelist_ids=None):
|
||||
self.bot = AsyncTeleBot(token)
|
||||
self.bot = Bot(token=token)
|
||||
self.dp = Dispatcher()
|
||||
self.router = Router()
|
||||
self.dp.include_router(self.router)
|
||||
self.tools = Tools()
|
||||
self.enabled = enabled
|
||||
self.db_pool = db_pool
|
||||
|
|
@ -20,104 +27,101 @@ class TelegramBot:
|
|||
self.msg_types = Static().BotMessageTypes()
|
||||
self.static = Static()
|
||||
self.sql_queries = SQLQueries()
|
||||
|
||||
self.router.message.register(self.handle_start, Command("start"))
|
||||
self.router.message.register(self.handle_register, Command("register"))
|
||||
|
||||
@self.bot.message_handler(commands=['start'])
|
||||
async def handle_start(message):
|
||||
tg_id = str(message.from_user.id)
|
||||
|
||||
# Ищем привязанный аккаунт пользователя
|
||||
async with self.db_pool.acquire() as conn:
|
||||
async with conn.cursor() as cursor:
|
||||
await cursor.execute(self.sql_queries.SELECT_USER_BY_TG_ID, (tg_id,))
|
||||
account = await cursor.fetchone()
|
||||
|
||||
if account:
|
||||
# Извлекаем id аккаунта с телефоном
|
||||
phone = account.get('phone')
|
||||
|
||||
await self.bot.send_message(
|
||||
message.chat.id,
|
||||
self.get_bot_message(self.msg_types.WELCOME_ALREADY_REGISTERED).format(phone=phone)
|
||||
)
|
||||
else:
|
||||
await self.bot.send_message(
|
||||
message.chat.id, self.get_bot_message(self.msg_types.WELCOME_NEW_USER)
|
||||
)
|
||||
|
||||
@self.bot.message_handler(commands=['register'])
|
||||
async def handle_register(message):
|
||||
tg_id = str(message.from_user.id)
|
||||
|
||||
# Проверка ID на наличие в белом списке
|
||||
if tg_id not in self.whitelist_ids:
|
||||
await self.bot.send_message(message.chat.id, self.get_bot_message(self.msg_types.ID_NOT_WHITELISTED))
|
||||
return
|
||||
|
||||
async with self.db_pool.acquire() as conn:
|
||||
async with conn.cursor() as cursor:
|
||||
# Проверка на существование
|
||||
await cursor.execute(self.sql_queries.SELECT_USER_BY_TG_ID, (tg_id,))
|
||||
if await cursor.fetchone():
|
||||
await self.bot.send_message(
|
||||
message.chat.id,
|
||||
self.get_bot_message(self.msg_types.ACCOUNT_ALREADY_EXISTS)
|
||||
)
|
||||
return
|
||||
|
||||
# Подготовка данных согласно схеме
|
||||
new_phone = f"7900{random.randint(1000000, 9999999)}"
|
||||
updatetime = str(int(time.time() * 1000))
|
||||
lastseen = str(int(time.time()))
|
||||
|
||||
try:
|
||||
# Создаем юзера
|
||||
await cursor.execute(
|
||||
self.sql_queries.INSERT_USER,
|
||||
(
|
||||
self.tools.generate_id(),
|
||||
new_phone, # phone
|
||||
tg_id, # telegram_id
|
||||
message.from_user.first_name[:59], # firstname
|
||||
(message.from_user.last_name or "")[:59], # lastname
|
||||
(message.from_user.username or "")[:60], # username
|
||||
json.dumps([]), # profileoptions
|
||||
json.dumps(["TT", "ONEME"]), # options
|
||||
0, # accountstatus
|
||||
updatetime,
|
||||
lastseen,
|
||||
)
|
||||
)
|
||||
|
||||
# Добавляем данные о аккаунте
|
||||
await cursor.execute(
|
||||
self.sql_queries.INSERT_USER_DATA,
|
||||
(
|
||||
new_phone, # phone
|
||||
json.dumps([]), # contacts
|
||||
json.dumps(self.static.USER_FOLDERS), # folders
|
||||
json.dumps(self.static.USER_SETTINGS), # user settings
|
||||
json.dumps({}), # chat_config
|
||||
)
|
||||
)
|
||||
|
||||
await self.bot.send_message(
|
||||
message.chat.id,
|
||||
self.get_bot_message(self.msg_types.REGISTRATION_SUCCESS).format(new_phone=new_phone)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Ошибка при регистрации: {e}")
|
||||
await self.bot.send_message(
|
||||
message.chat.id,
|
||||
self.get_bot_message(self.msg_types.INTERNAL_ERROR)
|
||||
)
|
||||
|
||||
def get_bot_message(self, msg_type):
|
||||
return dedent(self.static.BOT_MESSAGES.get(msg_type)).strip()
|
||||
|
||||
async def handle_start(self, message: Message):
|
||||
tg_id = str(message.from_user.id)
|
||||
|
||||
# Ищем привязанный аккаунт пользователя
|
||||
async with self.db_pool.acquire() as conn:
|
||||
async with conn.cursor() as cursor:
|
||||
await cursor.execute(self.sql_queries.SELECT_USER_BY_TG_ID, (tg_id,))
|
||||
account = await cursor.fetchone()
|
||||
|
||||
if account:
|
||||
# Извлекаем id аккаунта с телефоном
|
||||
phone = account.get('phone')
|
||||
|
||||
await message.answer(
|
||||
self.get_bot_message(self.msg_types.WELCOME_ALREADY_REGISTERED).format(phone=phone)
|
||||
)
|
||||
else:
|
||||
await message.answer(
|
||||
self.get_bot_message(self.msg_types.WELCOME_NEW_USER)
|
||||
)
|
||||
|
||||
async def handle_register(self, message: Message):
|
||||
tg_id = str(message.from_user.id)
|
||||
|
||||
# Проверка ID на наличие в белом списке
|
||||
if tg_id not in self.whitelist_ids:
|
||||
await message.answer(self.get_bot_message(self.msg_types.ID_NOT_WHITELISTED))
|
||||
return
|
||||
|
||||
async with self.db_pool.acquire() as conn:
|
||||
async with conn.cursor() as cursor:
|
||||
# Проверка на существование
|
||||
await cursor.execute(self.sql_queries.SELECT_USER_BY_TG_ID, (tg_id,))
|
||||
if await cursor.fetchone():
|
||||
await message.answer(
|
||||
self.get_bot_message(self.msg_types.ACCOUNT_ALREADY_EXISTS)
|
||||
)
|
||||
return
|
||||
|
||||
# Подготовка данных согласно схеме
|
||||
new_phone = f"7900{random.randint(1000000, 9999999)}"
|
||||
updatetime = str(int(time.time() * 1000))
|
||||
lastseen = str(int(time.time()))
|
||||
|
||||
try:
|
||||
# Создаем юзера
|
||||
await cursor.execute(
|
||||
self.sql_queries.INSERT_USER,
|
||||
(
|
||||
self.tools.generate_id(),
|
||||
new_phone, # phone
|
||||
tg_id, # telegram_id
|
||||
message.from_user.first_name[:59], # firstname
|
||||
(message.from_user.last_name or "")[:59], # lastname
|
||||
(message.from_user.username or "")[:60], # username
|
||||
json.dumps([]), # profileoptions
|
||||
json.dumps(["TT", "ONEME"]), # options
|
||||
0, # accountstatus
|
||||
updatetime,
|
||||
lastseen,
|
||||
)
|
||||
)
|
||||
|
||||
# Добавляем данные о аккаунте
|
||||
await cursor.execute(
|
||||
self.sql_queries.INSERT_USER_DATA,
|
||||
(
|
||||
new_phone, # phone
|
||||
json.dumps([]), # contacts
|
||||
json.dumps(self.static.USER_FOLDERS), # folders
|
||||
json.dumps(self.static.USER_SETTINGS), # user settings
|
||||
json.dumps({}), # chat_config
|
||||
)
|
||||
)
|
||||
|
||||
await message.answer(
|
||||
self.get_bot_message(self.msg_types.REGISTRATION_SUCCESS).format(new_phone=new_phone)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Ошибка при регистрации: {e}")
|
||||
await message.answer(
|
||||
self.get_bot_message(self.msg_types.INTERNAL_ERROR)
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
if self.enabled == True:
|
||||
try:
|
||||
await self.bot.polling()
|
||||
await self.dp.start_polling(self.bot)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Ошибка запуска Telegram бота: {e}")
|
||||
else:
|
||||
|
|
@ -126,7 +130,8 @@ class TelegramBot:
|
|||
async def send_auth_code(self, chat_id, phone, code):
|
||||
try:
|
||||
await self.bot.send_message(
|
||||
chat_id, self.get_bot_message(self.msg_types.INCOMING_CODE).format(phone=phone, code=code)
|
||||
chat_id,
|
||||
self.get_bot_message(self.msg_types.INCOMING_CODE).format(phone=phone, code=code)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Ошибка отправки кода в Telegram: {e}")
|
||||
Loading…
Reference in New Issue