diff --git a/app/dependencies.py b/app/dependencies.py index eae4b5a..3c95823 100644 --- a/app/dependencies.py +++ b/app/dependencies.py @@ -9,7 +9,7 @@ from app.models.admin import Admin, oauth2_scheme from app.utils.jwt import get_admin_payload -def get_db(): +def get_db(): # Dependency with GetDB() as db: yield db @@ -26,9 +26,10 @@ def get_admin( if not dbadmin: return - if bool(dbadmin.password_reset_at): - created_at = payload.get("created_at") - if not created_at or dbadmin.password_reset_at > created_at: + if dbadmin.password_reset_at: + if not payload.get("created_at"): + return + if dbadmin.password_reset_at > payload.get("created_at"): return return Admin.model_validate(dbadmin) @@ -60,13 +61,12 @@ def get_subscription_user( try: int(key, 16) except ValueError: - raise HTTPException(status_code=404) + raise HTTPException(404) db_user = crud.get_user(db, username) - if db_user is User and db_user.key == key: - return db_user - else: - raise HTTPException(status_code=404, detail="User not found") + if not db_user or db_user.key != key: + raise HTTPException(404) + return db_user def get_user( @@ -76,7 +76,8 @@ def get_user( ): db_user = crud.get_user(db, username) if not ( - admin.is_sudo or (db_user and db_user.admin.username == admin.username) + admin.is_sudo + or (db_user.admin and db_user.admin.username == admin.username) ): raise HTTPException(status_code=403, detail="You're not allowed") @@ -89,7 +90,7 @@ def get_user( def parse_start_date(start: str | None = None): if not start: return datetime.fromtimestamp( - datetime.now().timestamp() - 30 * 24 * 3600 + datetime.utcnow().timestamp() - 30 * 24 * 3600 ) else: return datetime.fromisoformat(start) @@ -97,7 +98,7 @@ def parse_start_date(start: str | None = None): def parse_end_date(end: str | None = None): if not end: - return datetime.now() + return datetime.utcnow() else: return datetime.fromisoformat(end) diff --git a/app/marzneshin.py b/app/marzneshin.py index db50858..667b949 100644 --- a/app/marzneshin.py +++ b/app/marzneshin.py @@ -3,9 +3,6 @@ import logging from datetime import datetime as dt from datetime import timedelta as td -from typing import AsyncGenerator -from contextlib import asynccontextmanager - from apscheduler.schedulers.asyncio import AsyncIOScheduler from fastapi import FastAPI, Request, status from fastapi.encoders import jsonable_encoder @@ -43,21 +40,10 @@ from .tasks import ( logger = logging.getLogger(__name__) - -@asynccontextmanager -async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: - await nodes_startup() - yield - scheduler.shutdown() - logger.info("Sending pending notifications before shutdown...") - await send_notifications() - - app = FastAPI( title="MarzneshinAPI", description="Unified GUI Censorship Resistant Solution Powered by Xray", version=__version__, - lifespan=lifespan, docs_url="/docs" if DOCS else None, redoc_url="/redoc" if DOCS else None, ) @@ -105,6 +91,18 @@ if WEBHOOK_ADDRESS: ) +@app.on_event("startup") +async def on_start(): + await nodes_startup() + + +@app.on_event("shutdown") +async def on_shutdown(): + scheduler.shutdown() + logger.info("Sending pending notifications before shutdown...") + await send_notifications() + + @app.exception_handler(RequestValidationError) def validation_exception_handler( request: Request, exc: RequestValidationError diff --git a/app/routes/admin.py b/app/routes/admin.py index 65e7a8a..b3340a1 100644 --- a/app/routes/admin.py +++ b/app/routes/admin.py @@ -31,7 +31,7 @@ def authenticate_admin( @router.get("", response_model=Page[Admin]) -def get_admins(db: DBDep, admin: SudoAdminDep, username: str | None = None): +def get_admins(db: DBDep, admin: SudoAdminDep, username: str = None): query = db.query(DBAdmin) if username: query = query.filter(DBAdmin.username.ilike(f"%{username}%")) @@ -95,10 +95,10 @@ def modify_admin( raise HTTPException(status_code=404, detail="Admin not found") # If a sudoer admin wants to edit another sudoer - if username != admin.username and bool(dbadmin.is_sudo): + if (username != admin.username) and dbadmin.is_sudo: raise HTTPException( status_code=403, - detail="You're not allowed to edit another sudoers account. Use marzneshin-cli instead.", + detail="You're not allowed to edit another sudoers account. Use marzban-cli instead.", ) dbadmin = crud.update_admin(db, dbadmin, modified_admin) @@ -111,10 +111,10 @@ def remove_admin(username: str, db: DBDep, admin: SudoAdminDep): if not dbadmin: raise HTTPException(status_code=404, detail="Admin not found") - if bool(dbadmin.is_sudo): + if dbadmin.is_sudo: raise HTTPException( status_code=403, - detail="You're not allowed to delete sudoers accounts. Use marzneshin-cli instead.", + detail="You're not allowed to delete sudoers accounts. Use marzban-cli instead.", ) crud.remove_admin(db, dbadmin) diff --git a/app/routes/inbounds.py b/app/routes/inbounds.py index 0c1bffc..d96221b 100644 --- a/app/routes/inbounds.py +++ b/app/routes/inbounds.py @@ -84,7 +84,7 @@ def get_inbound(id: int, db: DBDep): @router.get("/{id}/hosts", response_model=Page[InboundHostResponse]) -def get_inbound_hosts(id: int, db: DBDep): +def get_inbound_hosts(id: int | None, db: DBDep): """ Get hosts of a specific inbound """ diff --git a/cli/admin.py b/cli/admin.py index 5987e96..492fbd3 100644 --- a/cli/admin.py +++ b/cli/admin.py @@ -89,10 +89,7 @@ def create_admin( with GetDB() as db: try: crud.create_admin( - db, - AdminCreate( - username=username, password=password, is_sudo=is_sudo - ), + db, AdminCreate(username=username, password=password, is_sudo=is_sudo) ) utils.success(f'Admin "{username}" created successfully.') except IntegrityError: @@ -183,23 +180,16 @@ def import_from_env( utils.error("Aborted.") admin = crud.partial_update_admin( - db, - current_admin, - AdminPartialModify(password=password, is_sudo=True), + db, current_admin, AdminPartialModify(password=password, is_sudo=True) ) # If env admin does not exist yet else: admin = crud.create_admin( - db, - AdminCreate( - username=username, password=password, is_sudo=True - ), + db, AdminCreate(username=username, password=password, is_sudo=True) ) updated_user_count = ( - db.query(User) - .filter_by(admin_id=None) - .update({"admin_id": admin.id}) + db.query(User).filter_by(admin_id=None).update({"admin_id": admin.id}) ) db.commit() diff --git a/cli/subscription.py b/cli/subscription.py index f8f03c6..4b27c62 100644 --- a/cli/subscription.py +++ b/cli/subscription.py @@ -19,9 +19,7 @@ class ConfigFormat(str, Enum): @app.command(name="get-link") -def get_link( - username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True) -): +def get_link(username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True)): """ Prints the given user's subscription link. @@ -29,9 +27,7 @@ def get_link( in order to work correctly. """ with GetDB() as db: - user: UserResponse = UserResponse.from_orm( - utils.get_user(db, username) - ) + user: UserResponse = UserResponse.from_orm(utils.get_user(db, username)) print(user.subscription_url) @@ -62,9 +58,7 @@ def get_config( otherwise will be shown in the terminal. """ with GetDB() as db: - user: UserResponse = UserResponse.from_orm( - utils.get_user(db, username) - ) + user: UserResponse = UserResponse.from_orm(utils.get_user(db, username)) conf: str = generate_subscription( user=user, config_format=config_format.name, as_base64=as_base64 ) diff --git a/cli/user.py b/cli/user.py index c667415..f7bd605 100644 --- a/cli/user.py +++ b/cli/user.py @@ -19,9 +19,7 @@ def list_users( username: Optional[List[str]] = typer.Option( None, *utils.FLAGS["username"], help="Search by username(s)" ), - status: Optional[crud.UserStatus] = typer.Option( - None, *utils.FLAGS["status"] - ), + status: Optional[crud.UserStatus] = typer.Option(None, *utils.FLAGS["status"]), admin: Optional[str] = typer.Option( None, "--admin", "--owner", help="Search by owner admin's username" ), @@ -58,11 +56,7 @@ def list_users( user.username, user.status.value, readable_size(user.used_traffic), - ( - readable_size(user.data_limit) - if user.data_limit - else "Unlimited" - ), + readable_size(user.data_limit) if user.data_limit else "Unlimited", user.data_limit_reset_strategy.value, utils.readable_datetime(user.expire, include_time=False), user.admin.username if user.admin else "", @@ -89,8 +83,7 @@ def set_owner( """ with GetDB() as db: user: User = utils.raise_if_falsy( - crud.get_user(db, username=username), - f'User "{username}" not found.', + crud.get_user(db, username=username), f'User "{username}" not found.' ) dbadmin = utils.raise_if_falsy( diff --git a/cli/utils.py b/cli/utils.py index 3251f1a..0f563c0 100644 --- a/cli/utils.py +++ b/cli/utils.py @@ -51,9 +51,7 @@ def get_user(db, username: str) -> User: def print_table( - table: Table, - rows: Iterable[Iterable[Any]], - console: Optional[Console] = None, + table: Table, rows: Iterable[Iterable[Any]], console: Optional[Console] = None ): for row in rows: table.add_row(*row) diff --git a/marzneshin-cli.py b/marzneshin-cli.py index d416579..b904ac8 100755 --- a/marzneshin-cli.py +++ b/marzneshin-cli.py @@ -17,9 +17,7 @@ app.add_typer(cli.user.app, name="user") # Hidden completion app app_completion = typer.Typer( - no_args_is_help=True, - help="Generate and install completion scripts.", - hidden=True, + no_args_is_help=True, help="Generate and install completion scripts.", hidden=True ) app.add_typer(app_completion, name="completion")