1 files changed, 23 insertions, 18 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: