From a6cd4eabf30ec25fb4ddc6fa66ebf512344d0aae Mon Sep 17 00:00:00 2001 From: root Date: Tue, 7 Jun 2022 15:43:57 +0000 Subject: [PATCH] Added stat command --- bot.py | 68 ++++++++++++++++++++++--- gen.py | 126 ++++++++++++++++++++++++++++++++++++++++++----- requirements.txt | 1 + 3 files changed, 178 insertions(+), 17 deletions(-) diff --git a/bot.py b/bot.py index 9e558fb..9e81306 100755 --- a/bot.py +++ b/bot.py @@ -6,13 +6,17 @@ import logging import os import sys import configparser +from hurry.filesize import size +from subprocess import call from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram.ext import Updater, MessageHandler, CommandHandler, filters, CallbackQueryHandler, CallbackContext +from gen import wg_json from gen import add_peer as wg_add_peer from gen import update_configs from gen import list_peers as wg_list_peers from gen import del_peer as wg_del_peer +tg_max_len = 4096 logging.basicConfig( level=logging.INFO, @@ -30,13 +34,18 @@ else: config = "wg0" def _help(update, context): - update.message.reply_text('Help:\n * /add peer name\n * /del peer name\n * /list [peer name]', parse_mode='HTML', disable_web_page_preview=True) + update.message.reply_text( + 'Help:\n * /add peer name\n * /del peer name\n * /list [peer name]', + parse_mode='HTML', + disable_web_page_preview=True) def auth(handler): def wrapper(update, context): if update.message.chat.username not in admin: update.message.reply_text( - 'You are not allowed to do that.', parse_mode='HTML', disable_web_page_preview=True) + 'You are not allowed to do that.', + parse_mode='HTML', + disable_web_page_preview=True) return False handler(update, context) return wrapper @@ -49,7 +58,10 @@ def list_peers(update, context): for peer in wg_list_peers(): message += f"{n} * {peer['ip']}: {peer['name']}\n" n += 1 - update.message.reply_text(f"{message}", parse_mode='HTML', disable_web_page_preview=True) + update.message.reply_text( + f"{message}", + parse_mode='HTML', + disable_web_page_preview=True) else: peer_name = "_".join(update.message.text.split()[1:]) if peer_name.isnumeric(): @@ -61,8 +73,14 @@ def list_peers(update, context): n += 1 try: msg = open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'r').read() - update.message.reply_photo(open(f'/etc/wireguard/clients_{config}/{peer_name}-qr.png', 'rb'), parse_mode='HTML', filename=f'{peer_name} QR.png', quote=True, caption=f"Install Wireguard VPN app and scan or open config.\n{msg}") - update.message.reply_document(open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'rb')) + update.message.reply_photo( + open(f'/etc/wireguard/clients_{config}/{peer_name}-qr.png', 'rb'), + parse_mode='HTML', + filename=f'{peer_name} QR.png', + quote=True, + caption=f"Install Wireguard VPN app and scan or open config.\n{msg}") + update.message.reply_document( + open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'rb')) except: update.message.reply_text("Wrong client name.") @@ -76,6 +94,38 @@ def del_peer(update, context): wg_del_peer(peer_name) update.message.reply_text("Done.") +# {'preshared_key': 'kl0iFpC0jtqPwWaGbasVyQQsUgOYS9KyH917IwJ6Izs=', 'endpoint': '(none)', 'latest_handshake': 0, 'transfer_rx': 0, 'transfer_tx': 0, 'persistent_keepalive': 'off', 'allowed_ips': ['10.150.200.14/32']} +@auth +def status(update, context): + stat = wg_json() + msg = [] + for _if in stat.items(): + print(_if) + msg.append(f"{_if[0]}\nStarted {_if[1]['started']}") + peers = {} + for peer in _if[1]['peers']: + peers[peer['allowed_ips'][0]] = { + "tx": peer['transfer_rx'], + "rx": peer['transfer_tx'], + "total": peer['transfer_rx'] + peer['transfer_tx']} + 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[0]}\n Total {size(peer[1]['total'])} RX: {size(peer[1]['rx'])} TX: {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',) + msg = [] + msg.append(t_msg) + msg.append("Clients without any activity are skipped.") + msg = "\n".join(msg) + update.message.reply_text(f"{msg}", parse_mode='HTML',) + +@auth +def restart(update, context): + call(f"systemctl restart wg-quick@{config}.service", shell=True) + update.message.reply_text(f"Restarted {config} interface.") + @auth def add_peer(update, context): if len(update.message.text.split()) < 2: @@ -85,7 +135,11 @@ def add_peer(update, context): log.info("Creating peer %s", peer_name) wg_add_peer(peer_name) msg = open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'r').read() - update.message.reply_photo(open(f'/etc/wireguard/clients_{config}/{peer_name}-qr.png', 'rb'), parse_mode='HTML', filename=f'{peer_name} QR.png', quote=True, caption=f"Install Wireguard VPN app and scan or open config.\n{msg}") + update.message.reply_photo( + open(f'/etc/wireguard/clients_{config}/{peer_name}-qr.png', 'rb'), + parse_mode='HTML', + filename=f'{peer_name} QR.png', + quote=True, caption=f"Install Wireguard VPN app and scan or open config.\n{msg}") update.message.reply_document(open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'rb')) def error(update, context): @@ -97,6 +151,8 @@ def main(): updater.dispatcher.add_handler(CommandHandler('add', add_peer)) updater.dispatcher.add_handler(CommandHandler('list', list_peers)) updater.dispatcher.add_handler(CommandHandler('del', del_peer)) + updater.dispatcher.add_handler(CommandHandler('restart', restart)) + updater.dispatcher.add_handler(CommandHandler('status', status)) updater.dispatcher.add_handler(MessageHandler(filters.Filters.text, _help)) updater.start_polling() updater.idle() diff --git a/gen.py b/gen.py index e1c5ac8..8ea2fc6 100755 --- a/gen.py +++ b/gen.py @@ -4,11 +4,12 @@ import wgconfig # default iniparser cannot read WG configs. import logging -import json +import json as _json import ipaddress import argparse import configparser -from subprocess import call +from typing import TypedDict +from subprocess import call, Popen, PIPE from socket import getfqdn from os import path, mkdir from base64 import b64encode, b64decode @@ -25,10 +26,39 @@ log = logging.getLogger('generator') my_parser = argparse.ArgumentParser() # Add the arguments -my_parser.add_argument('--update', action='store_true', default=False) -my_parser.add_argument('--peer', action='store', type=str) -my_parser.add_argument('--delete', action='store', type=str) -my_parser.add_argument('--config', action='store', default='wg0', type=str) +args = { + "--update": { + "action": "store_true", + "default": False, + "desc": "Regenerate all client configs" + }, + "--json": { + "action": "store_true", + "default": False, + "desc": "Print all Wireguard statistics in JSON" + }, + "--peer": { + "action": "store", + "default": None, + "desc": "Add new peer" + }, + "--delete": { + "action": "store", + "default": None, + "desc": "Delete peer" + }, + "--config": { + "action": "store", + "default": "wg0", + "desc": "Config to use, default wg0" + }, +} +for arg in args.items(): + my_parser.add_argument(arg[0], action=arg[1]['action'], default=arg[1]['default']) + +help_msg = "" +for arg in args.items(): + help_msg += f" {arg[0]}\t{arg[1]['desc']}\n" ## Reading config # Execute the parse_args() method @@ -36,6 +66,7 @@ args = my_parser.parse_args() peer_name = args.peer del_name = args.delete is_update = args.update +json = args.json wpm_config = configparser.ConfigParser() client_dir = f"/etc/wireguard/clients_{args.config}" if not path.isdir(client_dir): @@ -51,9 +82,79 @@ else: dns = '8.8.8.8' hostname = getfqdn() config = args.config -log.info('Using %s WG config file.', config) +log.debug('Using %s WG config file.', config) +class WG_peer(TypedDict): + preshared_key: str + endpoint: str + latest_handshake: int + transfer_rx: int + transfer_tx: int + persistent_keepalive: bool + allowed_ips: list + + +class Interface(TypedDict): + name: str + private_key: str + public_key: str + listen_port: int + fwmark: str + peers: list + started: str + + +class Wireguard(TypedDict): + interfaces: dict + + +wg_state = Wireguard({}) + + +def wg_json(): + cmd = ["/usr/bin/wg", "show", "all", "dump"] + proc = Popen(cmd, + stdout=PIPE, + stderr=PIPE, + universal_newlines=True + ) + stdout, stderr = proc.communicate() + + for v in stdout.split('\n'): + cmd = ["systemctl", "show", "wg-quick@wg0", "--property", "InactiveEnterTimestamp"] + proc = Popen(cmd, + stdout=PIPE, + stderr=PIPE, + universal_newlines=True + ) + stdout, stderr = proc.communicate() + args = v.split('\t') + if len(args) == 5: + interface = Interface( + name=args[0], + private_key=args[1], + public_key=args[2], + listen_port=args[3], + fwmark=args[4], + started=stdout.strip().split("=")[1], + peers=[]) + wg_state[interface['name']] = interface + elif len(args) == 9: + allowed_ips = args[4].replace(' ', '').split(',') + peer = WG_peer( + preshared_key=args[1], + endpoint=args[3], + latest_handshake=int(args[5]), + transfer_rx=int(args[6]), + transfer_tx=int(args[7]), + persistent_keepalive=args[8], + allowed_ips=allowed_ips) + wg_state[args[0]]['peers'].append(peer) + else: + pass + #return _json.dumps(wg_state) + return wg_state class Peer: def __init__(self, peer=None, allowed_ips=None, comment='None'): @@ -201,9 +302,12 @@ def list_peers(): if __name__ == '__main__': if del_name: del_peer(del_name) - - if not is_update and peer_name: + elif not is_update and peer_name: add_peer(peer_name) - - if is_update: + elif is_update: update_configs() + elif json: + #print(_json.dumps(wg_json())) + print(wg_json()['wg0']['peers'][0]) + else: + print(help_msg) diff --git a/requirements.txt b/requirements.txt index 9654b62..e5cf52a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ wgconfig==0.2.2 python-telegram-bot==13.4 PyNaCl==1.4.0 +hurry.filesize==0.9