mirror of
https://github.com/house-of-vanity/Wireguard-Peer-Manager.git
synced 2025-07-06 09:14:07 +00:00
Added stat command
This commit is contained in:
root
68
bot.py
68
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_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<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()
|
||||
|
||||
126
gen.py
126
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)
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
wgconfig==0.2.2
|
||||
python-telegram-bot==13.4
|
||||
PyNaCl==1.4.0
|
||||
hurry.filesize==0.9
|
||||
|
||||
Reference in New Issue
Block a user