bot: fix login flow

2 files changed, 19 insertions, 1 deletions

diff --git a/callback.py b/callback.py
index 23b814b..d968523 100644
--- a/callback.py
+++ b/callback.py
@@ -97,10 +97,11 @@ async def callback_generate_fedi_login_url(update: Update, context: ContextTypes
                                                  scopes=BOT_SCOPE,
                                                  state=state)
 
-    await update.message.reply_text(PROMPT_FEDI_LOGIN,
+    msg = await update.message.reply_text(PROMPT_FEDI_LOGIN,
                                     reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Login", url=oauth_url)]]),
                                     parse_mode=ParseMode.MARKDOWN)
 
+    context.user_data[PROMPT_FEDI_LOGIN] = msg.message_id
     return FEDI_LOGIN
 
 
diff --git a/dbstore/peewee_store.py b/dbstore/peewee_store.py
index 89e4b4b..3d37140 100644
--- a/dbstore/peewee_store.py
+++ b/dbstore/peewee_store.py
@@ -24,6 +24,23 @@ class User(BaseModel):
     toot_visibility = CharField(max_length=128, default=TOOT_VISIBILITY_PRIVATE)
 
 
+def get_user_by_state(state: str) -> dict:
+    with db.connection_context():
+        try:
+            user = User.get(User.state == state)
+            return {
+                "telegram_user_id": user.telegram_user_id,
+                "access_key": user.access_key,
+                "home_instance": user.home_instance,
+                "state": user.state,
+                "client_id": user.client_id,
+                "client_secret": user.client_secret,
+                "toot_visibility": user.toot_visibility,
+            }
+        except DoesNotExist:
+            return {}
+
+
 class Location(BaseModel):
     fsq_id = CharField(unique=True, primary_key=True)
     name = CharField(max_length=128)