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,
55    rows: Iterable[Iterable[Any]],
56    console: Optional[Console] = None,
57):
58    for row in rows:
59        table.add_row(*row)
60
61    (console or rich_console).print(table)
62
63
64def readable_datetime(
65    date_time: Union[datetime, int, None],
66    include_date: bool = True,
67    include_time: bool = True,
68):
69    def get_datetime_format():
70        dt_format = ""
71        if include_date:
72            dt_format += "%d %B %Y"
73        if include_time:
74            if dt_format:
75                dt_format += ", "
76            dt_format += "%H:%M:%S"
77
78        return dt_format
79
80    if isinstance(date_time, int):
81        date_time = datetime.fromtimestamp(date_time)
82
83    return date_time.strftime(get_datetime_format()) if date_time else "-"
84
85
86def raise_if_falsy(obj: T, message: str) -> T:
87    if not obj:
88        error(text=message, auto_exit=True)
89
90    return obj