Next user planning (#1420)

* feat: Add auto_reset_usage field to user

* feat: next_user

* fix: removed somethings

* fix: namings and refactor

* fix(next_plan): api correct response

* fix: removed unnecessary next_plan deletes

app/db/__init__.py

@@ -29,7 +29,7 @@ def get_db():  # Dependency
                    get_or_create_inbound, get_system_usage,
                    get_tls_certificate, get_user, get_user_by_id, get_users,
                    get_users_count, remove_admin, remove_user, revoke_user_sub,
-                   set_owner, update_admin, update_user, update_user_status,
+                   set_owner, update_admin, update_user, update_user_status, reset_user_by_next,
                    update_user_sub, start_user_expire, get_admin_by_id,
                    get_admin_by_telegram_id)
 
@@ -47,6 +47,7 @@ def get_db():  # Dependency
     "update_user_status",
     "start_user_expire",
     "update_user_sub",
+    "reset_user_by_next",
     "revoke_user_sub",
     "set_owner",
     "get_system_usage",

app/db/crud.py

@@ -14,6 +14,7 @@
     JWT,
     TLS,
     Admin,
+    NextPlan,
     Node,
     NodeUsage,
     NodeUserUsage,
@@ -173,7 +174,7 @@ def get_user_queryset(db: Session) -> Query:
     Returns:
         Query: Base user query.
     """
-    return db.query(User).options(joinedload(User.admin))
+    return db.query(User).options(joinedload(User.admin)).options(joinedload(User.next_plan))
 
 
 def get_user(db: Session, username: str) -> Optional[User]:
@@ -386,7 +387,13 @@ def create_user(db: Session, user: UserCreate, admin: Admin = None) -> User:
         note=user.note,
         on_hold_expire_duration=(user.on_hold_expire_duration or None),
         on_hold_timeout=(user.on_hold_timeout or None),
-        auto_delete_in_days=user.auto_delete_in_days
+        auto_delete_in_days=user.auto_delete_in_days,
+        next_plan=NextPlan(
+            data_limit=user.next_plan.data_limit,
+            expire=user.next_plan.expire,
+            add_remaining_traffic=user.next_plan.add_remaining_traffic,
+            fire_on_either=user.next_plan.fire_on_either,
+        ) if user.next_plan else None
     )
     db.add(dbuser)
     db.commit()
@@ -504,7 +511,17 @@ def update_user(db: Session, dbuser: User, modify: UserModify) -> User:
 
     if modify.on_hold_expire_duration is not None:
         dbuser.on_hold_expire_duration = modify.on_hold_expire_duration
-
+
+    if modify.next_plan is not None:
+        dbuser.next_plan=NextPlan(
+            data_limit=modify.next_plan.data_limit,
+            expire=modify.next_plan.expire,
+            add_remaining_traffic=modify.next_plan.add_remaining_traffic,
+            fire_on_either=modify.next_plan.fire_on_either,
+        )
+    elif dbuser.next_plan is not None:
+        db.delete(dbuser.next_plan)
+
     dbuser.edit_at = datetime.utcnow()
 
     db.commit()
@@ -533,6 +550,46 @@ def reset_user_data_usage(db: Session, dbuser: User) -> User:
     dbuser.node_usages.clear()
     if dbuser.status not in (UserStatus.expired or UserStatus.disabled):
         dbuser.status = UserStatus.active.value
+
+    db.delete(dbuser.next_plan)
+    dbuser.next_plan = None
+    db.add(dbuser)
+
+    db.commit()
+    db.refresh(dbuser)
+    return dbuser
+
+
+def reset_user_by_next(db: Session, dbuser: User) -> User:
+    """
+    Resets the data usage of a user based on next user.
+
+    Args:
+        db (Session): Database session.
+        dbuser (User): The user object whose data usage is to be reset.
+
+    Returns:
+        User: The updated user object.
+    """
+
+    if (dbuser.next_plan is None):
+        return
+
+    usage_log = UserUsageResetLogs(
+        user=dbuser,
+        used_traffic_at_reset=dbuser.used_traffic,
+    )
+    db.add(usage_log)
+
+    dbuser.node_usages.clear()
+    dbuser.status = UserStatus.active.value
+
+    dbuser.data_limit = dbuser.next_plan.data_limit + (0 if dbuser.next_plan.add_remaining_traffic else dbuser.data_limit - dbuser.used_traffic)
+    dbuser.expire = dbuser.next_plan.expire
+
+    dbuser.used_traffic = 0
+    db.delete(dbuser.next_plan)
+    dbuser.next_plan = None
     db.add(dbuser)
 
     db.commit()
@@ -603,6 +660,8 @@ def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
             dbuser.status = UserStatus.active
         dbuser.usage_logs.clear()
         dbuser.node_usages.clear()
+        db.delete(dbuser.next_plan)
+        dbuser.next_plan = None
         db.add(dbuser)
 
     db.commit()

app/db/migrations/versions/c3cd674b9bcd_added_next_plan_for_user.py

@@ -0,0 +1,37 @@
+"""Added next_plan for user
+
+Revision ID: c3cd674b9bcd
+Revises: 21226bc711ac
+Create Date: 2024-11-07 12:45:51.159960
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'c3cd674b9bcd'
+down_revision = '21226bc711ac'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('next_plans',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('user_id', sa.Integer(), nullable=False),
+    sa.Column('data_limit', sa.BigInteger(), nullable=False),
+    sa.Column('expire', sa.Integer(), nullable=True),
+    sa.Column('add_remaining_traffic', sa.Boolean(), server_default='0', nullable=False),
+    sa.Column('fire_on_either', sa.Boolean(), server_default='0', nullable=False),
+    sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('next_plans')
+    # ### end Alembic commands ###

app/db/models.py

@@ -67,6 +67,13 @@ class User(Base):
 
     edit_at = Column(DateTime, nullable=True, default=None)
     last_status_change = Column(DateTime, default=datetime.utcnow, nullable=True)
+
+    next_plan = relationship(
+        "NextPlan",
+        uselist=False,
+        back_populates="user",
+        cascade="all, delete-orphan"
+    )
 
     @hybrid_property
     def reseted_usage(self):
@@ -126,6 +133,19 @@ def inbounds(self):
 )
 
 
+class NextPlan(Base):
+    __tablename__ = 'next_plans'
+
+    id = Column(Integer, primary_key=True)
+    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
+    data_limit = Column(BigInteger, nullable=False)
+    expire = Column(Integer, nullable=True)
+    add_remaining_traffic = Column(Boolean, nullable=False, default=False, server_default='0')
+    fire_on_either = Column(Boolean, nullable=False, default=True, server_default='0')
+
+    user = relationship("User", back_populates="next_plan")
+
+
 class UserTemplate(Base):
     __tablename__ = "user_templates"
 

app/discord/__init__.py

@@ -7,6 +7,7 @@
     report_user_deletion,
     report_status_change,
     report_user_usage_reset,
+    report_user_data_reset_by_next,
     report_user_subscription_revoked,
     report_login
 )
@@ -18,6 +19,7 @@
     "report_user_deletion",
     "report_status_change",
     "report_user_usage_reset",
+    "report_user_data_reset_by_next",
     "report_user_subscription_revoked",
     "report_login"
 ]

app/discord/handlers/report.py

@@ -1,5 +1,6 @@
 import requests
 from datetime import datetime
+from app.db.models import User
 from app.utils.system import readable_size
 from app.models.user import UserDataLimitResetStrategy
 from app.models.admin import Admin
@@ -56,7 +57,7 @@ def report_status_change(username: str, status: str, admin: Admin = None):
         admin_webhook=admin.discord_webhook if admin and admin.discord_webhook else None
         )
 
-def report_new_user(username: str, by: str, expire_date: int, data_limit: int, proxies: list, 
+def report_new_user(username: str, by: str, expire_date: int, data_limit: int, proxies: list, has_next_plan: bool,
                     data_limit_reset_strategy:UserDataLimitResetStrategy, admin: Admin = None):
 
     data_limit=readable_size(data_limit) if data_limit else "Unlimited"
@@ -73,7 +74,8 @@ def report_new_user(username: str, by: str, expire_date: int, data_limit: int, p
 **Traffic Limit:** {data_limit}
 **Expire Date:** {expire_date}
 **Proxies:** {proxies}
-**Data Limit Reset Strategy:**{data_limit_reset_strategy}""",
+**Data Limit Reset Strategy:**{data_limit_reset_strategy}
+**Has Next Plan:**{has_next_plan}""",
 
                 "footer": {
                     "text": f"Belongs To: {admin.username if admin else None}\nBy: {by}"
@@ -87,7 +89,7 @@ def report_new_user(username: str, by: str, expire_date: int, data_limit: int, p
         admin_webhook=admin.discord_webhook if admin and admin.discord_webhook else None
         )
 
-def report_user_modification(username: str, expire_date: int, data_limit: int, proxies: list, by: str, 
+def report_user_modification(username: str, expire_date: int, data_limit: int, proxies: list, by: str, has_next_plan: bool,
                     data_limit_reset_strategy:UserDataLimitResetStrategy, admin: Admin = None):
 
     data_limit=readable_size(data_limit) if data_limit else "Unlimited"
@@ -105,7 +107,8 @@ def report_user_modification(username: str, expire_date: int, data_limit: int, p
 **Traffic Limit:** {data_limit}
 **Expire Date:** {expire_date}
 **Proxies:** {proxies}
-**Data Limit Reset Strategy:**{data_limit_reset_strategy}""",
+**Data Limit Reset Strategy:**{data_limit_reset_strategy}
+**Has Next Plan:**{has_next_plan}""",
 
                 "footer": {
                     "text": f"Belongs To: {admin.username if admin else None}\nBy: {by}"
@@ -157,6 +160,29 @@ def report_user_usage_reset(username: str, by: str, admin: Admin = None):
         admin_webhook=admin.discord_webhook if admin and admin.discord_webhook else None
         )
 
+def report_user_data_reset_by_next(user: User, admin: Admin = None):
+    userUsageReset = {
+        'content': '',
+        'embeds': [
+            {
+                'title': ':repeat: AutoReset',
+                'description': f"""
+                **Username:** {user.username}
+**Traffic Limit:** {user.data_limit}
+**Expire Date:** {user.expire}""",
+
+                "footer": {
+                    "text": f"Belongs To: {admin.username if admin else None}"
+                },
+                'color': int('00ffff', 16)
+            }
+        ]
+    }
+    send_webhooks(
+        json_data=userUsageReset,
+        admin_webhook=admin.discord_webhook if admin and admin.discord_webhook else None
+        )
+
 def report_user_subscription_revoked(username: str, by: str, admin: Admin = None):
     subscriptionRevoked = {
         'content': '',
@@ -196,4 +222,4 @@ def report_login(username: str, password: str, client_ip: str, status: str):
     send_webhooks(
         json_data=login,
         admin_webhook=None
-        )
+        )

app/jobs/review_users.py

@@ -5,7 +5,7 @@
 
 from app import logger, scheduler, xray
 from app.db import (GetDB, get_notification_reminder, get_users,
-                    start_user_expire, update_user_status)
+                    start_user_expire, update_user_status, reset_user_by_next)
 from app.models.user import ReminderType, UserResponse, UserStatus
 from app.utils import report
 from app.utils.helpers import (calculate_expiration_days,
@@ -42,6 +42,12 @@ def add_notification_reminders(db: Session, user: "User", now: datetime = dateti
                     )
                 break
 
+def reset_user_by_next_report(db: Session, user: "User"):
+    user = reset_user_by_next(db, user)
+
+    xray.operations.update_user(user)
+
+    report.user_data_reset_by_next(user=UserResponse.from_orm(user), user_admin=user.admin)
 
 def review():
     now = datetime.utcnow()
@@ -51,6 +57,18 @@ def review():
 
             limited = user.data_limit and user.used_traffic >= user.data_limit
             expired = user.expire and user.expire <= now_ts
+
+            if (limited or expired) and user.next_plan is not None:
+                if user.next_plan is not None:
+
+                    if user.next_plan.fire_on_either:
+                        reset_user_by_next_report(db, user)
+                        continue
+
+                    elif limited and expired:
+                        reset_user_by_next_report(db, user)
+                        continue
+
             if limited:
                 status = UserStatus.limited
             elif expired:

app/models/user.py

@@ -48,6 +48,14 @@ class UserDataLimitResetStrategy(str, Enum):
     month = "month"
     year = "year"
 
+class NextPlanModel(BaseModel):
+    data_limit: Optional[int]
+    expire: Optional[int]
+    add_remaining_traffic: bool = False
+    fire_on_either: bool = True
+
+    class Config:
+        orm_mode = True
 
 class User(BaseModel):
     proxies: Dict[ProxyTypes, ProxySettings] = {}
@@ -67,7 +75,9 @@ class User(BaseModel):
     on_hold_timeout: Optional[Union[datetime, None]] = Field(None, nullable=True)
 
     auto_delete_in_days: Optional[int] = Field(None, nullable=True)
-
+
+    next_plan: Optional[NextPlanModel] = Field(None, nullable=True)
+
     @validator("proxies", pre=True, always=True)
     def validate_proxies(cls, v, values, **kwargs):
         if not v:
@@ -116,6 +126,12 @@ class Config:
                     "vmess": ["VMess TCP", "VMess Websocket"],
                     "vless": ["VLESS TCP REALITY", "VLESS GRPC REALITY"],
                 },
+                "next_plan": {
+                    "data_limit": 0,
+                    "expire": 0,
+                    "add_remaining_traffic": False,
+                    "fire_on_either": True
+                },
                 "expire": 0,
                 "data_limit": 0,
                 "data_limit_reset_strategy": "no_reset",
@@ -199,6 +215,12 @@ class Config:
                     "vmess": ["VMess TCP", "VMess Websocket"],
                     "vless": ["VLESS TCP REALITY", "VLESS GRPC REALITY"],
                 },
+                "next_plan": {
+                    "data_limit": 0,
+                    "expire": 0,
+                    "add_remaining_traffic": False,
+                    "fire_on_either": True
+                },
                 "expire": 0,
                 "data_limit": 0,
                 "data_limit_reset_strategy": "no_reset",
@@ -337,4 +359,4 @@ class UserUsagesResponse(BaseModel):
     usages: List[UserUsageResponse]
 
 class UsersUsagesResponse(BaseModel):
-    usages: List[UserUsageResponse]
+    usages: List[UserUsageResponse]