gitscope
← marzneshin / v0.3.0-alpha11 / app/routes/user.py
https://dev.kronoz.win/r/marzneshin.git
app/routes/user.py · Python
Blame Raw
  1import asyncio
  2import logging
  3from datetime import datetime, timezone
  4from enum import Enum
  5
  6import sqlalchemy
  7from fastapi import APIRouter
  8from fastapi import HTTPException, Query
  9from fastapi_pagination.ext.sqlalchemy import paginate
 10from fastapi_pagination.links import Page
 11
 12from app import marznode
 13from app.db import crud, User
 14from app.dependencies import (
 15    DBDep,
 16    AdminDep,
 17    SudoAdminDep,
 18    UserDep,
 19    StartDateDep,
 20    EndDateDep,
 21)
 22from app.models.user import (
 23    UserCreate,
 24    UserModify,
 25    UserResponse,
 26    UserStatus,
 27    UserUsagesResponse,
 28)
 29from app.utils import report
 30
 31logger = logging.getLogger(__name__)
 32router = APIRouter(prefix="/users", tags=["User"])
 33
 34
 35class UsersSortingOptions(str, Enum):
 36    username = "username"
 37    used_traffic = "used_traffic"
 38    data_limit = "data_limit"
 39    expire = "expire"
 40    created_at = "created_at"
 41
 42
 43@router.get("", response_model=Page[UserResponse])
 44def get_users(
 45    db: DBDep,
 46    admin: AdminDep,
 47    username: list[str] = Query(None),
 48    status: list[UserStatus] = Query(None),
 49    order_by: UsersSortingOptions = Query(None),
 50    descending: bool = Query(False),
 51):
 52    """
 53    Filters users based on the options
 54    """
 55    dbadmin = crud.get_admin(db, admin.username)
 56
 57    query = db.query(User)
 58
 59    admin = dbadmin if not admin.is_sudo else None
 60
 61    if username:
 62        if len(username) > 1:
 63            query = query.filter(User.username.in_(username))
 64        else:
 65            query = query.filter(User.username.ilike(f"%{username[0]}%"))
 66
 67    if status:
 68        query = query.filter(User.status.in_(status))
 69
 70    if admin:
 71        query = query.filter(User.admin == admin)
 72
 73    if order_by:
 74        order_column = getattr(User, order_by)
 75        if descending:
 76            order_column = order_column.desc()
 77        query = query.order_by(order_column)
 78
 79    return paginate(db, query)
 80
 81
 82@router.post("", response_model=UserResponse)
 83async def add_user(new_user: UserCreate, db: DBDep, admin: AdminDep):
 84    """
 85    Add a new user
 86
 87    - **username** must have 3 to 32 characters and is allowed to contain a-z, 0-9, and underscores in between
 88    - **expire** must be a UTC timestamp
 89    - **data_limit** must be in Bytes, e.g. 1073741824B = 1GB
 90    - **services** list of service ids
 91    """
 92    # TODO expire should be datetime instead of timestamp
 93
 94    try:
 95        db_user = crud.create_user(
 96            db, new_user, admin=crud.get_admin(db, admin.username)
 97        )
 98    except sqlalchemy.exc.IntegrityError:
 99        db.rollback()
100        raise HTTPException(status_code=409, detail="User already exists")
101
102    user = UserResponse.model_validate(db_user)
103    await marznode.operations.update_user(user=db_user)
104    asyncio.create_task(report.user_created(user=user, user_id=db_user.id, by=admin))
105    logger.info("New user `%s` added", db_user.username)
106    return user
107
108
109@router.post("/reset")
110async def reset_users_data_usage(db: DBDep, admin: SudoAdminDep):
111    """
112    Reset all users data usage,
113    You will need to restart for this to take effect for now
114    """
115    dbadmin = crud.get_admin(db, admin.username)
116    crud.reset_all_users_data_usage(db=db, admin=dbadmin)
117    return {}
118
119
120@router.delete("/expired")
121def delete_expired(passed_time: int, db: DBDep, admin: AdminDep):
122    """
123    Delete expired users
124    - **passed_time** must be a timestamp
125    - This function will delete all expired users that meet the specified number of days passed and can't be undone.
126    """
127
128    dbadmin = crud.get_admin(db, admin.username)
129
130    db_users = crud.get_users(
131        db=db,
132        status=[UserStatus.expired, UserStatus.limited],
133        admin=dbadmin if not admin.is_sudo else None,
134    )
135
136    current_time = int(datetime.now(timezone.utc).timestamp())
137    expiration_threshold = current_time - passed_time
138    expired_users = [
139        user
140        for user in db_users
141        if user.expire is not None and user.expire <= expiration_threshold
142    ]
143    if not expired_users:
144        raise HTTPException(status_code=404, detail="No expired user found.")
145
146    for db_user in expired_users:
147        crud.remove_user(db, db_user)
148
149        asyncio.create_task(report.user_deleted(username=db_user.username, by=admin))
150
151        logger.info("User `%s` removed", db_user.username)
152
153    return {}
154
155
156@router.get("/{username}", response_model=UserResponse)
157def get_user(db_user: UserDep):
158    """
159    Get users information
160    """
161    return db_user
162
163
164@router.put("/{username}", response_model=UserResponse)
165async def modify_user(
166    db_user: UserDep, modifications: UserModify, db: DBDep, admin: AdminDep
167):
168    """
169    Modify a user
170
171    - set **expire** to 0 to make the user unlimited in time, null for no change
172    - set **data_limit** to 0 to make the user unlimited in data, null for no change
173    """
174    old_status = db_user.status
175    status_change = bool(
176        modifications.status and modifications.status != db_user.status
177    )
178
179    old_inbounds = {(i.node_id, i.protocol, i.tag) for i in db_user.inbounds}
180    new_user = crud.update_user(db, db_user, modifications)
181    new_inbounds = {(i.node_id, i.protocol, i.tag) for i in new_user.inbounds}
182    user = UserResponse.model_validate(db_user)
183
184    inbound_change = old_inbounds != new_inbounds
185    if (
186        inbound_change
187        and user.enabled
188        and user.status in {UserStatus.active, UserStatus.on_hold}
189    ):
190        await marznode.operations.update_user(new_user, old_inbounds)
191
192    asyncio.create_task(report.user_updated(user=user, by=admin))
193
194    logger.info("User `%s` modified", user.username)
195
196    if status_change:
197        asyncio.create_task(
198            report.status_change(
199                username=user.username, status=user.status, user=user, by=admin
200            )
201        )
202
203        logger.info(
204            "User `%s` status changed from `%s` to `%s`",
205            user.username,
206            old_status,
207            user.status,
208        )
209
210    return user
211
212
213@router.delete("/{username}")
214async def remove_user(db_user: UserDep, db: DBDep, admin: AdminDep):
215    """
216    Remove a user
217    """
218    await marznode.operations.remove_user(db_user)
219
220    crud.remove_user(db, db_user)
221    db.flush()
222
223    asyncio.create_task(report.user_deleted(username=db_user.username, by=admin))
224    logger.info("User %s deleted", db_user.username)
225    return {}
226
227
228@router.post("/{username}/reset", response_model=UserResponse)
229async def reset_user_data_usage(db_user: UserDep, db: DBDep, admin: AdminDep):
230    """
231    Reset user data usage
232    """
233    previous_status = db_user.status
234    db_user = crud.reset_user_data_usage(db, db_user)
235
236    if db_user.status == UserStatus.active and previous_status == UserStatus.limited:
237        await marznode.operations.update_user(db_user)
238
239    user = UserResponse.model_validate(db_user)
240    asyncio.create_task(report.user_data_usage_reset(user=user, by=admin))
241
242    logger.info("User `%s`'s usage was reset", db_user.username)
243
244    return user
245
246
247@router.post("/{username}/enable", response_model=UserResponse)
248async def enable_user(db_user: UserDep, db: DBDep, admin: AdminDep):
249    """
250    Enables a user
251    """
252    if db_user.enabled:
253        raise HTTPException(409, "User is already enabled")
254    db_user.enabled = True
255    db.commit()
256
257    await marznode.operations.update_user(db_user)
258
259    user = UserResponse.model_validate(db_user)
260
261    logger.info("User `%s` has been enabled", db_user.username)
262
263    return user
264
265
266@router.post("/{username}/disable", response_model=UserResponse)
267async def disable_user(db_user: UserDep, db: DBDep, admin: AdminDep):
268    """
269    Disables a user
270    """
271    if not db_user.enabled:
272        raise HTTPException(409, "User is not enabled")
273    db_user.enabled = False
274    db.commit()
275
276    await marznode.operations.remove_user(db_user)
277
278    user = UserResponse.model_validate(db_user)
279
280    logger.info("User `%s` has been disabled", db_user.username)
281
282    return user
283
284
285@router.post("/{username}/revoke_sub", response_model=UserResponse)
286async def revoke_user_subscription(db_user: UserDep, db: DBDep, admin: AdminDep):
287    """
288    Revoke users subscription (Subscription link and proxies)
289    """
290    db_user = crud.revoke_user_sub(db, db_user)
291
292    if db_user.status in [UserStatus.active, UserStatus.on_hold]:
293        await marznode.operations.remove_user(db_user)
294        await marznode.operations.update_user(db_user)
295    user = UserResponse.model_validate(db_user)
296    asyncio.create_task(report.user_subscription_revoked(user=user, by=admin))
297
298    logger.info("User %s subscription revoked", db_user.username)
299
300    return user
301
302
303@router.get("/{username}/usage", response_model=UserUsagesResponse)
304def get_user_usage(
305    db: DBDep, db_user: UserDep, start_date: StartDateDep, end_date: EndDateDep
306):
307    """
308    Get users usage
309    """
310    usages = crud.get_user_usages(db, db_user, start_date, end_date)
311
312    return {"usages": usages, "username": db_user.username}
313
314
315@router.put("/{username}/set-owner", response_model=UserResponse)
316def set_owner(username: str, admin_username: str, db: DBDep, admin: SudoAdminDep):
317    db_user = crud.get_user(db, username)
318    if not db_user:
319        raise HTTPException(status_code=404, detail="User not found")
320
321    new_admin = crud.get_admin(db, username=admin_username)
322    if not new_admin:
323        raise HTTPException(status_code=404, detail="Admin not found")
324
325    db_user = crud.set_owner(db, db_user, new_admin)
326    user = UserResponse.model_validate(db_user)
327
328    logger.info("`%s`'s owner successfully set to `%s`", user.username, admin.username)
329
330    return user