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