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