traffic.py

#! /usr/bin/env python3.10

"""
This module is used for checking V2Ray traffic consumption.
Install Pandas beforehand via: `pip install pandas`

Sample Command:
$ xray api statsquery --server="127.0.0.1:10085" -pattern '' | python3 traffic.py -

GitHub: https://github.com/Mahyar24/V2Conf
[email protected], Sat 19 Aug 2023
"""

import enum
import json
import sys

import pandas as pd


class Unit(enum.Enum):
    KB = 1_000
    MB = 1_000_000
    GB = 1_000_000_000
    TB = 1_000_000_000_000


def calculate(stats: list[dict]) -> pd.DataFrame:
    return (
        pd.DataFrame(stats)
        .assign(
            **{
                "entity": lambda df_: df_.pop("name")
                .str.replace("user>>>", "*>>>")
                .str.replace(">>>traffic", "")
                # .str.replace("@email.com", "") 
                .str.replace("downlink", "Download")
                .str.replace("uplink", "Upload")
                .str.split(">>>"),
            }
        )
        .rename({"value": "traffic"}, axis=1)
        .pipe(
            lambda df_: pd.DataFrame(
                df_.pop("entity").to_list(), columns=["O/I", "user", "D/U"]
            ).join(df_)
        )
        .reset_index(drop=True)
        .drop("O/I", axis=1)
        .astype({"D/U": "category", "traffic": "int"})
        .sort_values(["user", "D/U"])
        .set_index(["user", "D/U"])
        .loc[
            lambda df_: df_.groupby("user")["traffic"]
            .sum()
            .sort_values(ascending=False)
            .index
        ]
    )


def show(df: pd.DataFrame, unit: Unit = Unit.MB, floating_points: int = 1) -> None:
    print(
        df.assign(
            **{
                "traffic": lambda df_: df_.pop("traffic")
                .div(unit.value)
                .map(lambda x: f"{x:,.{floating_points}f} {unit.name}")
            }
        )
    )


if __name__ == "__main__":
    show(calculate(json.load(sys.stdin)["stat"]), Unit.MB, 2)