gitscope
https://dev.kronoz.win/r/marzneshin.git
cli/admin.py · v0.1.0
Source
1 a47eb4c Dawsh 2y ago from typing import Optional, Union
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 from rich.console import Console
6 a47eb4c Dawsh 2y ago from rich.panel import Panel
7 a47eb4c Dawsh 2y ago from sqlalchemy.exc import IntegrityError
8 a47eb4c Dawsh 2y ago from decouple import config, UndefinedValueError
9 a47eb4c Dawsh 2y ago
10 a47eb4c Dawsh 2y ago from app.db import GetDB
11 a47eb4c Dawsh 2y ago from app.db import crud
12 a47eb4c Dawsh 2y ago from app.db.models import Admin, User
13 a47eb4c Dawsh 2y ago from app.models.admin import AdminCreate, AdminPartialModify
14 a47eb4c Dawsh 2y ago from . import utils
15 a47eb4c Dawsh 2y ago
16 a47eb4c Dawsh 2y ago app = typer.Typer(no_args_is_help=True)
17 a47eb4c Dawsh 2y ago
18 a47eb4c Dawsh 2y ago
19 a47eb4c Dawsh 2y ago @app.command(name="list")
20 a47eb4c Dawsh 2y ago def list_admins(
21 a47eb4c Dawsh 2y ago offset: Optional[int] = typer.Option(None, *utils.FLAGS["offset"]),
22 a47eb4c Dawsh 2y ago limit: Optional[int] = typer.Option(None, *utils.FLAGS["limit"]),
23 a47eb4c Dawsh 2y ago username: Optional[str] = typer.Option(None, *utils.FLAGS["username"], help="Search by username"),
24 a47eb4c Dawsh 2y ago ):
25 a47eb4c Dawsh 2y ago """Displays a table of admins"""
26 a47eb4c Dawsh 2y ago with GetDB() as db:
27 a47eb4c Dawsh 2y ago admins: list[Admin] = crud.get_admins(db, offset=offset, limit=limit, username=username)
28 a47eb4c Dawsh 2y ago utils.print_table(
29 a47eb4c Dawsh 2y ago table=Table("Username", "Is sudo", "Created at"),
30 a47eb4c Dawsh 2y ago rows=[
31 a47eb4c Dawsh 2y ago (str(admin.username),
32 a47eb4c Dawsh 2y ago "✔️" if admin.is_sudo else "✖️",
33 a47eb4c Dawsh 2y ago utils.readable_datetime(admin.created_at))
34 a47eb4c Dawsh 2y ago for admin in admins
35 a47eb4c Dawsh 2y ago ]
36 a47eb4c Dawsh 2y ago )
37 a47eb4c Dawsh 2y ago
38 a47eb4c Dawsh 2y ago
39 a47eb4c Dawsh 2y ago @app.command(name="delete")
40 a47eb4c Dawsh 2y ago def delete_admin(
41 a47eb4c Dawsh 2y ago username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True),
42 a47eb4c Dawsh 2y ago yes_to_all: bool = typer.Option(False, *utils.FLAGS["yes_to_all"], help="Skips confirmations")
43 a47eb4c Dawsh 2y ago ):
44 a47eb4c Dawsh 2y ago """
45 a47eb4c Dawsh 2y ago Deletes the specified admin
46 a47eb4c Dawsh 2y ago
47 a47eb4c Dawsh 2y ago Confirmations can be skipped using `--yes/-y` option.
48 a47eb4c Dawsh 2y ago """
49 a47eb4c Dawsh 2y ago with GetDB() as db:
50 a47eb4c Dawsh 2y ago admin: Union[Admin, None] = crud.get_admin(db, username=username)
51 a47eb4c Dawsh 2y ago if not admin:
52 a47eb4c Dawsh 2y ago utils.error(f"There's no admin with username \"{username}\"!")
53 a47eb4c Dawsh 2y ago
54 a47eb4c Dawsh 2y ago if yes_to_all or typer.confirm(f'Are you sure about deleting "{username}"?', default=False):
55 a47eb4c Dawsh 2y ago crud.remove_admin(db, admin)
56 a47eb4c Dawsh 2y ago utils.success(f'"{username}" deleted successfully.')
57 a47eb4c Dawsh 2y ago else:
58 a47eb4c Dawsh 2y ago utils.error("Operation aborted!")
59 a47eb4c Dawsh 2y ago
60 a47eb4c Dawsh 2y ago
61 a47eb4c Dawsh 2y ago @app.command(name="create")
62 a47eb4c Dawsh 2y ago def create_admin(
63 a47eb4c Dawsh 2y ago username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True),
64 a47eb4c Dawsh 2y ago is_sudo: bool = typer.Option(None, *utils.FLAGS["is_sudo"], prompt=True),
65 a47eb4c Dawsh 2y ago password: str = typer.Option(..., prompt=True, confirmation_prompt=True,
66 a47eb4c Dawsh 2y ago hide_input=True, hidden=True, envvar=utils.PASSWORD_ENVIRON_NAME)
67 a47eb4c Dawsh 2y ago ):
68 a47eb4c Dawsh 2y ago """
69 a47eb4c Dawsh 2y ago Creates an admin
70 a47eb4c Dawsh 2y ago
71 a47eb4c Dawsh 2y ago Password can also be set using the `MARZBAN_ADMIN_PASSWORD` environment variable for non-interactive usages.
72 a47eb4c Dawsh 2y ago """
73 a47eb4c Dawsh 2y ago with GetDB() as db:
74 a47eb4c Dawsh 2y ago try:
75 a47eb4c Dawsh 2y ago crud.create_admin(db, AdminCreate(username=username, password=password, is_sudo=is_sudo))
76 a47eb4c Dawsh 2y ago utils.success(f'Admin "{username}" created successfully.')
77 a47eb4c Dawsh 2y ago except IntegrityError:
78 a47eb4c Dawsh 2y ago utils.error(f'Admin "{username}" already exists!')
79 a47eb4c Dawsh 2y ago
80 a47eb4c Dawsh 2y ago
81 a47eb4c Dawsh 2y ago @app.command(name="update")
82 a47eb4c Dawsh 2y ago def update_admin(username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True)):
83 a47eb4c Dawsh 2y ago """
84 a47eb4c Dawsh 2y ago Updates the specified admin
85 a47eb4c Dawsh 2y ago
86 a47eb4c Dawsh 2y ago NOTE: This command CAN NOT be used non-interactively.
87 a47eb4c Dawsh 2y ago """
88 a47eb4c Dawsh 2y ago
89 a47eb4c Dawsh 2y ago def _get_modify_model(admin: Admin):
90 a47eb4c Dawsh 2y ago Console().print(
91 a47eb4c Dawsh 2y ago Panel(f'Editing "{username}". Just press "Enter" to leave each field unchanged.')
92 a47eb4c Dawsh 2y ago )
93 a47eb4c Dawsh 2y ago
94 a47eb4c Dawsh 2y ago is_sudo: bool = typer.confirm("Is sudo", default=admin.is_sudo)
95 a47eb4c Dawsh 2y ago new_password: Union[str, None] = typer.prompt(
96 a47eb4c Dawsh 2y ago "New password",
97 a47eb4c Dawsh 2y ago default="",
98 a47eb4c Dawsh 2y ago show_default=False,
99 a47eb4c Dawsh 2y ago confirmation_prompt=True,
100 a47eb4c Dawsh 2y ago hide_input=True
101 a47eb4c Dawsh 2y ago ) or None
102 a47eb4c Dawsh 2y ago
103 a47eb4c Dawsh 2y ago return AdminPartialModify(
104 a47eb4c Dawsh 2y ago is_sudo=is_sudo,
105 a47eb4c Dawsh 2y ago password=new_password,
106 a47eb4c Dawsh 2y ago )
107 a47eb4c Dawsh 2y ago
108 a47eb4c Dawsh 2y ago with GetDB() as db:
109 a47eb4c Dawsh 2y ago admin: Union[Admin, None] = crud.get_admin(db, username=username)
110 a47eb4c Dawsh 2y ago if not admin:
111 a47eb4c Dawsh 2y ago utils.error(f"There's no admin with username \"{username}\"!")
112 a47eb4c Dawsh 2y ago
113 a47eb4c Dawsh 2y ago crud.partial_update_admin(db, admin, _get_modify_model(admin))
114 a47eb4c Dawsh 2y ago utils.success(f'Admin "{username}" updated successfully.')
115 a47eb4c Dawsh 2y ago
116 a47eb4c Dawsh 2y ago
117 a47eb4c Dawsh 2y ago @app.command(name="import-from-env")
118 a47eb4c Dawsh 2y ago def import_from_env(yes_to_all: bool = typer.Option(False, *utils.FLAGS["yes_to_all"], help="Skips confirmations")):
119 a47eb4c Dawsh 2y ago """
120 a47eb4c Dawsh 2y ago Imports the sudo admin from env
121 a47eb4c Dawsh 2y ago
122 a47eb4c Dawsh 2y ago Confirmations can be skipped using `--yes/-y` option.
123 a47eb4c Dawsh 2y ago
124 a47eb4c Dawsh 2y ago What does it do?
125 a47eb4c Dawsh 2y ago - Creates a sudo admin according to `SUDO_USERNAME` and `SUDO_PASSWORD`.
126 a47eb4c Dawsh 2y ago - Links any user which doesn't have an `admin_id` to the imported sudo admin.
127 a47eb4c Dawsh 2y ago """
128 a47eb4c Dawsh 2y ago try:
129 a47eb4c Dawsh 2y ago username, password = config("SUDO_USERNAME"), config("SUDO_PASSWORD")
130 a47eb4c Dawsh 2y ago except UndefinedValueError:
131 a47eb4c Dawsh 2y ago utils.error(
132 a47eb4c Dawsh 2y ago "Unable to get SUDO_USERNAME and/or SUDO_PASSWORD.\n"
133 a47eb4c Dawsh 2y ago "Make sure you have set them in the env file or as environment variables."
134 a47eb4c Dawsh 2y ago )
135 a47eb4c Dawsh 2y ago
136 a47eb4c Dawsh 2y ago if not (username and password):
137 a47eb4c Dawsh 2y ago utils.error("Unable to retrieve username and password.\n"
138 a47eb4c Dawsh 2y ago "Make sure both SUDO_USERNAME and SUDO_PASSWORD are set.")
139 a47eb4c Dawsh 2y ago
140 a47eb4c Dawsh 2y ago with GetDB() as db:
141 a47eb4c Dawsh 2y ago admin: Union[None, Admin] = None
142 a47eb4c Dawsh 2y ago
143 a47eb4c Dawsh 2y ago # If env admin already exists
144 a47eb4c Dawsh 2y ago if current_admin := crud.get_admin(db, username=username):
145 a47eb4c Dawsh 2y ago if not yes_to_all and not typer.confirm(
146 a47eb4c Dawsh 2y ago f'Admin "{username}" already exists. Do you want to sync it with env?', default=None
147 a47eb4c Dawsh 2y ago ):
148 a47eb4c Dawsh 2y ago utils.error("Aborted.")
149 a47eb4c Dawsh 2y ago
150 a47eb4c Dawsh 2y ago admin = crud.partial_update_admin(
151 a47eb4c Dawsh 2y ago db,
152 a47eb4c Dawsh 2y ago current_admin,
153 a47eb4c Dawsh 2y ago AdminPartialModify(password=password, is_sudo=True)
154 a47eb4c Dawsh 2y ago )
155 a47eb4c Dawsh 2y ago # If env admin does not exist yet
156 a47eb4c Dawsh 2y ago else:
157 a47eb4c Dawsh 2y ago admin = crud.create_admin(db, AdminCreate(
158 a47eb4c Dawsh 2y ago username=username,
159 a47eb4c Dawsh 2y ago password=password,
160 a47eb4c Dawsh 2y ago is_sudo=True
161 a47eb4c Dawsh 2y ago ))
162 a47eb4c Dawsh 2y ago
163 a47eb4c Dawsh 2y ago updated_user_count = db.query(User).filter_by(admin_id=None).update({"admin_id": admin.id})
164 a47eb4c Dawsh 2y ago db.commit()
165 a47eb4c Dawsh 2y ago
166 a47eb4c Dawsh 2y ago utils.success(
167 a47eb4c Dawsh 2y ago f'Admin "{username}" imported successfully.\n'
168 a47eb4c Dawsh 2y ago f"{updated_user_count} users' admin_id set to the {username}'s id.\n"
169 a47eb4c Dawsh 2y ago 'You must delete SUDO_USERNAME and SUDO_PASSWORD from your env file now.'
170 a47eb4c Dawsh 2y ago )