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