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