Skip to content

Commit

Permalink
feat: add users inbounds manage (#88)
Browse files Browse the repository at this point in the history
* feat: add users inbounds edit

* fix: check deleted inbounds

* feat: bump to 0.2.1
  • Loading branch information
erfjab authored Oct 31, 2024
1 parent 86fea6b commit 4de2869
Show file tree
Hide file tree
Showing 13 changed files with 285 additions and 28 deletions.
25 changes: 17 additions & 8 deletions db/alembic/versions/ab1ce3ef2a57_add_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
4 changes: 1 addition & 3 deletions db/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions models/callback.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand All @@ -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"):
Expand Down
3 changes: 2 additions & 1 deletion routers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
15 changes: 11 additions & 4 deletions routers/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
14 changes: 11 additions & 3 deletions routers/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand All @@ -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(
Expand All @@ -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):
Expand Down
66 changes: 66 additions & 0 deletions routers/users.py
Original file line number Diff line number Diff line change
@@ -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(),
)
2 changes: 1 addition & 1 deletion utils/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
]
]
95 changes: 95 additions & 0 deletions utils/helpers.py
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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
Loading

0 comments on commit 4de2869

Please sign in to comment.