1import pydoc
 2from datetime import datetime
 3from typing import Any, Dict, Iterable, Optional, TypeVar, Union
 4
 5import typer
 6from rich.console import Console
 7from rich.table import Table
 8
 9from app.db import crud
10from app.db.models import User
11
12T = TypeVar("T")
13
14rich_console = Console()
15PASSWORD_ENVIRON_NAME = "MARZBAN_ADMIN_PASSWORD"
16
17FLAGS: Dict[str, tuple] = {
18    "username": ("--username", "-u"),
19    "limit": ("--limit", "-l"),
20    "offset": ("--offset", "-o"),
21    "yes_to_all": ("--yes", "-y"),
22    "is_sudo": ("--sudo/--no-sudo",),
23    "format": ("--format", "-f"),
24    "output_file": ("--output", "-o"),
25    "status": ("--status",),
26}
27
28
29def success(text: str, auto_exit: bool = True):
30    typer.echo(typer.style(text, fg=typer.colors.GREEN))
31    if auto_exit:
32        raise typer.Exit(0)
33
34
35def error(text: str, auto_exit: bool = True):
36    typer.echo(typer.style(text, fg=typer.colors.RED), err=True)
37    if auto_exit:
38        raise typer.Exit(-1)
39
40
41def paginate(text: str):
42    pydoc.pager(text)
43
44
45def get_user(db, username: str) -> User:
46    user: Union[User, None] = crud.get_user(db=db, username=username)
47    if not user:
48        error(f'User "{username}" not found!')
49
50    return user
51
52
53def print_table(
54    table: Table, rows: Iterable[Iterable[Any]], console: Optional[Console] = None
55):
56    for row in rows:
57        table.add_row(*row)
58
59    (console or rich_console).print(table)
60
61
62def readable_datetime(
63    date_time: Union[datetime, int, None],
64    include_date: bool = True,
65    include_time: bool = True,
66):
67    def get_datetime_format():
68        dt_format = ""
69        if include_date:
70            dt_format += "%d %B %Y"
71        if include_time:
72            if dt_format:
73                dt_format += ", "
74            dt_format += "%H:%M:%S"
75
76        return dt_format
77
78    if isinstance(date_time, int):
79        date_time = datetime.fromtimestamp(date_time)
80
81    return date_time.strftime(get_datetime_format()) if date_time else "-"
82
83
84def raise_if_falsy(obj: T, message: str) -> T:
85    if not obj:
86        error(text=message, auto_exit=True)
87
88    return obj