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