1from typing import List, Optional
  2
  3import typer
  4from rich.table import Table
  5
  6from app.db import GetDB, crud
  7from app.db.models import User
  8from app.utils.system import readable_size
  9
 10from . import utils
 11
 12app = typer.Typer(no_args_is_help=True)
 13
 14
 15@app.command(name="list")
 16def list_users(
 17    offset: Optional[int] = typer.Option(None, *utils.FLAGS["offset"]),
 18    limit: Optional[int] = typer.Option(None, *utils.FLAGS["limit"]),
 19    username: Optional[List[str]] = typer.Option(
 20        None, *utils.FLAGS["username"], help="Search by username(s)"
 21    ),
 22    status: Optional[crud.UserStatus] = typer.Option(None, *utils.FLAGS["status"]),
 23    admin: Optional[str] = typer.Option(
 24        None, "--admin", "--owner", help="Search by owner admin's username"
 25    ),
 26):
 27    """
 28    Displays a table of users
 29
 30    NOTE: Sorting is not currently available.
 31    """
 32    with GetDB() as db:
 33        users: list[User] = crud.get_users(
 34            db=db,
 35            offset=offset,
 36            limit=limit,
 37            usernames=username,
 38            status=status,
 39            admin=crud.get_admin(db, admin) if admin else None,
 40        )
 41
 42        utils.print_table(
 43            table=Table(
 44                "ID",
 45                "Username",
 46                "Status",
 47                "Used traffic",
 48                "Data limit",
 49                "Reset strategy",
 50                "Expires at",
 51                "Owner",
 52            ),
 53            rows=[
 54                (
 55                    str(user.id),
 56                    user.username,
 57                    user.status.value,
 58                    readable_size(user.used_traffic),
 59                    readable_size(user.data_limit) if user.data_limit else "Unlimited",
 60                    user.data_limit_reset_strategy.value,
 61                    utils.readable_datetime(user.expire, include_time=False),
 62                    user.admin.username if user.admin else "",
 63                )
 64                for user in users
 65            ],
 66        )
 67
 68
 69@app.command(name="set-owner")
 70def set_owner(
 71    username: str = typer.Option(None, *utils.FLAGS["username"], prompt=True),
 72    admin: str = typer.Option(
 73        None, "--admin", "--owner", prompt=True, help="Admin's username"
 74    ),
 75    yes_to_all: bool = typer.Option(
 76        False, *utils.FLAGS["yes_to_all"], help="Skips confirmations"
 77    ),
 78):
 79    """
 80    Transfers user's ownership
 81
 82    NOTE: This command needs additional confirmation for users who already have an owner.
 83    """
 84    with GetDB() as db:
 85        user: User = utils.raise_if_falsy(
 86            crud.get_user(db, username=username), f'User "{username}" not found.'
 87        )
 88
 89        dbadmin = utils.raise_if_falsy(
 90            crud.get_admin(db, username=admin), f'Admin "{admin}" not found.'
 91        )
 92
 93        # Ask for confirmation if user already has an owner
 94        if (
 95            user.admin
 96            and not yes_to_all
 97            and not typer.confirm(
 98                f'{username}\'s current owner is "{user.admin.username}".'
 99                f' Are you sure about transferring its ownership to "{admin}"?'
100            )
101        ):
102            utils.error("Aborted.")
103
104        crud.set_owner(db, user, dbadmin)
105
106        utils.success(f'{username}\'s owner successfully set to "{admin}".')