feat: add user charge in templates without reset (#173)

Author: erfjab

app/api/_clinet.py

@@ -30,7 +30,7 @@ async def generate_access(
             case ServerTypes.MARZNESHIN:
                 api = MarzneshinApiManager(host=data["host"])
                 token = await api.get_token(data["username"], data["password"])
-                token = token.access_token if token and token.is_sudo is True else None
+                token = token.access_token if token and token.is_sudo is True else False
 
             case ServerTypes.MARZBAN:
                 api = MarzbanApiManager(host=data["host"])

app/keys/_enums.py

@@ -22,5 +22,7 @@ class Actions(str, Enum):
 
 class YesOrNot(str, Enum):
     YES_USAGE = "✅ Yes (reset usage)"
+    YES_NORMAL = "✅ Yes (no reset/charge)"
+    YES_CHARGE = "✅ Yes (charge)"
     YES = "✅ Yes"
     NO = "❌ No"

app/routers/users/modify/charge.py

@@ -48,7 +48,7 @@ async def chargestart(
     await state.set_state(UserModifyForm.TEMPLATE)
     await state.update_data(username=callback_data.extra)
     return await callback.message.edit_text(
-        text=MessageTexts.ASK_ADMIN,
+        text=MessageTexts.ITEMS,
         reply_markup=BotKeys.selector(
             data=[(tem.button_remark, tem.id) for tem in templates],
             types=Pages.USERS,
@@ -79,11 +79,17 @@ async def chargeend(
     return await callback.message.edit_text(
         text=MessageTexts.ASK_SURE,
         reply_markup=BotKeys.selector(
-            data=[YesOrNot.YES_USAGE, YesOrNot.YES, YesOrNot.NO],
+            data=[
+                YesOrNot.YES_USAGE,
+                YesOrNot.YES_CHARGE,
+                YesOrNot.YES_NORMAL,
+                YesOrNot.NO,
+            ],
             types=Pages.USERS,
             action=Actions.MODIFY,
             extra=callback_data.extra,
             panel=server.id,
+            width=1,
         ),
     )
 
@@ -99,6 +105,7 @@ async def chargeend(
                     YesOrNot.YES_USAGE,
                     YesOrNot.YES,
                     YesOrNot.NO,
+                    YesOrNot.YES_CHARGE,
                 ]
             )
         )
@@ -128,17 +135,26 @@ async def confirmend(
         )
         return await tracker.add(track)
 
+    user = await ClinetManager.get_user(server, data["username"])
+    if not user:
+        track = await callback.message.edit_text(
+            text=MessageTexts.NOT_FOUND, reply_markup=BotKeys.cancel()
+        )
+        return await tracker.add(track)
+
     if callback_data.select == YesOrNot.YES_USAGE.value:
         await ClinetManager.reset_user(server, data["username"])
 
     datadict = charge_user_data(
         server.types,
-        data["username"],
+        user,
         template.data_limit,
         template.date_limit,
         template.date_types,
+        charge=True if callback_data.select == YesOrNot.YES_CHARGE else False,
     )
     action = await ClinetManager.modify_user(server, data["username"], datadict)
+
     return await callback.message.edit_text(
         text=MessageTexts.SUCCESS if action else MessageTexts.FAILED,
         reply_markup=BotKeys.cancel(),

app/settings/utils/user.py

@@ -9,7 +9,8 @@
     MarzneshinUserModify,
 )
 from app.models.server import ServerTypes
-from app.api.types.marzban import MarzbanProxyInbound
+from app.api.types.marzban import MarzbanProxyInbound, MarzbanUserResponse
+from app.api.types.marzneshin import MarzneshinUserResponse
 
 
 def _get_expire_strategy(types: ServerTypes, datetype: str):
@@ -107,7 +108,68 @@ def user_create_data(
 
 
 def charge_user_data(
-    types: ServerTypes, username: str, datalimit: int, datelimit: int, status: str
+    types: ServerTypes,
+    user: MarzbanUserResponse | MarzneshinUserResponse,
+    datalimit: int,
+    datelimit: int,
+    datetypes: str,
+    charge: bool = False,
+) -> dict:
+    if charge:
+        return advenced_charge_user_data(types, user, datalimit, datelimit, datetypes)
+    return normal_charge_user_data(types, user, datalimit, datelimit, datetypes)
+
+
+def advenced_charge_user_data(
+    types: ServerTypes,
+    user: MarzbanUserResponse | MarzneshinUserResponse,
+    datalimit: int,
+    datelimit: int,
+    datetypes: str,
+) -> dict:
+    match types:
+        case ServerTypes.MARZBAN:
+            data = MarzbanUserModify(
+                data_limit=int(datalimit) * (1024**3) + user.data_limit
+                if user.data_limit
+                else 0,
+                status="on_hold"
+                if datetypes == DateTypes.AFTER_FIRST_USE
+                else "active",
+                expire=int(int((datelimit) * (24 * 60 * 60)) + (user.expire))
+                if datetypes != DateTypes.AFTER_FIRST_USE
+                else None,
+                on_hold_expire_duration=int(
+                    int((datelimit) * (24 * 60 * 60)) + (user.on_hold_expire_duration)
+                )
+                if datetypes == DateTypes.AFTER_FIRST_USE
+                else None,
+            ).dict()
+        case ServerTypes.MARZNESHIN:
+            expire_strategy = _get_expire_strategy(types, datetypes)
+            data = MarzneshinUserModify(
+                username=user.username,
+                data_limit=int(datalimit) * (1024**3) + user.data_limit
+                if user.data_limit
+                else 0,
+                expire_strategy=expire_strategy,
+                expire_date=user.expire_date + timedelta(int(datelimit))
+                if expire_strategy == MarzneshinUserExpireStrategy.FIXED_DATE
+                else None,
+                usage_duration=user.usage_duration + (int(datelimit) * (24 * 60 * 60))
+                if expire_strategy == MarzneshinUserExpireStrategy.START_ON_FIRST_USE
+                else None,
+            ).dict()
+
+    return data
+
+
+def normal_charge_user_data(
+    types: ServerTypes,
+    user: MarzbanUserResponse | MarzneshinUserResponse,
+    datalimit: int,
+    datelimit: int,
+    status: str,
 ) -> dict:
     expire_strategy = _get_expire_strategy(types, status)
     expire_date = _get_expire_date(expire_strategy, datelimit)
@@ -126,7 +188,7 @@ def charge_user_data(
         ).dict()
     elif types == ServerTypes.MARZNESHIN.value:
         data = MarzneshinUserModify(
-            username=username,
+            username=user.username,
             data_limit=int(datalimit) * (1024**3),
             expire_strategy=expire_strategy,
             expire_date=expire_date,