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