#!/usr/bin/env python import asyncio import logging from dataclasses import dataclass from http import HTTPStatus from config import BOT_TOKEN, TELEGRAM_WEBHOOK_URL, HEALTHCHECK_URL, FEDI_LOGIN_CALLBACK_URL, BOT_DOMAIN, BOT_PORT, ENCRYPT_KEY import uvicorn from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import PlainTextResponse, Response, JSONResponse from starlette.routing import Route from telegram import Update from telegram.error import BadRequest from telegram.ext import ( Application, CallbackContext, ContextTypes, ExtBot, CallbackQueryHandler, CommandHandler, MessageHandler, filters, ConversationHandler, TypeHandler, ) from callback import ( callback_generate_fedi_login_url, callback_skip_media, callback_location_sharing, callback_manual_location, callback_location_confirmation, callback_location_keyword_search, callback_skip_location_keyword_search, callback_add_comment, callback_skip_comment, callback_add_media, callback_delayed_checkin ) from command import ( start_command, fedi_login_command, cancel_command, help_command, tos_command, toggle_visibility_command, callback_toggle_visibility, logout_command, list_command, delayed_checkin_command, ) from config import ( FEDI_LOGIN, WAIT_LOCATION, PROMPT_FEDI_LOGIN, LOCATION_SEARCH_KEYWORD, LOCATION_CONFIRMATION, ADD_MEDIA, ADD_COMMENT, BOT_TOKEN, BOT_SCOPE, MAIN_MENU, WAIT_VISIBILITY, DELAYED_CHECKIN, ) from prompt.string import PROMPT_CHOOSE_ACTION from mastodon import Mastodon from dbstore.peewee_store import db, User, get_user_by_state from util import encrypt, decrypt logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO ) logger = logging.getLogger(__name__) @dataclass class FediLoginCallbackUpdate: code: str state: str user_id: str class FediLoginCallbackContext(CallbackContext[ExtBot, dict, dict, dict]): """ Custom CallbackContext class that makes `user_data` available for updates of type `WebhookUpdate`. """ @classmethod def from_update(cls, update: object, application: "Application") -> "FediLoginCallbackContext": if isinstance(update, FediLoginCallbackUpdate): return cls(application=application, user_id=int(update.user_id)) return super().from_update(update, application) async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: FediLoginCallbackContext) -> None: state = update.state with db.connection_context(): user = User.get(User.state == state) client_id = user.client_id client_secret = user.client_secret home_instance = user.home_instance if len(user.access_key) == 0: mastodon_client = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance, version_check_mode="none") access_token = mastodon_client.log_in( code=update.code, redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL), scopes=BOT_SCOPE ) instance_info = mastodon_client.instance_nodeinfo() if instance_info["software"]["name"] == "pleroma": user.home_instance_type = "pleroma" user.access_key = encrypt(access_token, ENCRYPT_KEY) user.save() text = "You have successfully logged in to your Mastodon account!" try: await context.bot.delete_message(chat_id=user.telegram_user_id, message_id=context.user_data[PROMPT_FEDI_LOGIN]) await context.bot.send_message(chat_id=user.telegram_user_id, text=text) await context.bot.send_message(chat_id=user.telegram_user_id, text=PROMPT_CHOOSE_ACTION, reply_markup=MAIN_MENU) except BadRequest as e: if "not found" in str(e.message): pass async def main() -> None: context_types = ContextTypes(context=FediLoginCallbackContext) application = ( Application.builder().updater(None).token(BOT_TOKEN).context_types(context_types).build() ) # TODO: # check user login status before invoking commands checkin_handler = ConversationHandler( entry_points=[ CommandHandler("start", start_command), CommandHandler("login", fedi_login_command), MessageHandler(filters.LOCATION, callback_location_sharing), ], states={ FEDI_LOGIN: [ MessageHandler(filters.TEXT & ~filters.COMMAND, callback_generate_fedi_login_url), ], WAIT_LOCATION: [ MessageHandler(filters.LOCATION, callback_location_sharing), ], LOCATION_SEARCH_KEYWORD: [ MessageHandler(filters.TEXT & ~filters.COMMAND, callback_location_keyword_search), CallbackQueryHandler(callback_skip_location_keyword_search), ], LOCATION_CONFIRMATION: [ CallbackQueryHandler(callback_location_confirmation), MessageHandler(filters.TEXT & ~filters.COMMAND, callback_manual_location) ], ADD_COMMENT: [ MessageHandler(filters.TEXT & ~filters.COMMAND, callback_add_comment), CallbackQueryHandler(callback_skip_comment), ], ADD_MEDIA: [MessageHandler(filters.PHOTO, callback_add_media), CallbackQueryHandler(callback_skip_media)], }, fallbacks=[CommandHandler("cancel", cancel_command)], per_message=False, allow_reentry=True, ) # register handlers application.add_handler(CommandHandler("tos", tos_command)) visibility_conversation_handler = ConversationHandler( entry_points=[ CommandHandler("vis", toggle_visibility_command) ], states={ WAIT_VISIBILITY: [ CallbackQueryHandler(callback_toggle_visibility) ]}, fallbacks=[CommandHandler("cancel", cancel_command)], per_message=False, allow_reentry=True, ) delayed_checkin_handler = ConversationHandler( entry_points=[ CommandHandler("delay", delayed_checkin_command), ], states={ DELAYED_CHECKIN: [ MessageHandler(filters.TEXT & ~filters.COMMAND, callback_delayed_checkin), ], }, fallbacks=[CommandHandler("cancel", cancel_command)], per_message=False, allow_reentry=True, ) application.add_handler(CommandHandler("logout", logout_command)) application.add_handler(CommandHandler("list", list_command)) application.add_handler(CommandHandler("Help", help_command)) application.add_handler(TypeHandler(type=FediLoginCallbackUpdate, callback=process_oauth_login_callback)) application.add_handler(delayed_checkin_handler, 3) application.add_handler(visibility_conversation_handler, 2) application.add_handler(checkin_handler, 1) # Pass webhook settings to telegram await application.bot.set_webhook(url=f"{BOT_DOMAIN}{TELEGRAM_WEBHOOK_URL}") # Set up webserver async def telegram_webhook(request: Request) -> JSONResponse: """Handle incoming Telegram updates by putting them into the `update_queue`""" await application.update_queue.put( Update.de_json(data=await request.json(), bot=application.bot) ) return JSONResponse({'OK': 200}) async def fedi_oauth_login_callback(request: Request) -> PlainTextResponse: """ Handle incoming webhook updates by also putting them into the `update_queue` if the required parameters were passed correctly. """ try: code = request.query_params["code"] state = request.query_params.get("state") user = get_user_by_state(state) except KeyError: return PlainTextResponse( status_code=HTTPStatus.BAD_REQUEST, content="Mastodon callback request doesn't contain a valid OAuth code", ) await application.update_queue.put(FediLoginCallbackUpdate(state=state, code=code, user_id=user["telegram_user_id"])) return PlainTextResponse("Thank you for login! Now you can close the browser") async def healthcheck(_: Request) -> PlainTextResponse: return PlainTextResponse(content="OK") starlette_app = Starlette( routes=[ Route(TELEGRAM_WEBHOOK_URL, telegram_webhook, methods=["POST"]), Route(HEALTHCHECK_URL, healthcheck, methods=["GET"]), Route(FEDI_LOGIN_CALLBACK_URL, fedi_oauth_login_callback, methods=["GET"]), ] ) webserver = uvicorn.Server( config=uvicorn.Config( app=starlette_app, port=BOT_PORT, use_colors=False, host="0.0.0.0", ) ) # Run application and webserver together async with application: await application.start() await webserver.serve() await application.stop() if __name__ == "__main__": asyncio.run(main())