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