Added stat command

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('<b>Help:</b>\n <b>*</b> /add <i>peer name</i>\n <b>*</b> /del <i>peer name</i>\n <b>*</b> /list [<i>peer name</i>]', parse_mode='HTML', disable_web_page_preview=True)
+    update.message.reply_text(
+        '<b>Help:</b>\n <b>*</b> /add <i>peer name</i>\n <b>*</b> /del <i>peer name</i>\n <b>*</b> /list [<i>peer name</i>]',
+        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}</code>", parse_mode='HTML', disable_web_page_preview=True)
+        update.message.reply_text(
+            f"{message}</code>",
+            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<code>{msg}</code>")
+            update.message.reply_photo(
-            update.message.reply_document(open(f'/etc/wireguard/clients_{config}/{peer_name}.conf', 'rb'))
+                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<code>{msg}</code>")
+            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"<b>{_if[0]}</b>\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" * <b>{peer[0]}</b>\n    <b>Total</b> {size(peer[1]['total'])}<b> RX</b>: {size(peer[1]['rx'])} <b>TX</b>: {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("<b>Clients without any activity are skipped.</b>")
+    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<code>{msg}</code>")
+    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<code>{msg}</code>")
     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()

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)
+args = {
-my_parser.add_argument('--peer', action='store', type=str)
+    "--update": {
-my_parser.add_argument('--delete', action='store', type=str)
+        "action": "store_true",
-my_parser.add_argument('--config', action='store', default='wg0', type=str)
+        "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)
-
+    elif not is_update and peer_name:
-    if not is_update and peer_name:
         add_peer(peer_name)
-
+    elif is_update:
-    if is_update:
         update_configs()
+    elif json:
+    	#print(_json.dumps(wg_json()))
+    	print(wg_json()['wg0']['peers'][0])
+    else:
+        print(help_msg)

requirements.txt

@@ -1,3 +1,4 @@
 wgconfig==0.2.2
 python-telegram-bot==13.4
 PyNaCl==1.4.0
+hurry.filesize==0.9