From ea52529e063ae0d3ca9ae77b37ebeaf059cba2ba Mon Sep 17 00:00:00 2001
From: donaldzou <donaldzou@live.hk>
Date: Fri, 23 Oct 2020 01:31:10 -0400
Subject: [PATCH] Able to add peer

---
 .vscode/sftp.json                |  9 ++++
 src/dashboard.py                 | 44 ++++++++++++++++--
 src/json/conf.json               |  0
 src/templates/configuration.html | 78 ++++++++++++++++++++++++++++----
 src/templates/index.html         | 23 +++++++---
 src/test.py                      |  9 ++--
 6 files changed, 136 insertions(+), 27 deletions(-)
 create mode 100644 .vscode/sftp.json
 create mode 100644 src/json/conf.json

diff --git a/.vscode/sftp.json b/.vscode/sftp.json
new file mode 100644
index 0000000..30290fe
--- /dev/null
+++ b/.vscode/sftp.json
@@ -0,0 +1,9 @@
+{
+    "name": "My Server",
+    "host": "38.106.21.18",
+    "protocol": "sftp",
+    "port": 22,
+    "username": "root",
+    "remotePath": "/root/wireguard-dashboard",
+    "uploadOnSave": true
+}
diff --git a/src/dashboard.py b/src/dashboard.py
index 6cee06d..f187e43 100644
--- a/src/dashboard.py
+++ b/src/dashboard.py
@@ -2,12 +2,30 @@ import os
 from flask import Flask, request, render_template, redirect, url_for
 import subprocess
 from datetime import datetime, date, time, timedelta
+from tinydb import TinyDB, Query
+import time
+import requests
+
 
 conf_location = "/etc/wireguard"
 
 app = Flask("Wireguard Dashboard")
 app.config['TEMPLATES_AUTO_RELOAD'] = True
 css = ""
+conf_data = {}
+
+
+def get_conf_peer_key(config_name):
+    keys = []
+    try: peer_key = subprocess.check_output("wg show "+config_name+" peers", shell=True)
+    except Exception: return "stopped"
+    peer_key = peer_key.decode("UTF-8").split()
+    for i in peer_key: keys.append(i)
+    return keys
+
+
+
+
 
 def get_conf_peers_data(config_name):
     peer_data = {}
@@ -26,9 +44,9 @@ def get_conf_peers_data(config_name):
     download_total = 0
     total = 0
     for i in range(int(len(data_usage)/3)):
-        peer_data[data_usage[count]]['total_recive'] = round(int(data_usage[count+1])/(1024**3), 4)
-        peer_data[data_usage[count]]['total_sent'] = round(int(data_usage[count+2])/(1024**3),4)
-        peer_data[data_usage[count]]['total_data'] = round((int(data_usage[count+2])+int(data_usage[count+1]))/(1024**3),4)
+        peer_data[data_usage[count]]['total_recive'] = int(data_usage[count+1])/(1024**3)
+        peer_data[data_usage[count]]['total_sent'] = int(data_usage[count+2])/(1024**3)
+        peer_data[data_usage[count]]['total_data'] = (int(data_usage[count+2])+int(data_usage[count+1]))/(1024**3)
         count += 3
 
     #Get endpoint
@@ -49,7 +67,6 @@ def get_conf_peers_data(config_name):
     b = timedelta(minutes=2)
     for i in range(int(len(data_usage)/2)):
         minus = now - datetime.fromtimestamp(int(data_usage[count+1]))
-        
         if minus < b:
             peer_data[data_usage[count]]['status'] = "running"
         else:
@@ -137,6 +154,7 @@ def conf(config_name):
         "checked": ""
     }
     if conf_data['status'] == "stopped":
+        print(conf_data)
         return redirect('/')
     else:
         conf_data['checked'] = "checked"
@@ -155,5 +173,21 @@ def switch(config_name):
     return redirect('/')
 
 
+@app.route('/add_peer/<config_name>', methods=['POST'])
+def add_peer(config_name):
+    data = request.get_json()
+    public_key = data['public_key']
+    allowed_ips = data['allowed_ips']
+    keys = get_conf_peer_key(config_name)
+    if public_key in keys:
+        return "Key already exist."
+    else:
+        try: 
+            status = subprocess.check_output("wg set "+config_name+" peer "+public_key+" allowed-ips "+allowed_ips, shell=True)
+            status = subprocess.check_output("wg-quick save "+config_name, shell=True)
+            return "Good"
+        except Exception: return redirect('/configuration/'+config_name)
+        
 
-app.run(host='0.0.0.0',debug=False, port=10086)
+
+app.run(host='0.0.0.0',debug=False, port=10086)
\ No newline at end of file
diff --git a/src/json/conf.json b/src/json/conf.json
new file mode 100644
index 0000000..e69de29
diff --git a/src/templates/configuration.html b/src/templates/configuration.html
index cb0a4c9..5c79481 100644
--- a/src/templates/configuration.html
+++ b/src/templates/configuration.html
@@ -40,7 +40,6 @@
 		</div>
 		<main role="main" class="col-md-9 ml-sm-auto col-lg-10 px-md-4 mt-4">
 			<div class="info mt-4">
-				
 				<div class="row">
 					<div class="col-sm">
 						<small class="text-muted"><strong>CONFIGURATION</strong></small>
@@ -48,7 +47,7 @@
 					</div>
 					<div class="col-sm">
 						<small class="text-muted"><strong>ACTION</strong></small><br>
-						<input class="mt-2 switch" id="{{conf_data['name']}}" type="checkbox" data-toggle="toggle"  {{conf_data['checked']}} data-height="15">
+						<input class="mt-2 switch" id="{{conf_data['name']}}" type="checkbox" data-toggle="toggle"  {{conf_data['checked']}} data-size="sm">
 					</div>
 					<div class="w-100"></div>
 					<div class="col-sm">
@@ -78,6 +77,12 @@
 					</div>
 				</div>
 				<hr>
+				<div class="button-div" style="text-align: right;">
+					<button type="button" class="btn btn-outline-primary btn-sm mb-3" data-toggle="modal" data-target="#staticBackdrop">
+						<strong>ADD PEER</strong>
+					</button>
+					
+				</div>
 			</div>
 			
 			{% for i in conf_data['peer_data']%}
@@ -135,9 +140,37 @@
 			{%endfor%}
 		</main>
 	</div>
+	<div class="modal fade" id="staticBackdrop" data-backdrop="static" data-keyboard="false" tabindex="-1"
+		aria-labelledby="staticBackdropLabel" aria-hidden="true">
+		<div class="modal-dialog modal-dialog-centered">
+			<div class="modal-content">
+				<div class="modal-header">
+					<h5 class="modal-title" id="staticBackdropLabel">Add a new peer</h5>
+					<button type="button" class="close" data-dismiss="modal" aria-label="Close">
+						<span aria-hidden="true">&times;</span>
+					</button>
+				</div>
+				<div class="modal-body">
+					<form>
+						<div class="form-group">
+							<label for="public_key">Public Key</label>
+							<input type="text" class="form-control" id="public_key" aria-describedby="public_key">
+						</div>
+						<div class="form-group">
+							<label for="allowed_ips">Allowed IPs</label>
+							<input type="text" class="form-control" id="allowed_ips">
+						</div>
+					</form>
+				</div>
+				<div class="modal-footer">
+					<button type="button" class="btn btn-secondary" data-dismiss="modal">Cancel</button>
+					<button type="button" class="btn btn-primary" id="save_peer" conf_id={{conf_data['name']}}>Save</button>
+				</div>
+			</div>
+		</div>
+	</div>
 </body>
-<script src="https://code.jquery.com/jquery-3.5.1.min.js"
-	integrity="sha256-9/aliU8dGd2tb6OSsuzixeV4y/faTqgFtohetphbbj0=" crossorigin="anonymous"></script>
+<script src="https://code.jquery.com/jquery-3.5.1.min.js" integrity="sha256-9/aliU8dGd2tb6OSsuzixeV4y/faTqgFtohetphbbj0=" crossorigin="anonymous"></script>
 <script src="https://cdn.jsdelivr.net/npm/popper.js@1.16.1/dist/umd/popper.min.js"
 	integrity="sha384-9/reFTGAW83EW2RDu2S0VKaIzap3H66lZH81PoYlFhbGU+6BZp6G7niu735Sk7lN"
 	crossorigin="anonymous"></script>
@@ -147,12 +180,37 @@
 <script src="https://cdn.jsdelivr.net/gh/gitbrent/bootstrap4-toggle@3.6.1/js/bootstrap4-toggle.min.js"></script>
 <script>
 	$('.switch').change(function() {
-		location.replace("/switch/"+$(this).attr('id'))
-    })
+		if ($(this).prop('checked') == false){
+			if (confirm('Are you sure you want to turn off this connection?')){
+				location.replace("/switch/"+$(this).attr('id'))
+			}
+		}
+		else{
+			location.replace("/switch/"+$(this).attr('id'))
+		}
+    });
 
+	$("#save_peer").click(function(){
+		if ($("#allowed_ips") != "" && $("#public_key") != ""){
+			var conf = $(this).attr('conf_id')
 
-	setInterval(function(){
-		location.reload();
-	}, 10000)
+			$.ajax({
+				method: "POST",
+				url: "/add_peer/"+conf,
+				headers:{
+					"Content-Type": "application/json"
+				},
+				data: JSON.stringify({"public_key":$("#public_key").val(), "allowed_ips": $("#allowed_ips").val()}),
+				success: function (response){
+					console.log(response);
+				}
+			})
+		}
+		
+	})
+
+	// setInterval(function(){
+	// 	location.reload();
+	// }, 10000)
 </script>
-</html>
\ No newline at end of file
+</html>	
\ No newline at end of file
diff --git a/src/templates/index.html b/src/templates/index.html
index d4170e1..8e85099 100644
--- a/src/templates/index.html
+++ b/src/templates/index.html
@@ -43,12 +43,15 @@
 					<div class="card-body">
 						
 						<div class="row">
-							<div class="col-sm">
+							<div class="col">
+								<small class="text-muted"><strong>CONFIGURATION</strong></small>
 								<a href="/configuration/{{i['conf']}}">
 									<h5 class="card-title">{{i['conf']}}</h5>
 								</a>
 							</div>
-							
+							<div class="col">
+								<input class="mt-2 switch" id="{{i['conf']}}"  type="checkbox" data-toggle="toggle"  {{i['checked']}} data-size="sm">
+							</div>
 							
 							<div class="w-100"></div>
 							<div class="col-sm">
@@ -59,9 +62,8 @@
 								<small class="text-muted"><strong>PUBLIC KEY</strong></small>
 								<h6 style="text-transform: uppercase;"><samp>{{i['public_key']}}</samp></h6>
 							</div>
-							<div class="col-sm">
-								<input class="mt-2 switch" id="{{i['conf']}}"  type="checkbox" data-toggle="toggle"  {{i['checked']}} data-height="15">
-							</div>
+							
+
 						</div>
 					</div>
 				</div>
@@ -80,7 +82,14 @@
 <script src="https://cdn.jsdelivr.net/gh/gitbrent/bootstrap4-toggle@3.6.1/js/bootstrap4-toggle.min.js"></script>
 <script>
 	$('.switch').change(function() {
-		location.replace("/switch/"+$(this).attr('id'))
-    })
+		if ($(this).prop('checked') == false){
+			if (confirm('Are you sure you want to turn off this connection?')){
+				location.replace("/switch/"+$(this).attr('id'))
+			}
+		}
+		else{
+			location.replace("/switch/"+$(this).attr('id'))
+		}
+    });
 </script>
 </html>
\ No newline at end of file
diff --git a/src/test.py b/src/test.py
index bbca88c..03ecd9e 100644
--- a/src/test.py
+++ b/src/test.py
@@ -1,5 +1,4 @@
-import os
-import subprocess
-try: status = subprocess.check_output("wg show wg0", shell=True)
-except Exception: print("false")
-else:print(status.decode("UTF-8"))
\ No newline at end of file
+from tinydb import TinyDB, Query
+conf_db = TinyDB("json/conf.json")
+
+print(conf_db.all())
\ No newline at end of file