Fix webhook key-created notify / extract key view

handlers/keys/key_mode/key_cluster_mode.py

@@ -61,6 +61,150 @@ router = Router()
 moscow_tz = pytz.timezone("Europe/Moscow")
 
 
+async def send_or_edit_key_created_view(
+    session,
+    tg_id: int,
+    *,
+    key_record: dict,
+    client_id: str,
+    email: str,
+    final_link: str,
+    cluster_id: str,
+    key_name: str | None = None,
+    plan: int | None = None,
+    selected_device_limit: int | None = None,
+    selected_traffic_gb: int | None = None,
+    target_message: Message | CallbackQuery | None = None,
+) -> None:
+    if key_name is None:
+        key_name = email
+
+    target = None
+    safe_to_edit = False
+    if isinstance(target_message, CallbackQuery) and target_message.message:
+        target = target_message.message
+        safe_to_edit = True
+    elif isinstance(target_message, Message):
+        target = target_message
+        safe_to_edit = True
+
+    tg_notify = await notify_telegram_chat_id(session, tg_id)
+
+    try:
+        vless_enabled = False
+        try:
+            if plan:
+                vless_enabled = await get_vless_enabled(session, plan)
+            elif key_record.get("tariff_id"):
+                vless_enabled = await get_vless_enabled(session, key_record["tariff_id"])
+        except Exception:
+            vless_enabled = False
+
+        tv_button_enabled = bool(BUTTONS_CONFIG.get("ANDROID_TV_BUTTON_ENABLE"))
+
+        builder = InlineKeyboardBuilder()
+        if vless_enabled:
+            builder.row(
+                InlineKeyboardButton(
+                    text=ROUTER_BUTTON, callback_data=build_key_callback("connect_router", client_id, key_name)
+                )
+            )
+        else:
+            if await is_full_remnawave_cluster(cluster_id, session):
+                use_webapp = bool(MODES_CONFIG.get("REMNAWAVE_WEBAPP_ENABLED", REMNAWAVE_WEBAPP))
+                open_in_browser = bool(
+                    MODES_CONFIG.get("REMNAWAVE_WEBAPP_OPEN_IN_BROWSER", REMNAWAVE_WEBAPP_OPEN_IN_BROWSER)
+                )
+                if use_webapp and final_link:
+                    use_webapp = await process_remnawave_webapp_override(
+                        remnawave_webapp=use_webapp,
+                        final_link=final_link,
+                        session=session,
+                    )
+
+                if (
+                    use_webapp
+                    and final_link
+                    and isinstance(final_link, str)
+                    and final_link.startswith(("http://", "https://"))
+                ):
+                    if open_in_browser:
+                        builder.row(InlineKeyboardButton(text=CONNECT_DEVICE, url=final_link))
+                    else:
+                        builder.row(InlineKeyboardButton(text=CONNECT_DEVICE, web_app=WebAppInfo(url=final_link)))
+                    if tv_button_enabled:
+                        builder.row(
+                            InlineKeyboardButton(
+                                text=TV_BUTTON, callback_data=build_key_callback("connect_tv", client_id, email)
+                            )
+                        )
+                else:
+                    builder.row(
+                        InlineKeyboardButton(
+                            text=CONNECT_DEVICE, callback_data=build_key_callback("connect_device", client_id, key_name)
+                        )
+                    )
+            else:
+                builder.row(
+                    InlineKeyboardButton(
+                        text=CONNECT_DEVICE, callback_data=build_key_callback("connect_device", client_id, key_name)
+                    )
+                )
+
+        builder.row(InlineKeyboardButton(text=MY_SUB, callback_data=build_key_callback("view_key", client_id, key_name)))
+        builder.row(InlineKeyboardButton(text=SUPPORT, url=SUPPORT_CHAT_URL))
+        builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile"))
+
+        if tg_notify is not None and await process_intercept_key_creation_message(
+            chat_id=tg_notify,
+            session=session,
+            target_message=target_message,
+        ):
+            return
+
+        hook_commands = (
+            await process_key_creation_complete(
+                chat_id=tg_notify,
+                admin=False,
+                session=session,
+                email=email,
+                key_name=key_name,
+            )
+            if tg_notify is not None
+            else []
+        )
+        if hook_commands:
+            builder = insert_hook_buttons(builder, hook_commands)
+
+        key_message_text = await build_key_created_message(
+            session=session,
+            key_record=key_record,
+            final_link=final_link,
+            selected_device_limit=selected_device_limit,
+            selected_traffic_gb=selected_traffic_gb,
+        )
+
+        default_media_path = "img/pic.jpg"
+        if safe_to_edit and target is not None:
+            await edit_or_send_message(
+                target_message=target,
+                text=key_message_text,
+                reply_markup=builder.as_markup(),
+                media_path=default_media_path,
+            )
+        elif tg_notify is not None:
+            await bot.send_message(
+                chat_id=tg_notify,
+                text=key_message_text,
+                reply_markup=builder.as_markup(),
+            )
+    except Exception as e:
+        logger.error(
+            f"[Key Created View] Ошибка отправки/редактирования окна о создании ключа "
+            f"для пользователя {tg_id}: {e}"
+        )
+
+
 async def key_cluster_mode(
     tg_id: int,
     expiry_time: datetime,
@@ -216,113 +360,20 @@ async def key_cluster_mode(
             await bot.send_message(chat_id=tg_notify, text=error_message)
         return
 
-    vless_enabled = False
-    try:
-        if plan:
-            vless_enabled = await get_vless_enabled(session, plan)
-        elif key_record.get("tariff_id"):
-            vless_enabled = await get_vless_enabled(session, key_record["tariff_id"])
-    except Exception:
-        vless_enabled = False
-
-    tv_button_enabled = bool(BUTTONS_CONFIG.get("ANDROID_TV_BUTTON_ENABLE"))
-
-    builder = InlineKeyboardBuilder()
-    if vless_enabled:
-        builder.row(
-            InlineKeyboardButton(
-                text=ROUTER_BUTTON, callback_data=build_key_callback("connect_router", client_id, key_name)
-            )
-        )
-    else:
-        if await is_full_remnawave_cluster(least_loaded_cluster, session):
-            use_webapp = bool(MODES_CONFIG.get("REMNAWAVE_WEBAPP_ENABLED", REMNAWAVE_WEBAPP))
-            open_in_browser = bool(
-                MODES_CONFIG.get("REMNAWAVE_WEBAPP_OPEN_IN_BROWSER", REMNAWAVE_WEBAPP_OPEN_IN_BROWSER)
-            )
-            if use_webapp and final_link:
-                use_webapp = await process_remnawave_webapp_override(
-                    remnawave_webapp=use_webapp,
-                    final_link=final_link,
-                    session=session,
-                )
-
-            if (
-                use_webapp
-                and final_link
-                and isinstance(final_link, str)
-                and final_link.startswith(("http://", "https://"))
-            ):
-                if open_in_browser:
-                    builder.row(InlineKeyboardButton(text=CONNECT_DEVICE, url=final_link))
-                else:
-                    builder.row(InlineKeyboardButton(text=CONNECT_DEVICE, web_app=WebAppInfo(url=final_link)))
-                if tv_button_enabled:
-                    builder.row(
-                        InlineKeyboardButton(
-                            text=TV_BUTTON, callback_data=build_key_callback("connect_tv", client_id, email)
-                        )
-                    )
-            else:
-                builder.row(
-                    InlineKeyboardButton(
-                        text=CONNECT_DEVICE, callback_data=build_key_callback("connect_device", client_id, key_name)
-                    )
-                )
-        else:
-            builder.row(
-                InlineKeyboardButton(
-                    text=CONNECT_DEVICE, callback_data=build_key_callback("connect_device", client_id, key_name)
-                )
-            )
-
-    builder.row(InlineKeyboardButton(text=MY_SUB, callback_data=build_key_callback("view_key", client_id, key_name)))
-    builder.row(InlineKeyboardButton(text=SUPPORT, url=SUPPORT_CHAT_URL))
-    builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile"))
-
-    if tg_notify is not None and await process_intercept_key_creation_message(
-        chat_id=tg_notify,
-        session=session,
-        target_message=message_or_query,
-    ):
-        return
-
-    hook_commands = (
-        await process_key_creation_complete(
-            chat_id=tg_notify,
-            admin=False,
-            session=session,
-            email=email,
-            key_name=key_name,
-        )
-        if tg_notify is not None
-        else []
-    )
-    if hook_commands:
-        builder = insert_hook_buttons(builder, hook_commands)
-
-    key_message_text = await build_key_created_message(
+    await send_or_edit_key_created_view(
         session=session,
+        tg_id=tg_id,
         key_record=key_record,
+        client_id=client_id,
+        email=email,
+        key_name=key_name,
         final_link=final_link,
+        cluster_id=least_loaded_cluster,
+        plan=plan,
         selected_device_limit=selected_device_limit,
         selected_traffic_gb=selected_traffic_gb,
+        target_message=message_or_query,
     )
 
-    default_media_path = "img/pic.jpg"
-    if safe_to_edit:
-        await edit_or_send_message(
-            target_message=target_message,
-            text=key_message_text,
-            reply_markup=builder.as_markup(),
-            media_path=default_media_path,
-        )
-    elif tg_notify is not None:
-        await bot.send_message(
-            chat_id=tg_notify,
-            text=key_message_text,
-            reply_markup=builder.as_markup(),
-        )
-
     if state:
         await state.clear()