src.handlers.wg_service

  1"""Действия с конфигурациями WireGuard"""
  2
  3import logging
  4import os
  5from contextlib import suppress
  6from typing import Union
  7
  8from aiogram import Bot, F, Router
  9from aiogram.exceptions import TelegramBadRequest
 10from aiogram.filters.command import Command
 11from aiogram.types import CallbackQuery, FSInputFile, Message
 12from pytils.numeral import get_plural
 13
 14import text
 15from core import exceptions as exc
 16from core.config import settings
 17from core.err import bot_except
 18from core.metric import async_speed_metric
 19from db import utils
 20from db.models import FreezeSteps, UserActivity, UserData, WgConfig
 21from handlers.utils import find_config, find_user
 22from kb import get_config_keyboard, remove_config, static_pay_button, why_freezed_button
 23from wg.utils import WgServerTools
 24
 25logger = logging.getLogger()
 26router = Router()
 27router.message.filter(F.chat.type == "private")
 28
 29
 30@router.message(Command("me"))
 31@router.message(Command("config"))
 32@router.message(F.text.lower().in_(text.me))
 33@router.callback_query(F.data == "user_configurations")
 34@async_speed_metric
 35@bot_except
 36async def post_user_data(trigger: Union[Message, CallbackQuery]):
 37    """Отправляет данные о конфигурациях пользователя.
 38
 39    Args:
 40        trigger (Union[Message, CallbackQuery]): Сообщение или событие обратного вызова, инициировавшее команду.
 41
 42    Если у пользователя есть конфигурации, отправляет их список. Если конфигурации заморожены, отправляет сообщение об этом.
 43    Также информирует пользователя о максимальном количестве конфигураций для его тарифа.
 44    """
 45    user_data: UserData = await find_user(trigger, configs=True)
 46    if not user_data:
 47        return
 48    else:
 49        create_cfg_btn, create_output_cfg_btn = get_config_keyboard()
 50
 51        if user_data.configs:
 52            await getattr(trigger, "message", trigger).answer("Ваши конфигурации:")
 53
 54            user_data.configs.sort(key=lambda conf: conf.id)
 55            for i, config in enumerate(user_data.configs, 1):
 56                if config.user_private_key:
 57                    config_text = f"({i}/{settings.acceptable_config[user_data.stage]}) - Name: {config.name} | id: {config.user_private_key[:4]}"
 58
 59                    if config.freeze != FreezeSteps.no:
 60                        config_text = f'🥶<b>FREEZED</b>❄️\n<span class="tg-spoiler">{config_text}</span>'
 61
 62                        await getattr(trigger, "message", trigger).answer(
 63                            config_text, reply_markup=why_freezed_button
 64                        )
 65                    else:
 66                        await getattr(trigger, "message", trigger).answer(
 67                            config_text, reply_markup=create_output_cfg_btn
 68                        )
 69
 70        if user_data.active == UserActivity.active:
 71            cfg_number = settings.acceptable_config[user_data.stage] - len(
 72                user_data.configs
 73            )
 74
 75            if cfg_number <= 0:
 76                await getattr(trigger, "message", trigger).answer(
 77                    "Достигнуто максимальное количество конфигураций для данного тарифа.",
 78                    reply_markup=static_pay_button,
 79                )
 80            else:
 81                await getattr(trigger, "message", trigger).answer(
 82                    f"Вы можете создать еще {get_plural(cfg_number, 'конфигурацию, конфигурации, конфигураций')}",
 83                    reply_markup=create_cfg_btn,
 84                )
 85        elif user_data.active == UserActivity.inactive:
 86            await getattr(trigger, "message", trigger).answer(
 87                text.UNPAY, reply_markup=static_pay_button
 88            )
 89
 90
 91@router.message(Command("create"))
 92@router.callback_query(F.data == "create_configuration")
 93@async_speed_metric
 94@bot_except
 95async def create_config_data(trigger: Union[Message, CallbackQuery], bot: Bot):
 96    """Создает новую конфигурацию для пользователя.
 97
 98    Args:
 99        trigger (Union[Message, CallbackQuery]): Сообщение или событие обратного вызова, инициировавшее команду.
100        bot (Bot): Экземпляр бота для выполнения действий.
101
102    Проверяет, может ли пользователь создать новую конфигурацию, и создает ее, если это возможно.
103    В случае ошибок отправляет соответствующие сообщения.
104    """
105    try:
106        user_data: UserData = await find_user(trigger, configs=True)
107        if not user_data:
108            return
109        elif user_data.active != UserActivity.active:
110            raise exc.PayError
111
112        elif len(user_data.configs) < settings.acceptable_config[user_data.stage]:
113            wg = WgServerTools()
114            conf = await wg.move_user(move="add", user_id=trigger.from_user.id)
115            config: WgConfig = await utils.add_wg_config(
116                conf, user_id=trigger.from_user.id
117            )
118
119        else:
120            raise exc.StagePayError
121
122    except exc.DatabaseError:
123        await trigger.answer(text=text.DB_ERROR, show_alert=True)
124    except exc.WireguardError:
125        await trigger.answer(text=text.WG_ERROR, show_alert=True)
126    except exc.StagePayError:
127        await getattr(trigger, "message", trigger).answer(
128            "Достигнуто максимальное количество конфигураций для данного тарифа.",
129            reply_markup=static_pay_button,
130        )
131        await bot.delete_message(
132            trigger.from_user.id, getattr(trigger, "message", trigger).message_id
133        )
134    except exc.PayError:
135        await getattr(trigger, "message", trigger).answer(
136            text.UNPAY, reply_markup=static_pay_button
137        )
138    else:
139        create_cfg_btn, create_output_cfg_btn = get_config_keyboard()
140
141        await trigger.answer(text="Конфигурация успешно создана", show_alert=True)
142        await getattr(trigger, "message", trigger).answer(
143            f"Конфигурация: {config.name} | id: {config.user_private_key[:4]}",
144            reply_markup=create_output_cfg_btn,
145        )
146
147        with suppress(TelegramBadRequest):
148            cfg_number = get_plural(
149                settings.acceptable_config[user_data.stage]
150                - len(user_data.configs)
151                - 1,
152                "конфигурацию, конфигурации, конфигураций",
153            )
154            await getattr(trigger, "message", trigger).edit_text(
155                f"Вы можете создать еще {cfg_number}", reply_markup=create_cfg_btn
156            )
157
158        if cfg_number.startswith("0"):
159            await bot.delete_message(
160                trigger.from_user.id, getattr(trigger, "message", trigger).message_id
161            )
162
163
164@router.callback_query(F.data == "remove_config_confirm")
165@bot_except
166async def remove_config_data_confirm(callback: CallbackQuery):
167    """Подтверждение удаления конфигурации.
168
169    Args:
170        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
171    """
172
173    with suppress(TelegramBadRequest):
174        await callback.message.edit_text(
175            "После удаления конфигурации вы больше не сможете подключаться по ней к VPN-серверу. "
176            "Для восстановления подключения вам необходимо будет создать новую конфигурацию и заменить ей старую в vpn-приложении"
177            f"\n\n? Вы уверены, что хотите удалить конфигурацию {callback.message.text}",
178            reply_markup=remove_config(),
179        )
180
181
182@router.callback_query(F.data == "remove_config_cancel")
183@bot_except
184async def remove_config_data_cancel(callback: CallbackQuery):
185    """Отмена удаления конфигурации.
186
187    Args:
188        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
189    """
190    _, create_output_cfg_btn = get_config_keyboard()
191
192    *_, old_cfg_message = callback.message.text.split("удалить конфигурацию ")
193
194    with suppress(TelegramBadRequest):
195        await callback.message.edit_text(
196            old_cfg_message, reply_markup=create_output_cfg_btn
197        )
198
199
200@router.callback_query(F.data == "remove_config")
201@async_speed_metric
202@bot_except
203async def remove_config_data(callback: CallbackQuery, bot: Bot):
204    """Удаляет конфигурацию с WG сервера, а затем из БД.
205
206    Args:
207        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
208    """
209
210    try:
211        user_config: WgConfig = await find_config(callback)
212        if not user_config:
213            return
214
215        wg = WgServerTools()
216        await wg.move_user(move="del", server_pubkey=user_config.server_public_key)
217
218        await utils.delete_wg_config(user_config)
219
220    except exc.DatabaseError:
221        await callback.answer(text=text.DB_ERROR, show_alert=True)
222    except exc.WireguardError:
223        await callback.answer(text=text.WG_ERROR, show_alert=True)
224    else:
225        await callback.answer(text="Конфигурация успешно удалена", show_alert=True)
226    finally:
227        await bot.delete_message(callback.from_user.id, callback.message.message_id)
228
229
230@router.callback_query(F.data == "create_conf_text")
231@async_speed_metric
232@bot_except
233async def get_config_text(callback: CallbackQuery, bot: Bot):
234    """Отправляет текст конфигурации пользователю.
235
236    Args:
237        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
238
239    Отправляет текст конфигурации, связанный с пользователем, в формате, удобном для чтения.
240    """
241    user_config: WgConfig = await find_config(callback)
242    if not user_config:
243        await bot.delete_message(callback.from_user.id, callback.message.message_id)
244        return
245
246    config = text.get_config_data(user_config)
247
248    await callback.message.answer("Конфигурация " + callback.message.text)
249    await callback.message.answer(f"<pre>{config}</pre>")
250
251
252@router.callback_query(F.data == "create_conf_file")
253@async_speed_metric
254@bot_except
255async def get_config_file(callback: CallbackQuery, bot: Bot):
256    """Отправляет файл конфигурации пользователю.
257
258    Args:
259        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
260
261    Создает файл конфигурации и отправляет его пользователю в виде документа.
262    """
263    user_config: WgConfig = await find_config(callback)
264    if not user_config:
265        await bot.delete_message(callback.from_user.id, callback.message.message_id)
266        return
267
268    config = text.get_config_data(user_config)
269    config_file = await text.create_config_file(config)
270
271    await callback.message.answer("Конфигурация " + callback.message.text)
272    await callback.message.answer_document(FSInputFile(config_file, "myVpn.conf"))
273    await callback.message.answer(
274        "‼️ Если у вас возникает ошибка: <b>Неправильное имя</b> "
275        "Проверьте имя файла: в нем не должно быть пробелов или спецсимволов (подчеркиваний, скобок ...), "
276        "длина имени файла должна быть не более 10 символов"
277    )
278    os.remove(config_file)
279
280
281@router.callback_query(F.data == "create_conf_qr")
282@async_speed_metric
283@bot_except
284async def get_config_qr(callback: CallbackQuery, bot: Bot):
285    """Отправляет QR-код конфигурации пользователю.
286
287    Args:
288        callback (CallbackQuery): Событие обратного вызова, инициировавшее команду.
289
290    Создает QR-код конфигурации и отправляет его пользователю в виде изображения.
291    """
292    user_config: WgConfig = await find_config(callback)
293    if not user_config:
294        await bot.delete_message(callback.from_user.id, callback.message.message_id)
295        return
296
297    config = text.get_config_data(user_config)
298    config_qr = text.create_config_qr(config)
299
300    await callback.message.answer("Конфигурация " + callback.message.text)
301    await callback.message.answer_photo(
302        FSInputFile(config_qr, f"{user_config.name}_wg.conf")
303    )
304    os.remove(config_qr)