From 037a6c56d3f7ce93a2e2a5f73a00f6ac4f7fbf5f Mon Sep 17 00:00:00 2001
From: Maxim Kochurov <maxim.v.kochurov@gmail.com>
Date: Sun, 13 Mar 2022 19:33:37 +0300
Subject: [PATCH] Implement Optional Private Keys (#161)

---
 custom/js/helper.js    |  2 +-
 handler/routes.go      | 73 +++++++++++++++++++++++++++++++-----------
 store/jsondb/jsondb.go |  4 +--
 templates/base.html    | 22 ++++++++++++-
 4 files changed, 78 insertions(+), 23 deletions(-)

diff --git a/custom/js/helper.js b/custom/js/helper.js
index 6fe8847e..dd4b4974 100644
--- a/custom/js/helper.js
+++ b/custom/js/helper.js
@@ -31,7 +31,7 @@ function renderClientList(data) {
                                 <div class="btn-group">      
                                     <button type="button" class="btn btn-outline-secondary btn-sm" data-toggle="modal"
                                         data-target="#modal_qr_client" data-clientid="${obj.Client.id}"
-                                        data-clientname="${obj.Client.name}">Scan</button>
+                                        data-clientname="${obj.Client.name}" ${obj.QRCode != "" ? '' : ' disabled'}>Scan</button>
                                 </div>
                                 <div class="btn-group">      
                                     <button type="button" class="btn btn-outline-secondary btn-sm" data-toggle="modal"
diff --git a/handler/routes.go b/handler/routes.go
index 9c18fc40..da7ecd2f 100644
--- a/handler/routes.go
+++ b/handler/routes.go
@@ -171,23 +171,51 @@ func NewClient(db store.IStore) echo.HandlerFunc {
 		client.ID = guid.String()
 
 		// gen Wireguard key pair
-		key, err := wgtypes.GeneratePrivateKey()
-		if err != nil {
-			log.Error("Cannot generate wireguard key pair: ", err)
-			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot generate Wireguard key pair"})
-		}
+		if client.PublicKey == "" {
+			key, err := wgtypes.GeneratePrivateKey()
+			if err != nil {
+				log.Error("Cannot generate wireguard key pair: ", err)
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot generate Wireguard key pair"})
+			}
+			client.PrivateKey = key.String()
+			client.PublicKey = key.PublicKey().String()
+		} else {
+			_, err := wgtypes.ParseKey(client.PublicKey)
+			if err != nil {
+				log.Error("Cannot verify wireguard public key: ", err)
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot verify Wireguard public key"})
+			}
+			// check for duplicates
+			clients, err := db.GetClients(false)
+			if err != nil {
+				log.Error("Cannot get clients for duplicate check")
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot get clients for duplicate check"})
+			}
+			for _, other := range clients {
+				if other.Client.PublicKey == client.PublicKey {
+					log.Error("Duplicate Public Key")
+					return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Duplicate Public Key"})
+				}
+			}
 
-		presharedKey, err := wgtypes.GenerateKey()
-		if err != nil {
-			log.Error("Cannot generated preshared key: ", err)
-			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{
-				false, "Cannot generate Wireguard preshared key",
-			})
 		}
 
-		client.PrivateKey = key.String()
-		client.PublicKey = key.PublicKey().String()
-		client.PresharedKey = presharedKey.String()
+		if client.PresharedKey == "" {
+			presharedKey, err := wgtypes.GenerateKey()
+			if err != nil {
+				log.Error("Cannot generated preshared key: ", err)
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{
+					false, "Cannot generate Wireguard preshared key",
+				})
+			}
+			client.PresharedKey = presharedKey.String()
+		} else {
+			_, err := wgtypes.ParseKey(client.PresharedKey)
+			if err != nil {
+				log.Error("Cannot verify wireguard preshared key: ", err)
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "Cannot verify Wireguard preshared key"})
+			}
+		}
 		client.CreatedAt = time.Now().UTC()
 		client.UpdatedAt = client.CreatedAt
 
@@ -227,18 +255,25 @@ func EmailClient(db store.IStore, mailer emailer.Emailer, emailSubject, emailCon
 		config := util.BuildClientConfig(*clientData.Client, server, globalSettings)
 
 		cfg_att := emailer.Attachment{"wg0.conf", []byte(config)}
-		qrdata, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(clientData.QRCode, "data:image/png;base64,"))
-		if err != nil {
-			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "decoding: " + err.Error()})
+		var attachments []emailer.Attachment
+		if clientData.Client.PrivateKey != "" {
+			qrdata, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(clientData.QRCode, "data:image/png;base64,"))
+			if err != nil {
+				return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "decoding: " + err.Error()})
+			}
+			qr_att := emailer.Attachment{"wg.png", qrdata}
+			attachments = []emailer.Attachment{cfg_att, qr_att}
+		} else {
+			attachments = []emailer.Attachment{cfg_att}
 		}
-		qr_att := emailer.Attachment{"wg.png", qrdata}
 		err = mailer.Send(
 			clientData.Client.Name,
 			payload.Email,
 			emailSubject,
 			emailContent,
-			[]emailer.Attachment{cfg_att, qr_att},
+			attachments,
 		)
+
 		if err != nil {
 			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})
 		}
diff --git a/store/jsondb/jsondb.go b/store/jsondb/jsondb.go
index 763756ef..4ee35cac 100644
--- a/store/jsondb/jsondb.go
+++ b/store/jsondb/jsondb.go
@@ -155,7 +155,7 @@ func (o *JsonDB) GetClients(hasQRCode bool) ([]model.ClientData, error) {
 		}
 
 		// generate client qrcode image in base64
-		if hasQRCode {
+		if hasQRCode && client.PrivateKey != "" {
 			server, _ := o.GetServer()
 			globalSettings, _ := o.GetGlobalSettings()
 
@@ -185,7 +185,7 @@ func (o *JsonDB) GetClientByID(clientID string, hasQRCode bool) (model.ClientDat
 	}
 
 	// generate client qrcode image in base64
-	if hasQRCode {
+	if hasQRCode && client.PrivateKey != "" {
 		server, _ := o.GetServer()
 		globalSettings, _ := o.GetGlobalSettings()
 
diff --git a/templates/base.html b/templates/base.html
index b3c254cd..acb0fe3d 100644
--- a/templates/base.html
+++ b/templates/base.html
@@ -184,6 +184,21 @@ <h4 class="modal-title">New Wireguard Client</h4>
                                     </label>
                                 </div>
                             </div>
+                            <details>
+                                <summary>Public and Preshared Keys</summary>
+                                <div class="form-group" style="margin-top: 1rem">
+                                    <label for="client_public_key" class="control-label">
+                                        Public Key
+                                    </label>
+                                    <input type="text" class="form-control" id="client_public_key" name="client_public_key" placeholder="Autogenerated (insecure)" aria-invalid="false">
+                                </div>
+                                <div class="form-group">
+                                    <label for="client_preshared_key" class="control-label">
+                                        Preshared Key
+                                    </label>
+                                    <input type="text" class="form-control" id="client_preshared_key" name="client_preshared_key" placeholder="Autogenerated">
+                                </div>
+                            </details>
                         </div>
                         <div class="modal-footer justify-content-between">
                             <button type="button" class="btn btn-default" data-dismiss="modal">Cancel</button>
@@ -314,9 +329,12 @@ <h1>{{template "page_title" .}}</h1>
             if ($("#enabled").is(':checked')){
                 enabled = true;
             }
+            const public_key = $("#client_public_key").val();
+            const preshared_key = $("#client_preshared_key").val();
 
             const data = {"name": name, "email": email, "allocated_ips": allocated_ips, "allowed_ips": allowed_ips,
-                "extra_allowed_ips": extra_allowed_ips, "use_server_dns": use_server_dns, "enabled": enabled};
+                "extra_allowed_ips": extra_allowed_ips, "use_server_dns": use_server_dns, "enabled": enabled,
+                "public_key": public_key, "preshared_key": preshared_key};
 
             $.ajax({
                 cache: false,
@@ -434,6 +452,8 @@ <h1>{{template "page_title" .}}</h1>
             $("#modal_new_client").on('shown.bs.modal', function (e) {
                 $("#client_name").val("");
                 $("#client_email").val("");
+                $("#client_public_key").val("");
+                $("#client_preshared_key").val("");
                 $("#client_allocated_ips").importTags('');
                 $("#client_extra_allowed_ips").importTags('');
                 updateIPAllocationSuggestion();