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)