gitscope
https://dev.kronoz.win/r/marzneshin.git
cli/admin.py · v0.3.0-alpha3
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 01cdf2e Dawsh 2y ago db, AdminCreate(username=username, password=password, is_sudo=is_sudo)
93 01cdf2e Dawsh 2y ago )
94 a47eb4c Dawsh 2y ago utils.success(f'Admin "{username}" created successfully.')
95 a47eb4c Dawsh 2y ago except IntegrityError:
96 a47eb4c Dawsh 2y ago utils.error(f'Admin "{username}" already exists!')
97 a47eb4c Dawsh 2y ago
98 a47eb4c Dawsh 2y ago
99 a47eb4c Dawsh 2y ago @app.command(name="update")
100 01cdf2e Dawsh 2y ago def update_admin(
101 01cdf2e Dawsh 2y ago username: str = typer.Option(..., *utils.FLAGS["username"], prompt=True)
102 01cdf2e Dawsh 2y ago ):
103 a47eb4c Dawsh 2y ago """
104 a47eb4c Dawsh 2y ago Updates the specified admin
105 a47eb4c Dawsh 2y ago
106 a47eb4c Dawsh 2y ago NOTE: This command CAN NOT be used non-interactively.
107 a47eb4c Dawsh 2y ago """
108 a47eb4c Dawsh 2y ago
109 a47eb4c Dawsh 2y ago def _get_modify_model(admin: Admin):
110 a47eb4c Dawsh 2y ago Console().print(
111 01cdf2e Dawsh 2y ago Panel(
112 01cdf2e Dawsh 2y ago f'Editing "{username}". Just press "Enter" to leave each field unchanged.'
113 01cdf2e Dawsh 2y ago )
114 a47eb4c Dawsh 2y ago )
115 a47eb4c Dawsh 2y ago
116 a47eb4c Dawsh 2y ago is_sudo: bool = typer.confirm("Is sudo", default=admin.is_sudo)
117 01cdf2e Dawsh 2y ago new_password: Union[str, None] = (
118 01cdf2e Dawsh 2y ago typer.prompt(
119 01cdf2e Dawsh 2y ago "New password",
120 01cdf2e Dawsh 2y ago default="",
121 01cdf2e Dawsh 2y ago show_default=False,
122 01cdf2e Dawsh 2y ago confirmation_prompt=True,
123 01cdf2e Dawsh 2y ago hide_input=True,
124 01cdf2e Dawsh 2y ago )
125 01cdf2e Dawsh 2y ago or None
126 01cdf2e Dawsh 2y ago )
127 a47eb4c Dawsh 2y ago
128 a47eb4c Dawsh 2y ago return AdminPartialModify(
129 a47eb4c Dawsh 2y ago is_sudo=is_sudo,
130 a47eb4c Dawsh 2y ago password=new_password,
131 a47eb4c Dawsh 2y ago )
132 a47eb4c Dawsh 2y ago
133 a47eb4c Dawsh 2y ago with GetDB() as db:
134 a47eb4c Dawsh 2y ago admin: Union[Admin, None] = crud.get_admin(db, username=username)
135 a47eb4c Dawsh 2y ago if not admin:
136 01cdf2e Dawsh 2y ago utils.error(f'There\'s no admin with username "{username}"!')
137 a47eb4c Dawsh 2y ago
138 a47eb4c Dawsh 2y ago crud.partial_update_admin(db, admin, _get_modify_model(admin))
139 a47eb4c Dawsh 2y ago utils.success(f'Admin "{username}" updated successfully.')
140 a47eb4c Dawsh 2y ago
141 a47eb4c Dawsh 2y ago
142 a47eb4c Dawsh 2y ago @app.command(name="import-from-env")
143 01cdf2e Dawsh 2y ago def import_from_env(
144 01cdf2e Dawsh 2y ago yes_to_all: bool = typer.Option(
145 01cdf2e Dawsh 2y ago False, *utils.FLAGS["yes_to_all"], help="Skips confirmations"
146 01cdf2e Dawsh 2y ago )
147 01cdf2e Dawsh 2y ago ):
148 a47eb4c Dawsh 2y ago """
149 a47eb4c Dawsh 2y ago Imports the sudo admin from env
150 a47eb4c Dawsh 2y ago
151 a47eb4c Dawsh 2y ago Confirmations can be skipped using `--yes/-y` option.
152 a47eb4c Dawsh 2y ago
153 a47eb4c Dawsh 2y ago What does it do?
154 a47eb4c Dawsh 2y ago - Creates a sudo admin according to `SUDO_USERNAME` and `SUDO_PASSWORD`.
155 a47eb4c Dawsh 2y ago - Links any user which doesn't have an `admin_id` to the imported sudo admin.
156 a47eb4c Dawsh 2y ago """
157 a47eb4c Dawsh 2y ago try:
158 a47eb4c Dawsh 2y ago username, password = config("SUDO_USERNAME"), config("SUDO_PASSWORD")
159 a47eb4c Dawsh 2y ago except UndefinedValueError:
160 a47eb4c Dawsh 2y ago utils.error(
161 a47eb4c Dawsh 2y ago "Unable to get SUDO_USERNAME and/or SUDO_PASSWORD.\n"
162 a47eb4c Dawsh 2y ago "Make sure you have set them in the env file or as environment variables."
163 a47eb4c Dawsh 2y ago )
164 a47eb4c Dawsh 2y ago
165 a47eb4c Dawsh 2y ago if not (username and password):
166 01cdf2e Dawsh 2y ago utils.error(
167 01cdf2e Dawsh 2y ago "Unable to retrieve username and password.\n"
168 01cdf2e Dawsh 2y ago "Make sure both SUDO_USERNAME and SUDO_PASSWORD are set."
169 01cdf2e Dawsh 2y ago )
170 a47eb4c Dawsh 2y ago
171 a47eb4c Dawsh 2y ago with GetDB() as db:
172 a47eb4c Dawsh 2y ago admin: Union[None, Admin] = None
173 a47eb4c Dawsh 2y ago
174 a47eb4c Dawsh 2y ago # If env admin already exists
175 a47eb4c Dawsh 2y ago if current_admin := crud.get_admin(db, username=username):
176 a47eb4c Dawsh 2y ago if not yes_to_all and not typer.confirm(
177 01cdf2e Dawsh 2y ago f'Admin "{username}" already exists. Do you want to sync it with env?',
178 01cdf2e Dawsh 2y ago default=None,
179 a47eb4c Dawsh 2y ago ):
180 a47eb4c Dawsh 2y ago utils.error("Aborted.")
181 a47eb4c Dawsh 2y ago
182 a47eb4c Dawsh 2y ago admin = crud.partial_update_admin(
183 01cdf2e Dawsh 2y ago db, current_admin, AdminPartialModify(password=password, is_sudo=True)
184 a47eb4c Dawsh 2y ago )
185 a47eb4c Dawsh 2y ago # If env admin does not exist yet
186 a47eb4c Dawsh 2y ago else:
187 01cdf2e Dawsh 2y ago admin = crud.create_admin(
188 01cdf2e Dawsh 2y ago db, AdminCreate(username=username, password=password, is_sudo=True)
189 01cdf2e Dawsh 2y ago )
190 a47eb4c Dawsh 2y ago
191 01cdf2e Dawsh 2y ago updated_user_count = (
192 01cdf2e Dawsh 2y ago db.query(User).filter_by(admin_id=None).update({"admin_id": admin.id})
193 01cdf2e Dawsh 2y ago )
194 a47eb4c Dawsh 2y ago db.commit()
195 a47eb4c Dawsh 2y ago
196 a47eb4c Dawsh 2y ago utils.success(
197 a47eb4c Dawsh 2y ago f'Admin "{username}" imported successfully.\n'
198 a47eb4c Dawsh 2y ago f"{updated_user_count} users' admin_id set to the {username}'s id.\n"
199 01cdf2e Dawsh 2y ago "You must delete SUDO_USERNAME and SUDO_PASSWORD from your env file now."
200 a47eb4c Dawsh 2y ago )