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