diff --git a/app/dependencies.py b/app/dependencies.py index 3c95823..3e522b3 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(): # Dependency +def get_db(): with GetDB() as db: yield db @@ -27,9 +27,8 @@ def get_admin( return if dbadmin.password_reset_at: - if not payload.get("created_at"): - return - if dbadmin.password_reset_at > payload.get("created_at"): + created_at = payload.get("created_at") + if not created_at or dbadmin.password_reset_at > created_at: return return Admin.model_validate(dbadmin) @@ -61,12 +60,13 @@ def get_subscription_user( try: int(key, 16) except ValueError: - raise HTTPException(404) + raise HTTPException(status_code=404) db_user = crud.get_user(db, username) - if not db_user or db_user.key != key: - raise HTTPException(404) - return db_user + if db_user and db_user.key == key: + return db_user + else: + raise HTTPException(status_code=404) def get_user( @@ -76,8 +76,7 @@ def get_user( ): db_user = crud.get_user(db, username) if not ( - admin.is_sudo - or (db_user.admin and db_user.admin.username == admin.username) + admin.is_sudo or (db_user and db_user.admin.username == admin.username) ): raise HTTPException(status_code=403, detail="You're not allowed") diff --git a/app/marzneshin.py b/app/marzneshin.py index 667b949..db50858 100644 --- a/app/marzneshin.py +++ b/app/marzneshin.py @@ -3,6 +3,9 @@ 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 @@ -40,10 +43,21 @@ 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, ) @@ -91,18 +105,6 @@ 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 b3340a1..23e1a09 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): +def get_admins(db: DBDep, admin: SudoAdminDep, username: str | None = 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 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 marzban-cli instead.", + detail="You're not allowed to edit another sudoers account. Use marzneshin-cli instead.", ) dbadmin = crud.update_admin(db, dbadmin, modified_admin) @@ -114,7 +114,7 @@ def remove_admin(username: str, db: DBDep, admin: SudoAdminDep): if dbadmin.is_sudo: raise HTTPException( status_code=403, - detail="You're not allowed to delete sudoers accounts. Use marzban-cli instead.", + detail="You're not allowed to delete sudoers accounts. Use marzneshin-cli instead.", ) crud.remove_admin(db, dbadmin) diff --git a/app/routes/inbounds.py b/app/routes/inbounds.py index d96221b..0c1bffc 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 | None, db: DBDep): +def get_inbound_hosts(id: int, db: DBDep): """ Get hosts of a specific inbound """ diff --git a/cli/admin.py b/cli/admin.py index 492fbd3..5987e96 100644 --- a/cli/admin.py +++ b/cli/admin.py @@ -89,7 +89,10 @@ 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: @@ -180,16 +183,23 @@ 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 4b27c62..f8f03c6 100644 --- a/cli/subscription.py +++ b/cli/subscription.py @@ -19,7 +19,9 @@ 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. @@ -27,7 +29,9 @@ def get_link(username: str = typer.Option(..., *utils.FLAGS["username"], prompt= 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) @@ -58,7 +62,9 @@ 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 f7bd605..c667415 100644 --- a/cli/user.py +++ b/cli/user.py @@ -19,7 +19,9 @@ 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" ), @@ -56,7 +58,11 @@ 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 "", @@ -83,7 +89,8 @@ 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 0f563c0..3251f1a 100644 --- a/cli/utils.py +++ b/cli/utils.py @@ -51,7 +51,9 @@ 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 b904ac8..d416579 100755 --- a/marzneshin-cli.py +++ b/marzneshin-cli.py @@ -17,7 +17,9 @@ 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")