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