diff --git a/.gitignore b/.gitignore index 60452cd..5e52ef1 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ __pycache__/ clients/ server.key server.pub +stats.bin *.conf *.conf_old !wpm_example.conf diff --git a/bot.py b/bot.py index 792b08a..0517541 100755 --- a/bot.py +++ b/bot.py @@ -6,6 +6,9 @@ import logging import os import sys import configparser +import threading +import pickle +from time import sleep from hurry.filesize import size from subprocess import call from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update @@ -39,6 +42,36 @@ def _help(update, context): parse_mode='HTML', disable_web_page_preview=True) + +def save_stats(): + while True: + stat = wg_json(config) + peer_names = dict() + peers = dict() + for peer in wg_list_peers(): + peer_names[peer['ip']] = peer['name'] + for _if in stat.items(): + for peer in _if[1]['peers']: + peers[peer['allowed_ips'][0]] = { + "name": peer_names[peer['allowed_ips'][0]], + "rx": int(peer["transfer_rx"]), + "tx": int(peer["transfer_tx"]), + } + try: + read_stats = pickle.load(open("stats.bin", "rb")) + except FileNotFoundError: + read_stats = None + if read_stats: + for peer in peers: + if read_stats[peer]["rx"] >= peers[peer]["rx"]: + peers[peer]["rx"] += read_stats[peer]["rx"] + peers[peer]["tx"] += read_stats[peer]["tx"] + pickle.dump(peers, open("stats.bin", "wb")) + sleep(60) + +save_stat_thread = threading.Thread(target=save_stats) +save_stat_thread.start() + def auth(handler): def wrapper(update, context): if update.message.chat.username not in admin: @@ -97,6 +130,7 @@ def del_peer(update, context): @auth def status(update, context): stat = wg_json(config) + overall_stats = pickle.load(open("stats.bin", "rb")) peer_names = dict() for peer in wg_list_peers(): peer_names[peer['ip']] = peer['name'] @@ -113,7 +147,12 @@ def status(update, context): peers_sorted = sorted(peers.items(), key=lambda x: x[1]['total'], reverse=True) peers_sorted = list(filter(lambda x: (x[1]['total'] != 0), peers_sorted)) for peer in peers_sorted: - t_msg = f" * {peer[1]['name']}\n {peer[0]}\n Total {size(peer[1]['total'])} RX: {size(peer[1]['rx'])} TX: {size(peer[1]['tx'])}" + t_msg = ( + f" 🔹 {peer[1]['name']}\n {peer[0]}\n " + f"Overall stats {size(overall_stats[peer[0]]['rx'] + overall_stats[peer[0]]['tx'])} " + f"RX/TX: {size(overall_stats[peer[0]]['rx'])}/{size(overall_stats[peer[0]]['tx'])}\n " + f"Since last run {size(peer[1]['total'])} RX/TX: {size(peer[1]['rx'])}" + f"/{size(peer[1]['tx'])}") if len(t_msg + "\n".join(msg)) >= tg_max_len: msg = "\n".join(msg) update.message.reply_text(f"{msg}", parse_mode='HTML',)