feat: add users inbounds manage

db/alembic/versions/ab1ce3ef2a57_add_settings.py

@@ -5,27 +5,36 @@
 Create Date: 2024-10-13 01:42:55.733416
 
 """
+
 from typing import Sequence, Union
 
 from alembic import op
 import sqlalchemy as sa
 
 
 # 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')
+    op.drop_table("settings")

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

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"):

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

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)
+    await node_monitoring_menu(callback)

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):

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(),
+    )

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()
-]
+]

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