From adb574ab528fbc4401fb7cd50b8748d1fb529450 Mon Sep 17 00:00:00 2001 From: clarkzjw Date: Thu, 23 Feb 2023 00:30:52 -0800 Subject: bot: multi user working --- bot.py | 41 +++++++++++++++++++++++------------------ command.py | 14 ++++++++++---- dbstore/peewee_store.py | 9 +++++++++ 3 files changed, 42 insertions(+), 22 deletions(-) diff --git a/bot.py b/bot.py index a059c9c..35884ff 100644 --- a/bot.py +++ b/bot.py @@ -5,32 +5,24 @@ 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 - import uvicorn from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import PlainTextResponse, Response from starlette.routing import Route - - from telegram import Update from telegram.ext import ( Application, CallbackContext, ContextTypes, ExtBot, - TypeHandler, -) - -from telegram.ext import ( - Application, CallbackQueryHandler, CommandHandler, MessageHandler, filters, - ConversationHandler + ConversationHandler, + TypeHandler, ) - from callback import ( callback_generate_fedi_login_url, callback_skip_media, @@ -52,18 +44,20 @@ from command import ( from config import ( FEDI_LOGIN, WAIT_LOCATION, + PROMPT_FEDI_LOGIN, LOCATION_SEARCH_KEYWORD, LOCATION_CONFIRMATION, ADD_MEDIA, ADD_COMMENT, BOT_TOKEN, - BOT_SCOPE + BOT_SCOPE, + MAIN_MENU ) -from mastodon import Mastodon -from dbstore.peewee_store import db, User +from prompt.string import PROMPT_CHOOSE_ACTION +from mastodon import Mastodon +from dbstore.peewee_store import db, User, get_user_by_state -# Enable logging logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO ) @@ -74,10 +68,19 @@ logger = logging.getLogger(__name__) class FediLoginCallbackUpdate: code: str state: str + user_id: str class FediLoginCallbackContext(CallbackContext[ExtBot, dict, dict, dict]): - pass + """ + 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: @@ -101,13 +104,13 @@ async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: user.save() text = "You have successfully logged in to your Mastodon account!" + 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) async def main() -> None: context_types = ContextTypes(context=FediLoginCallbackContext) - # Here we set updater to None because we want our custom webhook server to handle the updates - # and hence we don't need an Updater instance application = ( Application.builder().updater(None).token(BOT_TOKEN).context_types(context_types).build() ) @@ -169,13 +172,15 @@ async def main() -> None: 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)) + 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: diff --git a/command.py b/command.py index 901b792..7426901 100644 --- a/command.py +++ b/command.py @@ -2,15 +2,21 @@ from telegram import Update from telegram.constants import ParseMode from telegram.error import BadRequest from telegram.ext import ContextTypes, ConversationHandler - +from dbstore.peewee_store import get_user_access_key from config import * async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: await update.message.reply_text(PROMPT_START, parse_mode=ParseMode.MARKDOWN) - await update.message.reply_text(PROMPT_CHOOSE_ACTION, reply_markup=MAIN_MENU) - - return WAIT_LOCATION + user_access_key = get_user_access_key(str(update.effective_user.id)) + # TODO + # verify user access key still valid + if len(user_access_key) == 0: + await update.message.reply_text(PROMPT_FEDI_LOGIN_WHERE_IS_INSTANCE, parse_mode=ParseMode.MARKDOWN) + return FEDI_LOGIN + else: + await update.message.reply_text(PROMPT_CHOOSE_ACTION, reply_markup=MAIN_MENU) + return WAIT_LOCATION async def fedi_login_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: diff --git a/dbstore/peewee_store.py b/dbstore/peewee_store.py index 3d37140..272e230 100644 --- a/dbstore/peewee_store.py +++ b/dbstore/peewee_store.py @@ -41,6 +41,15 @@ def get_user_by_state(state: str) -> dict: return {} +def get_user_access_key(telegram_user_id: str) -> str: + with db.connection_context(): + try: + user = User.get(User.telegram_user_id == telegram_user_id) + return user.access_key + except DoesNotExist: + return "" + + class Location(BaseModel): fsq_id = CharField(unique=True, primary_key=True) name = CharField(max_length=128) -- cgit v1.2.3