Source code for

# Copyright (C) 2019  alfred richardsn
# This file is part of TellerBot.
# TellerBot is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# GNU Affero General Public License for more details.
# You should have received a copy of the GNU Affero General Public License
# along with TellerBot.  If not, see <>.
import asyncio
import logging
import typing
from time import time

from aiogram import Bot
from aiogram import types
from import api
from aiogram.contrib.middlewares.logging import LoggingMiddleware
from aiogram.dispatcher import Dispatcher
from aiogram.dispatcher.middlewares import BaseMiddleware
from pymongo import ReturnDocument

from src.config import config
from src.database import database
from src.database import database_user
from src.database import MongoStorage
from src.i18n import i18n

[docs]class IncomingHistoryMiddleware(BaseMiddleware): """Middleware for storing incoming history."""
[docs] async def trigger(self, action, args): """Save incoming data in the database.""" if ( "update" not in action and "error" not in action and action.startswith("pre_process_") ): await database.logs.insert_one( { "direction": "in", "type": action.split("pre_process_", 1)[1], "data": args[0].to_python(), } )
[docs]class TellerBot(Bot): """Custom bot class."""
[docs] async def request(self, method, data=None, *args, **kwargs): """Make a request and save it in the database.""" result = await super().request(method, data, *args, **kwargs) if ( config.DATABASE_LOGGING_ENABLED and result and method not in ( api.Methods.GET_UPDATES, api.Methods.SET_WEBHOOK, api.Methods.DELETE_WEBHOOK, api.Methods.GET_WEBHOOK_INFO, api.Methods.GET_ME, ) ): # On requests Telegram either returns True on success or relevant object. # To store only useful information, method's payload is saved if result is # a boolean and result is saved otherwise. await database.logs.insert_one( { "direction": "out", "type": method, "data": data if isinstance(result, bool) else result, } ) return result
[docs]class DispatcherManual(Dispatcher): """Dispatcher with user availability in database check."""
[docs] async def process_update(self, update: types.Update): """Process update object with user availability in database check. If bot doesn't know the user, it pretends they sent /start message. """ user = None if update.message: user = update.message.from_user chat = elif update.callback_query and update.callback_query.message: user = update.callback_query.from_user chat = if user: await database.users.update_many( {"id": {"$ne":}, "mention": user.mention}, {"$set": {"has_username": False}}, ) document = await database.users.find_one_and_update( {"id":, "chat":}, { "$set": { "mention": user.mention, "has_username": bool(user.username), } }, return_document=ReturnDocument.AFTER, ) if document is None: if update.message: if not update.message.text.startswith("/start "): update.message.text = "/start" elif update.callback_query: await update.callback_query.answer() update = types.Update( update_id=update.update_id, message={ "message_id": -1, "from": user.to_python(), "chat": chat.to_python(), "date": int(time()), "text": "/start", }, ) database_user.set(document) return await super().process_update(update)
tg = TellerBot(None, loop=asyncio.get_event_loop(), validate_token=False) dp = DispatcherManual(tg)
[docs]def setup(): """Set API token from config to bot and setup dispatcher.""" with open(config.TOKEN_FILENAME, "r") as token_file: tg._ctx_token.set( = MongoStorage() i18n.reload() dp.middleware.setup(i18n) logging.basicConfig(level=config.LOGGER_LEVEL) dp.middleware.setup(LoggingMiddleware()) if config.DATABASE_LOGGING_ENABLED: dp.middleware.setup(IncomingHistoryMiddleware())
[docs]def private_handler(*args, **kwargs): """Register handler only for private message.""" def decorator(handler: typing.Callable): dp.register_message_handler( handler, lambda message: == types.ChatType.PRIVATE, # noqa: E721 *args, **kwargs ) return handler return decorator
state_handlers = {}
[docs]def state_handler(state): """Associate ``state`` with decorated handler.""" def decorator(handler): state_handlers[state.state] = handler return handler return decorator