From 44961e91dc11a75b4658ca317af40796c681c462 Mon Sep 17 00:00:00 2001
From: HeshamTB <hishaminv@gmail.com>
Date: Mon, 18 Mar 2024 01:17:02 +0300
Subject: [PATCH] cli: allow generating private key on startup

Signed-off-by: HeshamTB <hishaminv@gmail.com>
---
 cmd/hvpn-node/hvpn-node.go | 56 ++++++++++++++++++++++++--------------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/cmd/hvpn-node/hvpn-node.go b/cmd/hvpn-node/hvpn-node.go
index 70a2abe..7d69ee5 100644
--- a/cmd/hvpn-node/hvpn-node.go
+++ b/cmd/hvpn-node/hvpn-node.go
@@ -114,7 +114,6 @@ func createCliApp() *cli.App {
 
     privateKeyFileFlag := cli.PathFlag{
         Name: "private-key",
-        Required: true,
         Usage: "Path to file with private key",
         Destination: &PrivateKeyPath,
     }
@@ -204,29 +203,45 @@ func setup(ctx *cli.Context) error {
     }
     slog.Info("Node UUID: " + uuid.String())
 
-    privKeyFile, err := os.Open(PrivateKeyPath)
-    if err != nil {
-        return cli.Exit(err, 1)
+    var privateKey wgtypes.Key
+    createPrivKey := func() error {
+        slog.Info("Creating a private key")
+        privateKey, err = wgtypes.GeneratePrivateKey()
+        if err != nil {
+            return cli.Exit(err, 1)
+        }
+        slog.Debug(fmt.Sprintf("Private key: %s", privateKey.String()))
+        return nil
     }
-    defer privKeyFile.Close()
-    slog.Debug("Keyfile opened for reading")
 
+    if PrivateKeyPath == "" {
+        err := createPrivKey()
+        if err != nil {
+            return err
+        }
+    } else {
+        privKeyFile, err := os.Open(PrivateKeyPath)
+        defer privKeyFile.Close()
+        if err != nil {
+            return cli.Exit(err, 1)
+        }
+        privateKeyStr := make([]byte, 45)
+        n, err := privKeyFile.Read(privateKeyStr)
+        if err != nil {
+            return cli.Exit(err, 1)
+        }
+        if n != 45 {
+            slog.Warn("Private key length did not math the expected 45!")
+        }
+        slog.Debug(fmt.Sprintf("Read %d bytes from keyfile", n))
 
-    privateKeyStr := make([]byte, 45)
-    n, err := privKeyFile.Read(privateKeyStr)
-    if err != nil {
-        return cli.Exit(err, 1)
+        privateKey, err = wgtypes.ParseKey(string(privateKeyStr))
+        slog.Debug("Keyfile opened for reading")
+        if err != nil {
+            return cli.Exit(err, 1)
+        }
+        slog.Debug("Private key parsed and is correct")
     }
-    if n != 45 {
-        slog.Warn("Private key length did not math the expected 45!")
-    }
-    slog.Debug(fmt.Sprintf("Read %d bytes from keyfile", n))
-
-    privateKey, err := wgtypes.ParseKey(string(privateKeyStr))
-    if err != nil {
-        return cli.Exit(err, 1)
-    }
-    slog.Debug("Private key parsed and is correct")
 
     wg, err := hvpnnode3.InitWGLink(
         InterfaceName, 
@@ -366,6 +381,7 @@ func handleStdin(c chan struct{}) {
         in = strings.ReplaceAll(in, "\n", "")
         if in == "q" || in == "exit" {
             c <- struct{}{}
+            break
         }
     }
 }