From 4de28695d8192febe9943167c93c14d7000ebd71 Mon Sep 17 00:00:00 2001 From: Erfan <143827987+erfjab@users.noreply.github.com> Date: Thu, 31 Oct 2024 16:53:27 +0330 Subject: [PATCH] feat: add users inbounds manage (#88) * feat: add users inbounds edit * fix: check deleted inbounds * feat: bump to 0.2.1 --- .../versions/ab1ce3ef2a57_add_settings.py | 25 +++-- db/models.py | 4 +- models/callback.py | 3 + routers/__init__.py | 3 +- routers/node.py | 15 ++- routers/user.py | 14 ++- routers/users.py | 66 +++++++++++++ utils/config.py | 2 +- utils/helpers.py | 95 +++++++++++++++++++ utils/keys.py | 35 ++++++- utils/lang.py | 10 +- utils/log.py | 2 +- utils/panel.py | 39 +++++++- 13 files changed, 285 insertions(+), 28 deletions(-) create mode 100644 routers/users.py diff --git a/db/alembic/versions/ab1ce3ef2a57_add_settings.py b/db/alembic/versions/ab1ce3ef2a57_add_settings.py index 844c2d0..feedee6 100644 --- a/db/alembic/versions/ab1ce3ef2a57_add_settings.py +++ b/db/alembic/versions/ab1ce3ef2a57_add_settings.py @@ -5,6 +5,7 @@ Create Date: 2024-10-13 01:42:55.733416 """ + from typing import Sequence, Union from alembic import op @@ -12,20 +13,28 @@ # revision identifiers, used by Alembic. -revision: str = 'ab1ce3ef2a57' -down_revision: Union[str, None] = '3e5deef43bf0' +revision: str = "ab1ce3ef2a57" +down_revision: Union[str, None] = "3e5deef43bf0" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None + def upgrade(): # Create settings table - op.create_table('settings', - sa.Column('key', sa.String(256), primary_key=True), - sa.Column('value', sa.String(2048), nullable=True), - sa.Column('created_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), - sa.Column('updated_at', sa.DateTime(), nullable=True), + op.create_table( + "settings", + sa.Column("key", sa.String(256), primary_key=True), + sa.Column("value", sa.String(2048), nullable=True), + sa.Column( + "created_at", + sa.DateTime(), + server_default=sa.func.current_timestamp(), + nullable=False, + ), + sa.Column("updated_at", sa.DateTime(), nullable=True), ) + def downgrade(): # Drop settings table - op.drop_table('settings') \ No newline at end of file + op.drop_table("settings") diff --git a/db/models.py b/db/models.py index 833cd07..62af69f 100644 --- a/db/models.py +++ b/db/models.py @@ -20,9 +20,7 @@ class Token(Base): class Setting(Base): __tablename__ = "settings" - key: Mapped[str] = mapped_column( - String(256), primary_key=True - ) + key: Mapped[str] = mapped_column(String(256), primary_key=True) value: Mapped[str] = mapped_column(String(2048)) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=func.now(), nullable=False diff --git a/models/callback.py b/models/callback.py index 45a189f..5faeac1 100644 --- a/models/callback.py +++ b/models/callback.py @@ -12,12 +12,14 @@ class AdminActions(str, Enum): class BotActions(str, Enum): NodeChecker = "node_checker" NodeAutoRestart = "node_auto_restart" + UsersInbound = "users_inbound" class PagesActions(str, Enum): Home = "home" UserCreate = "user_create" NodeMonitoring = "node_monitoring" + UsersMenu = "users_menu" class PagesCallbacks(CallbackData, prefix="pages"): @@ -41,6 +43,7 @@ class UserInboundsCallbacks(CallbackData, prefix="user_inbounds"): is_selected: bool | None = None action: AdminActions is_done: bool = False + just_one_inbound: bool = False class AdminSelectCallbacks(CallbackData, prefix="admin_select"): diff --git a/routers/__init__.py b/routers/__init__.py index 4e986a8..54dad37 100644 --- a/routers/__init__.py +++ b/routers/__init__.py @@ -3,12 +3,13 @@ def setup_routers() -> Router: - from . import base, user, node + from . import base, user, node, users router = Router() router.include_router(base.router) router.include_router(user.router) router.include_router(node.router) + router.include_router(users.router) return router diff --git a/routers/node.py b/routers/node.py index d052766..066df10 100644 --- a/routers/node.py +++ b/routers/node.py @@ -10,24 +10,29 @@ SettingKeys, SettingUpsert, ConfirmCallbacks, - BotActions + BotActions, ) router = Router() + async def get_setting_status(key: SettingKeys) -> str: return "ON" if await SettingManager.get(key) else "OFF" + async def toggle_setting(key: SettingKeys): current_value = await SettingManager.get(key) new_value = None if current_value else "True" await SettingManager.upsert(SettingUpsert(key=key, value=new_value)) + @router.callback_query(PagesCallbacks.filter(F.page.is_(PagesActions.NodeMonitoring))) async def node_monitoring_menu(callback: CallbackQuery): checker_status = await get_setting_status(SettingKeys.NodeMonitoringIsActive) - auto_restart_status = await get_setting_status(SettingKeys.NodeMonitoringAutoRestart) - + auto_restart_status = await get_setting_status( + SettingKeys.NodeMonitoringAutoRestart + ) + text = MessageTexts.NodeMonitoringMenu.format( checker=checker_status, auto_restart=auto_restart_status, @@ -36,12 +41,14 @@ async def node_monitoring_menu(callback: CallbackQuery): text=text, reply_markup=BotKeyboards.node_monitoring() ) + @router.callback_query(ConfirmCallbacks.filter(F.page.is_(BotActions.NodeAutoRestart))) async def node_monitoring_auto_restart(callback: CallbackQuery): await toggle_setting(SettingKeys.NodeMonitoringAutoRestart) await node_monitoring_menu(callback) + @router.callback_query(ConfirmCallbacks.filter(F.page.is_(BotActions.NodeChecker))) async def node_monitoring_checker(callback: CallbackQuery): await toggle_setting(SettingKeys.NodeMonitoringIsActive) - await node_monitoring_menu(callback) \ No newline at end of file + await node_monitoring_menu(callback) diff --git a/routers/user.py b/routers/user.py index 1c205be..a2c95c4 100644 --- a/routers/user.py +++ b/routers/user.py @@ -124,7 +124,7 @@ async def user_create_status( ) -@router.callback_query(AdminSelectCallbacks.filter()) +@router.callback_query(AdminSelectCallbacks.filter(F.just_one_inbound.is_(False))) async def user_create_owner_select( callback: CallbackQuery, callback_data: AdminSelectCallbacks, state: FSMContext ): @@ -139,7 +139,11 @@ async def user_create_owner_select( @router.callback_query( UserInboundsCallbacks.filter( - (F.action.is_(AdminActions.Add) & (F.is_done.is_(False))) + ( + F.action.is_(AdminActions.Add) + & (F.is_done.is_(False)) + & (F.just_one_inbound.is_(False)) + ) ) ) async def user_create_inbounds( @@ -163,7 +167,11 @@ async def user_create_inbounds( @router.callback_query( UserInboundsCallbacks.filter( - (F.action.is_(AdminActions.Add) & (F.is_done.is_(True))) + ( + F.action.is_(AdminActions.Add) + & (F.is_done.is_(True)) + & (F.just_one_inbound.is_(False)) + ) ) ) async def user_create_inbounds_save(callback: CallbackQuery, state: FSMContext): diff --git a/routers/users.py b/routers/users.py new file mode 100644 index 0000000..e439d41 --- /dev/null +++ b/routers/users.py @@ -0,0 +1,66 @@ +from aiogram import Router, F +from aiogram.types import CallbackQuery +from models import ( + PagesActions, + PagesCallbacks, + AdminActions, + ConfirmCallbacks, + BotActions, + UserInboundsCallbacks, +) +from utils.lang import MessageTexts +from utils.keys import BotKeyboards +from utils import panel, helpers + +router = Router() + + +@router.callback_query(PagesCallbacks.filter(F.page == PagesActions.UsersMenu)) +async def menu(callback: CallbackQuery): + return await callback.message.edit_text( + text=MessageTexts.UsersMenu, reply_markup=BotKeyboards.users() + ) + + +@router.callback_query(ConfirmCallbacks.filter(F.page == BotActions.UsersInbound)) +async def inbound_add(callback: CallbackQuery, callback_data: ConfirmCallbacks): + inbounds = await panel.inbounds() + return await callback.message.edit_text( + text=MessageTexts.UsersInboundSelect, + reply_markup=BotKeyboards.inbounds( + inbounds=inbounds, action=callback_data.action, just_one_inbound=True + ), + ) + + +@router.callback_query( + UserInboundsCallbacks.filter( + ( + F.action.in_([AdminActions.Add, AdminActions.Delete]) + & (F.is_done.is_(True)) + & (F.just_one_inbound.is_(True)) + ) + ) +) +async def inbound_confirm( + callback: CallbackQuery, callback_data: UserInboundsCallbacks +): + working_message = await callback.message.edit_text(text=MessageTexts.Working) + result = await helpers.manage_panel_inbounds( + callback_data.tag, + callback_data.protocol, + ( + AdminActions.Add + if callback_data.action.value == AdminActions.Add.value + else AdminActions.Delete + ), + ) + + return await working_message.edit_text( + text=( + MessageTexts.UsersInboundSuccessUpdated + if result + else MessageTexts.UsersInboundErrorUpdated + ), + reply_markup=BotKeyboards.home(), + ) diff --git a/utils/config.py b/utils/config.py index 0765a94..7bd5736 100644 --- a/utils/config.py +++ b/utils/config.py @@ -36,4 +36,4 @@ def require_setting(setting_name, value): x.strip() for x in config("EXCLUDED_MONITORINGS", default="", cast=str).split(",") if x.strip() -] \ No newline at end of file +] diff --git a/utils/helpers.py b/utils/helpers.py index 95bf373..089f77c 100644 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -1,5 +1,10 @@ import qrcode +import asyncio from io import BytesIO +from models import AdminActions +from utils import panel +from utils.log import logger +from marzban import UserModify, UserResponse async def create_qr(text: str) -> bytes: @@ -20,3 +25,93 @@ async def create_qr(text: str) -> bytes: img_bytes_io = BytesIO() qr_img.save(img_bytes_io, "PNG") return img_bytes_io.getvalue() + + +async def process_user( + semaphore: asyncio.Semaphore, + user: UserResponse, + tag: str, + protocol: str, + action: AdminActions, + max_retries: int = 3, +) -> bool: + """Process a single user with semaphore for rate limiting and retry mechanism""" + async with semaphore: + current_inbounds = user.inbounds.copy() if user.inbounds else {} + current_proxies = user.proxies.copy() if user.proxies else {} + + needs_update = False + + if action == AdminActions.Delete: + if protocol in current_inbounds and tag in current_inbounds[protocol]: + current_inbounds[protocol].remove(tag) + needs_update = True + + if protocol in current_inbounds and not current_inbounds[protocol]: + current_inbounds.pop(protocol, None) + current_proxies.pop(protocol, None) + + elif action == AdminActions.Add: + if protocol not in current_inbounds: + current_inbounds[protocol] = [] + current_proxies[protocol] = {} + needs_update = True + + if tag not in current_inbounds.get(protocol, []): + if protocol not in current_inbounds: + current_inbounds[protocol] = [] + current_inbounds[protocol].append(tag) + needs_update = True + + if not needs_update: + return True + + update_data = UserModify( + proxies=current_proxies, + inbounds=current_inbounds, + ) + + success = await panel.user_modify(user.username, update_data) + + if success: + return True + + +async def process_batch( + users: list[UserResponse], tag: str, protocol: str, action: AdminActions +) -> int: + """Process a batch of users concurrently with rate limiting""" + semaphore = asyncio.Semaphore(5) + tasks = [] + + for user in users: + task = asyncio.create_task(process_user(semaphore, user, tag, protocol, action)) + tasks.append(task) + + results = await asyncio.gather(*tasks) + return sum(results) + + +async def manage_panel_inbounds(tag: str, protocol: str, action: AdminActions) -> bool: + try: + offset = 0 + batch_size = 25 + + while True: + users = await panel.get_users(offset) + if not users: + break + + await process_batch(users, tag, protocol, action) + + if len(users) < batch_size: + break + offset += batch_size + + await asyncio.sleep(1.0) + + return True + + except Exception as e: + logger.error(f"Error in manage panel inbounds: {e}") + return False diff --git a/utils/keys.py b/utils/keys.py index 9b59e7b..1509600 100644 --- a/utils/keys.py +++ b/utils/keys.py @@ -29,7 +29,11 @@ def home() -> InlineKeyboardMarkup: text=KeyboardTexts.NodeMonitoring, callback_data=PagesCallbacks(page=PagesActions.NodeMonitoring).pack(), ) - return kb.as_markup() + kb.button( + text=KeyboardTexts.UsersMenu, + callback_data=PagesCallbacks(page=PagesActions.UsersMenu).pack(), + ) + return kb.adjust(2).as_markup() @staticmethod def cancel() -> InlineKeyboardMarkup: @@ -72,18 +76,21 @@ def inbounds( inbounds: dict[str, list[ProxyInbound]], selected: set[str] = [], action: AdminActions = AdminActions.Add, + just_one_inbound: bool = False, ): kb = InlineKeyboardBuilder() for protocol_list in inbounds.values(): for inbound in protocol_list: is_selected = inbound["tag"] in selected kb.button( - text=f"{'โœ…' if is_selected else 'โŒ'} {inbound['tag']} ({inbound['protocol']})", + text=f"{('โœ…' if is_selected else 'โŒ') if not just_one_inbound else '๐Ÿ”˜'} {inbound['tag']} ({inbound['protocol']})", callback_data=UserInboundsCallbacks( tag=inbound["tag"], protocol=inbound["protocol"], is_selected=is_selected, action=action, + just_one_inbound=just_one_inbound, + is_done=just_one_inbound, ), ) kb.row( @@ -141,3 +148,27 @@ def node_monitoring() -> InlineKeyboardMarkup: ), ) return kb.adjust(2).as_markup() + + @staticmethod + def users() -> InlineKeyboardMarkup: + kb = InlineKeyboardBuilder() + + kb.button( + text=KeyboardTexts.UsersAddInbound, + callback_data=ConfirmCallbacks( + page=BotActions.UsersInbound, action=AdminActions.Add + ), + ) + kb.button( + text=KeyboardTexts.UsersDeleteInbound, + callback_data=ConfirmCallbacks( + page=BotActions.UsersInbound, action=AdminActions.Delete + ), + ) + kb.row( + InlineKeyboardButton( + text=KeyboardTexts.Home, + callback_data=PagesCallbacks(page=PagesActions.Home).pack(), + ), + ) + return kb.adjust(2).as_markup() diff --git a/utils/lang.py b/utils/lang.py index 0bc6fd7..947b362 100644 --- a/utils/lang.py +++ b/utils/lang.py @@ -1,6 +1,6 @@ from enum import Enum -VERSION = "0.2.0" +VERSION = "0.2.1" OWNER = "@ErfJabs" @@ -13,6 +13,9 @@ class KeyboardTexts(str, Enum): Finish = "โœ”๏ธ Finish" NodeMonitoringChecker = "๐Ÿงจ Checker" NodeMonitoringAutoRestart = "๐Ÿ” AutoRestart" + UsersMenu = "๐Ÿ‘ฅ Users" + UsersAddInbound = "โž• Add inbound" + UsersDeleteInbound = "โž– Delete inbound" class MessageTexts(str, Enum): @@ -45,3 +48,8 @@ class MessageTexts(str, Enum): "๐Ÿงจ Checker is {checker}\n" "๐Ÿ” AutoRestart is {auto_restart}" ) + UsersMenu = "๐Ÿ‘ฅ What do you need?" + UsersInboundSelect = "๐ŸŒ Select Your Inbound:" + Working = "โณ" + UsersInboundSuccessUpdated = "โœ… Users Inbounds is Updated!" + UsersInboundErrorUpdated = "โŒ Users Inbounds not Updated!" diff --git a/utils/log.py b/utils/log.py index 5d74612..bed7531 100644 --- a/utils/log.py +++ b/utils/log.py @@ -23,4 +23,4 @@ def setup_logger(bot_name, level=logging.INFO): return logger -logger = setup_logger("HamedBot") +logger = setup_logger("HolderBot") diff --git a/utils/panel.py b/utils/panel.py index 30990f1..ae79524 100644 --- a/utils/panel.py +++ b/utils/panel.py @@ -1,4 +1,11 @@ -from marzban import MarzbanAPI, ProxyInbound, UserResponse, UserCreate, Admin +from marzban import ( + MarzbanAPI, + ProxyInbound, + UserResponse, + UserCreate, + Admin, + UserModify, +) from datetime import datetime, timedelta from utils.config import MARZBAN_ADDRESS from db import TokenManager @@ -13,7 +20,7 @@ async def inbounds() -> dict[str, list[ProxyInbound]]: inbounds = await marzban_panel.get_inbounds(get_token.token) return inbounds or False except Exception as e: - logger.error(f"Error getting token: {e}") + logger.error(f"Error getting panel inbounds: {e}") return False @@ -61,7 +68,7 @@ async def admins() -> list[Admin]: admins = await marzban_panel.get_admins(get_token.token) return admins or False except Exception as e: - logger.error(f"Error getting token: {e}") + logger.error(f"Error getting admins list: {e}") return False @@ -73,5 +80,29 @@ async def set_owner(admin: str, user: str) -> bool: ) return user or False except Exception as e: - logger.error(f"Error getting token: {e}") + logger.error(f"Error set owner: {e}") + return False + + +async def user_modify(username: str, data: UserModify) -> bool: + try: + get_token = await TokenManager.get() + user = await marzban_panel.modify_user( + username=username, user=data, token=get_token.token + ) + return True if user else False + except Exception as e: + logger.error(f"Error user modify: {e}") + return False + + +async def get_users(offset: int = 0) -> list[UserResponse]: + try: + get_token = await TokenManager.get() + users = await marzban_panel.get_users( + token=get_token.token, offset=offset, limit=offset + ) + return users.users if users else False + except Exception as e: + logger.error(f"Error getting all users: {e}") return False