1 files changed, 27 insertions, 22 deletions

diff --git a/callback.py b/callback.py
index 7530afb..5fe4593 100644
--- a/callback.py
+++ b/callback.py
@@ -12,7 +12,7 @@ from config import BOT_SCOPE, ENCRYPT_KEY
 from dbstore.peewee_store import User, db, TOOT_VISIBILITY_PRIVATE, TOOT_VISIBILITY_PUBLIC, TOOT_VISIBILITY_UNLISTED
 import uuid
 from mastodon import Mastodon
-from util import decrypt
+from util import decrypt, check_user
 
 
 def generate_uuid():
@@ -118,7 +118,8 @@ async def callback_generate_fedi_login_url(update: Update, context: ContextTypes
     return FEDI_LOGIN
 
 
-async def callback_location_sharing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_location_sharing(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     if update.message.venue is not None:
         context.user_data["fsq_id"] = update.message.venue.foursquare_id
         context.user_data["title"] = update.message.venue.title
@@ -128,11 +129,10 @@ async def callback_location_sharing(update: Update, context: ContextTypes.DEFAUL
         poi = query_poi_by_fsq_id(context.user_data.get("fsq_id"))
         content = generate_toot_text(poi["name"], poi["locality"], poi["region"], poi["latitude"], poi["longitude"])
 
-        u = get_user_by_id(str(update.effective_user.id))
-        content_type = "text/markdown" if u["home_instance_type"] == "pleroma" else None
+        content_type = "text/markdown" if user["home_instance_type"] == "pleroma" else None
 
         status = get_mastodon_client(update.effective_user.id).status_post(content,
-                                                                           visibility=u["tool_visibility"],
+                                                                           visibility=user["tool_visibility"],
                                                                            content_type=content_type,
                                                                            media_ids=[])
 
@@ -174,7 +174,8 @@ async def _process_location_search(keyword, lat, lon) -> list:
     return keyboard
 
 
-async def callback_location_keyword_search(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_location_keyword_search(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     await context.bot.delete_message(update.effective_chat.id, context.user_data.get(PROMPT_LOCATION_KEYWORD))
     key = update.effective_message.text
 
@@ -191,7 +192,8 @@ async def callback_location_keyword_search(update: Update, context: ContextTypes
     return LOCATION_CONFIRMATION
 
 
-async def callback_skip_location_keyword_search(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_skip_location_keyword_search(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     query = update.callback_query
     await query.answer()
     await query.message.delete()
@@ -202,7 +204,7 @@ async def callback_skip_location_keyword_search(update: Update, context: Context
     return LOCATION_CONFIRMATION
 
 
-async def _process_location_selection(context: ContextTypes.DEFAULT_TYPE) -> int:
+async def _process_location_selection(context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     poi_name = context.user_data.get("poi_name")
     if context.user_data.get("fsq_id") is not None:
         poi = get_poi_by_fsq_id(context.user_data.get("fsq_id"))
@@ -212,11 +214,10 @@ async def _process_location_selection(context: ContextTypes.DEFAULT_TYPE) -> int
         content = generate_toot_text(poi_name, "", "", context.user_data.get("latitude"),
                                      context.user_data.get("longitude"))
 
-    u = get_user_by_id(context.user_data["user_id"])
-    content_type = "text/markdown" if u["home_instance_type"] == "pleroma" else None
+    content_type = "text/markdown" if user["home_instance_type"] == "pleroma" else None
 
     status = get_mastodon_client(context.user_data["user_id"]).status_post(content,
-                                                                           visibility=u["toot_visibility"],
+                                                                           visibility=user["toot_visibility"],
                                                                            content_type=content_type,
                                                                            media_ids=[])
 
@@ -236,7 +237,8 @@ async def _process_location_selection(context: ContextTypes.DEFAULT_TYPE) -> int
     return ADD_COMMENT
 
 
-async def callback_location_confirmation(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_location_confirmation(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     query = update.callback_query
     await query.answer()
     context.user_data["fsq_id"] = query.data
@@ -245,15 +247,16 @@ async def callback_location_confirmation(update: Update, context: ContextTypes.D
 
     context.user_data["chat_id"] = update.effective_chat.id
 
-    return await _process_location_selection(context)
+    return await _process_location_selection(context, user)
 
 
-async def callback_manual_location(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_manual_location(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     context.user_data["poi_name"] = update.effective_message.text
     context.user_data["chat_id"] = update.effective_chat.id
     context.user_data["user_id"] = update.effective_user.id
 
-    return await _process_location_selection(context)
+    return await _process_location_selection(context, user)
 
 
 async def _process_comment(context: ContextTypes.DEFAULT_TYPE) -> int:
@@ -264,14 +267,13 @@ async def _process_comment(context: ContextTypes.DEFAULT_TYPE) -> int:
     return ADD_MEDIA
 
 
-async def callback_add_comment(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_add_comment(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     context.user_data["chat_id"] = update.effective_chat.id
     context.user_data["user_id"] = update.effective_user.id
     await context.bot.delete_message(update.effective_chat.id, context.user_data.get(PROMPT_ADD_COMMENT))
 
-    with db.connection_context():
-        u = User.get(User.telegram_user_id == context.user_data["user_id"])
-        content_type = "text/markdown" if u.home_instance_type == "pleroma" else None
+    content_type = "text/markdown" if user["home_instance_type"] == "pleroma" else None
 
     comment = update.effective_message.text
     get_mastodon_client(update.effective_user.id).status_update(id=context.user_data.get(KEY_TOOT_STATUS_ID),
@@ -283,14 +285,16 @@ async def callback_add_comment(update: Update, context: ContextTypes.DEFAULT_TYP
     return await _process_comment(context)
 
 
-async def callback_skip_comment(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_skip_comment(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     context.user_data["chat_id"] = update.effective_chat.id
     await context.bot.delete_message(update.effective_chat.id, context.user_data.get(PROMPT_ADD_COMMENT))
 
     return await _process_comment(context)
 
 
-async def callback_add_media(update: Update, context: CallbackContext):
+@check_user
+async def callback_add_media(update: Update, context: CallbackContext, user: User):
     await update.message.reply_chat_action(ChatAction.TYPING)
 
     try:
@@ -342,7 +346,8 @@ async def callback_add_media(update: Update, context: CallbackContext):
         await update.message.reply_text(text=PROMPT_DONE, reply_markup=MAIN_MENU)
 
 
-async def callback_skip_media(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+@check_user
+async def callback_skip_media(update: Update, context: ContextTypes.DEFAULT_TYPE, user: User) -> int:
     query = update.callback_query
     await query.answer()