v2.0-beta-6 Commit

src/dashboard.py

@@ -1,9 +1,7 @@
-dashboard_version = 'v2.0'
-
-
 # Python Built-in Library
 import os
-from flask import Flask, request, render_template, redirect, url_for, session, abort
+from flask import Flask, request, render_template, redirect, url_for, session, abort, jsonify
+from icmplib import ping, multiping, traceroute, resolve, Host, Hop
 import subprocess
 from datetime import datetime, date, time, timedelta
 from operator import itemgetter
@@ -11,28 +9,30 @@ import secrets
 import hashlib
 import json, urllib.request
 import configparser
+import re
 # PIP installed library
 import ifcfg
 from tinydb import TinyDB, Query
+
+# Dashboard Version
+dashboard_version = 'v2.0.1'
+# Dashboard Config Name
 dashboard_conf = 'wg-dashboard.ini'
+# Upgrade Required
 update = ""
+# Flask App Configuration
 app = Flask("Wireguard Dashboard")
 app.secret_key = secrets.token_urlsafe(16)
 app.config['TEMPLATES_AUTO_RELOAD'] = True
-conf_data = {}
-
 
 
 def get_conf_peer_key(config_name):
-    keys = []
     try:
         peer_key = subprocess.check_output("wg show " + config_name + " peers", shell=True)
+        peer_key = peer_key.decode("UTF-8").split()
+        return peer_key
     except Exception:
-        return "stopped"
-    peer_key = peer_key.decode("UTF-8").split()
-    for i in peer_key: keys.append(i)
-    return keys
-
+        return config_name+" is not running."
 
 def get_conf_running_peer_number(config_name):
     running = 0
@@ -52,9 +52,14 @@ def get_conf_running_peer_number(config_name):
         count += 2
     return running
 
+
+def is_match(regex, text):
+    pattern = re.compile(regex)
+    return pattern.search(text) is not None
+
 def read_conf_file(config_name):
     # Read Configuration File Start
-    conf_location = wg_conf_path+"/" + config_name + ".conf"
+    conf_location = wg_conf_path + "/" + config_name + ".conf"
     f = open(conf_location, 'r')
     file = f.read().split("\n")
     conf_peer_data = {
@@ -63,31 +68,32 @@ def read_conf_file(config_name):
     }
     peers_start = 0
     for i in range(len(file)):
-        if file[i] == "[Peer]":
-            peers_start = i
-            break
-        else:
-            if len(file[i]) > 0:
-                if file[i] != "[Interface]":
-                    tmp = file[i].replace(" ", "").split("=", 1)
-                    if len(tmp) == 2:
-                        conf_peer_data['Interface'][tmp[0]] = tmp[1]
+        if not is_match("^#(.*)",file[i]):
+            if file[i] == "[Peer]":
+                peers_start = i
+                break
+            else:
+                if len(file[i]) > 0:
+                    if file[i] != "[Interface]":
+                        tmp = re.split(r'\s*=\s*', file[i], 1)
+                        if len(tmp) == 2:
+                            conf_peer_data['Interface'][tmp[0]] = tmp[1]
     conf_peers = file[peers_start:]
     peer = -1
     for i in conf_peers:
-        if i == "[Peer]":
-            peer += 1
-            conf_peer_data["Peers"].append({})
-        else:
-            if len(i) > 0:
-                tmp = i.replace(" ", "").split("=", 1)
-                if len(tmp) == 2:
-                    conf_peer_data["Peers"][peer][tmp[0]] = tmp[1]
+        if not is_match("^#(.*)", i):
+            if i == "[Peer]":
+                peer += 1
+                conf_peer_data["Peers"].append({})
+            else:
+                if len(i) > 0:
+                    tmp = re.split('\s*=\s*', i,1)
+                    if len(tmp) == 2:
+                        conf_peer_data["Peers"][peer][tmp[0]] = tmp[1]
     # Read Configuration File End
     return conf_peer_data
 
 
-
 def get_conf_peers_data(config_name):
     db = TinyDB('db/' + config_name + '.json')
     peers = Query()
@@ -141,26 +147,25 @@ def get_conf_peers_data(config_name):
         cur_i = db.search(peers.id == data_usage[count])
         total_sent = cur_i[0]['total_sent']
         total_receive = cur_i[0]['total_receive']
+        traffic = cur_i[0]['traffic']
         cur_total_sent = round(int(data_usage[count + 2]) / (1024 ** 3), 4)
         cur_total_receive = round(int(data_usage[count + 1]) / (1024 ** 3), 4)
         if cur_i[0]["status"] == "running":
-            if total_sent <= cur_total_sent:
+            if total_sent <= cur_total_sent and total_receive <= cur_total_receive:
                 total_sent = cur_total_sent
-            else: total_sent += cur_total_sent
-
-            if total_receive <= cur_total_receive:
                 total_receive = cur_total_receive
-            else: total_receive += cur_total_receive
-            db.update({"total_receive": round(total_receive,4),
-                       "total_sent": round(total_sent,4),
+            else:
+                now = datetime.now()
+                ctime = now.strftime("%d/%m/%Y %H:%M:%S")
+                traffic.append({"time": ctime, "total_receive": round(total_receive, 4),"total_sent": round(total_sent, 4),
+                                "total_data": round(total_receive + total_sent, 4)})
+                total_sent = 0
+                total_receive = 0
+                db.update({"traffic": traffic}, peers.id == data_usage[count])
+            db.update({"total_receive": round(total_receive, 4),
+                       "total_sent": round(total_sent, 4),
                        "total_data": round(total_receive + total_sent, 4)}, peers.id == data_usage[count])
 
-        # Will get implement in the future
-        # traffic = db.search(peers.id == data_usage[count])[0]['traffic']
-        # traffic.append({"time": current_time, "total_receive": round(int(data_usage[count + 1]) / (1024 ** 3), 4),
-        #                 "total_sent": round(int(data_usage[count + 2]) / (1024 ** 3), 4)})
-        # db.update({"traffic": traffic}, peers.id == data_usage[count])
-
         count += 3
 
     # Get endpoint
@@ -176,7 +181,8 @@ def get_conf_peers_data(config_name):
 
     # Get allowed ip
     for i in conf_peer_data["Peers"]:
-        db.update({"allowed_ip":i.get('AllowedIPs', '(None)')}, peers.id == i["PublicKey"])
+        db.update({"allowed_ip": i.get('AllowedIPs', '(None)')}, peers.id == i["PublicKey"])
+    db.close()
 
 
 def get_peers(config_name):
@@ -184,6 +190,7 @@ def get_peers(config_name):
     db = TinyDB('db/' + config_name + '.json')
     result = db.all()
     result = sorted(result, key=lambda d: d['status'])
+    db.close()
     return result
 
 
@@ -209,9 +216,15 @@ def get_conf_total_data(config_name):
     upload_total = 0
     download_total = 0
     for i in db.all():
-        upload_total += round(i['total_sent'],4)
-        download_total += round(i['total_receive'],4)
+        upload_total += i['total_sent']
+        download_total += i['total_receive']
+        for k in i['traffic']:
+            upload_total += k['total_sent']
+            download_total += k['total_receive']
     total = round(upload_total + download_total, 4)
+    upload_total = round(upload_total, 4)
+    download_total = round(download_total, 4)
+    db.close()
     return [total, upload_total, download_total]
 
 
@@ -240,8 +253,6 @@ def get_conf_list():
     return conf
 
 
-
-
 @app.before_request
 def auth_req():
     conf = configparser.ConfigParser(strict=False)
@@ -259,9 +270,11 @@ def auth_req():
             session['message'] = "You need to sign in first!"
             return redirect(url_for("signin"))
     else:
-        if request.endpoint in ['signin', 'signout', 'auth', 'settings', 'update_acct', 'update_pwd', 'update_app_ip_port', 'update_wg_conf_path']:
+        if request.endpoint in ['signin', 'signout', 'auth', 'settings', 'update_acct', 'update_pwd',
+                                'update_app_ip_port', 'update_wg_conf_path']:
             return redirect(url_for("index"))
 
+
 @app.route('/signin', methods=['GET'])
 def signin():
     message = ""
@@ -279,7 +292,6 @@ def signout():
     return render_template('signin.html', message=message)
 
 
-
 @app.route('/settings', methods=['GET'])
 def settings():
     message = ""
@@ -292,14 +304,18 @@ def settings():
         session.pop("message")
         session.pop("message_status")
     required_auth = config.get("Server", "auth_req")
-    return render_template('settings.html',conf=get_conf_list(),message=message, status=status, app_ip=config.get("Server", "app_ip"), app_port=config.get("Server", "app_port"), required_auth=required_auth, wg_conf_path=config.get("Server", "wg_conf_path"))
+    return render_template('settings.html', conf=get_conf_list(), message=message, status=status,
+                           app_ip=config.get("Server", "app_ip"), app_port=config.get("Server", "app_port"),
+                           required_auth=required_auth, wg_conf_path=config.get("Server", "wg_conf_path"))
+
 
 @app.route('/auth', methods=['POST'])
 def auth():
     config = configparser.ConfigParser(strict=False)
     config.read(dashboard_conf)
     password = hashlib.sha256(request.form['password'].encode())
-    if password.hexdigest() == config["Account"]["password"] and request.form['username'] == config["Account"]["username"]:
+    if password.hexdigest() == config["Account"]["password"] and request.form['username'] == config["Account"][
+        "username"]:
         session['username'] = request.form['username']
         config.clear()
         return redirect(url_for("index"))
@@ -308,6 +324,7 @@ def auth():
         config.clear()
         return redirect(url_for("signin"))
 
+
 @app.route('/update_acct', methods=['POST'])
 def update_acct():
     config = configparser.ConfigParser(strict=False)
@@ -326,12 +343,14 @@ def update_acct():
         config.clear()
         return redirect(url_for("settings"))
 
+
 @app.route('/update_pwd', methods=['POST'])
 def update_pwd():
     config = configparser.ConfigParser(strict=False)
     config.read(dashboard_conf)
     if hashlib.sha256(request.form['currentpass'].encode()).hexdigest() == config.get("Account", "password"):
-        if hashlib.sha256(request.form['newpass'].encode()).hexdigest() ==  hashlib.sha256(request.form['repnewpass'].encode()).hexdigest():
+        if hashlib.sha256(request.form['newpass'].encode()).hexdigest() == hashlib.sha256(
+                request.form['repnewpass'].encode()).hexdigest():
             config.set("Account", "password", hashlib.sha256(request.form['repnewpass'].encode()).hexdigest())
             try:
                 config.write(open(dashboard_conf, "w"))
@@ -355,6 +374,7 @@ def update_pwd():
         config.clear()
         return redirect(url_for("settings"))
 
+
 @app.route('/update_app_ip_port', methods=['POST'])
 def update_app_ip_port():
     config = configparser.ConfigParser(strict=False)
@@ -365,6 +385,7 @@ def update_app_ip_port():
     config.clear()
     os.system('bash wgd.sh restart')
 
+
 @app.route('/update_wg_conf_path', methods=['POST'])
 def update_wg_conf_path():
     config = configparser.ConfigParser(strict=False)
@@ -376,15 +397,71 @@ def update_wg_conf_path():
     config.clear()
     os.system('bash wgd.sh restart')
 
-# @app.route('/check_update_dashboard', methods=['GET'])
-# def check_update_dashboard():
-#    return have_update
+@app.route('/update_dashboard_refresh_interval', methods=['POST'])
+def update_dashboard_refresh_interval():
+    config = configparser.ConfigParser(strict=False)
+    config.read(dashboard_conf)
+    config.set("Server", "dashboard_refresh_interval", str(request.form['interval']))
+    config.write(open(dashboard_conf, "w"))
+    config.clear()
+    return "true"
+
+@app.route('/get_ping_ip', methods=['POST'])
+def get_ping_ip():
+    config = request.form['config']
+    db = TinyDB('db/' + config + '.json')
+    html = ""
+    for i in db.all():
+        html += '<optgroup label="'+i['name']+' - '+i['id']+'">'
+        allowed_ip = str(i['allowed_ip']).split(",")
+        for k in allowed_ip:
+            k = k.split("/")
+            if len(k) == 2:
+                html += "<option value="+k[0]+">"+k[0]+"</option>"
+        endpoint = str(i['endpoint']).split(":")
+        if len(endpoint) == 2:
+            html += "<option value=" + endpoint[0] + ">" + endpoint[0] + "</option>"
+        html += "</optgroup>"
+    return html
+
+@app.route('/ping_ip', methods=['POST'])
+def ping_ip():
+    try:
+        result = ping(''+request.form['ip']+'', count=int(request.form['count']),privileged=True, source=None)
+        returnjson = {
+            "address": result.address,
+            "is_alive": result.is_alive,
+            "min_rtt": result.min_rtt,
+            "avg_rtt": result.avg_rtt,
+            "max_rtt": result.max_rtt,
+            "package_sent": result.packets_sent,
+            "package_received": result.packets_received,
+            "package_loss": result.packet_loss
+        }
+        return jsonify(returnjson)
+    except Exception:
+        return "Error"
+
+@app.route('/traceroute_ip', methods=['POST'])
+def traceroute_ip():
+    try:
+        result = traceroute(''+request.form['ip']+'', first_hop=1, max_hops=30, count=1, fast=True)
+        returnjson = []
+        last_distance = 0
+        for hop in result:
+            if last_distance + 1 != hop.distance:
+                returnjson.append({"hop":"*", "ip":"*", "avg_rtt":"", "min_rtt":"", "max_rtt":""})
+            returnjson.append({"hop": hop.distance, "ip": hop.address, "avg_rtt": hop.avg_rtt, "min_rtt": hop.min_rtt, "max_rtt": hop.max_rtt})
+            last_distance = hop.distance
+        return jsonify(returnjson)
+    except Exception:
+        return "Error"
 
 @app.route('/', methods=['GET'])
 def index():
-    print(request.referrer)
     return render_template('index.html', conf=get_conf_list())
 
+
 @app.route('/configuration/<config_name>', methods=['GET'])
 def conf(config_name):
     conf_data = {
@@ -396,12 +473,13 @@ def conf(config_name):
         conf_data['checked'] = "nope"
     else:
         conf_data['checked'] = "checked"
-    return render_template('configuration.html', conf=get_conf_list(), conf_data=conf_data)
+    config = configparser.ConfigParser(strict=False)
+    config.read(dashboard_conf)
+    return render_template('configuration.html', conf=get_conf_list(), conf_data=conf_data, dashboard_refresh_interval=int(config.get("Server","dashboard_refresh_interval")))
 
 
 @app.route('/get_config/<config_name>', methods=['GET'])
 def get_conf(config_name):
-    db = TinyDB('db/' + config_name + '.json')
     conf_data = {
         "peer_data": get_peers(config_name),
         "name": config_name,
@@ -445,6 +523,8 @@ def add_peer(config_name):
     public_key = data['public_key']
     allowed_ips = data['allowed_ips']
     keys = get_conf_peer_key(config_name)
+    if public_key is not list:
+        return config_name+" is not running."
     if public_key in keys:
         return "Key already exist."
     else:
@@ -454,23 +534,30 @@ def add_peer(config_name):
                 "wg set " + config_name + " peer " + public_key + " allowed-ips " + allowed_ips, shell=True,
                 stderr=subprocess.STDOUT)
             status = subprocess.check_output("wg-quick save " + config_name, shell=True, stderr=subprocess.STDOUT)
+            get_conf_peers_data(config_name)
+            db = TinyDB("db/" + config_name + ".json")
+            peers = Query()
+            db.update({"name": data['name']}, peers.id == public_key)
+            db.close()
             return "true"
         except subprocess.CalledProcessError as exc:
             return exc.output.strip()
-            # return redirect('/configuration/'+config_name)
 
 
 @app.route('/remove_peer/<config_name>', methods=['POST'])
 def remove_peer(config_name):
     if get_conf_status(config_name) == "stopped":
-        return "Your need to turn on "+config_name+" first."
+        return "Your need to turn on " + config_name + " first."
 
     db = TinyDB("db/" + config_name + ".json")
     peers = Query()
     data = request.get_json()
     delete_key = data['peer_id']
     keys = get_conf_peer_key(config_name)
+    if keys is not list:
+        return config_name+" is not running."
     if delete_key not in keys:
+        db.close()
         return "This key does not exist"
     else:
         try:
@@ -478,6 +565,7 @@ def remove_peer(config_name):
                                              stderr=subprocess.STDOUT)
             status = subprocess.check_output("wg-quick save " + config_name, shell=True, stderr=subprocess.STDOUT)
             db.remove(peers.id == delete_key)
+            db.close()
             return "true"
         except subprocess.CalledProcessError as exc:
             return exc.output.strip()
@@ -490,8 +578,8 @@ def save_peer_name(config_name):
     name = data['name']
     db = TinyDB("db/" + config_name + ".json")
     peers = Query()
-
     db.update({"name": name}, peers.id == id)
+    db.close()
     return id + " " + name
 
 
@@ -502,8 +590,12 @@ def get_peer_name(config_name):
     db = TinyDB("db/" + config_name + ".json")
     peers = Query()
     result = db.search(peers.id == id)
+    db.close()
     return result[0]['name']
 
+
+
+
 def init_dashboard():
     # Set Default INI File
     if not os.path.isfile("wg-dashboard.ini"):
@@ -530,9 +622,12 @@ def init_dashboard():
         config['Server']['auth_req'] = 'true'
     if 'version' not in config['Server'] or config['Server']['version'] != dashboard_version:
         config['Server']['version'] = dashboard_version
+    if 'dashboard_refresh_interval' not in config['Server']:
+        config['Server']['dashboard_refresh_interval'] = '15000'
     config.write(open(dashboard_conf, "w"))
     config.clear()
 
+
 def check_update():
     conf = configparser.ConfigParser(strict=False)
     conf.read(dashboard_conf)
@@ -554,4 +649,3 @@ if __name__ == "__main__":
     wg_conf_path = config.get("Server", "wg_conf_path")
     config.clear()
     app.run(host=app_ip, debug=False, port=app_port)
-