From 18cef8280ad7dc0eebf38a57aad7e6a3988a8825 Mon Sep 17 00:00:00 2001
From: David Merris <annika@adhdgirl.dev>
Date: Fri, 9 Aug 2024 11:32:09 -0400
Subject: [PATCH] [client] Allow setup keys to be provided in a file (#2337)

Adds a flag and a bit of logic to allow a setup key to be passed in using a file. The flag should be exclusive with the standard --setup-key flag.
---
 client/cmd/login.go          |  7 +++++++
 client/cmd/root.go           |  8 ++++++++
 client/cmd/up.go             |  7 +++++++
 client/cmd/up_daemon_test.go | 31 +++++++++++++++++++++++++++++++
 4 files changed, 53 insertions(+)

diff --git a/client/cmd/login.go b/client/cmd/login.go
index 14c973d91..512fbb081 100644
--- a/client/cmd/login.go
+++ b/client/cmd/login.go
@@ -39,6 +39,13 @@ var loginCmd = &cobra.Command{
 			ctx = context.WithValue(ctx, system.DeviceNameCtxKey, hostName)
 		}
 
+		if setupKeyPath != "" && setupKey == "" {
+			setupKey, err = getSetupKeyFromFile(setupKeyPath)
+			if err != nil {
+				return err
+			}
+		}
+
 		// workaround to run without service
 		if logFile == "console" {
 			err = handleRebrand(cmd)
diff --git a/client/cmd/root.go b/client/cmd/root.go
index db02ff5ea..b6d6694ee 100644
--- a/client/cmd/root.go
+++ b/client/cmd/root.go
@@ -56,6 +56,7 @@ var (
 	managementURL           string
 	adminURL                string
 	setupKey                string
+	setupKeyPath            string
 	hostName                string
 	preSharedKey            string
 	natExternalIPs          []string
@@ -128,6 +129,8 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&logLevel, "log-level", "l", "info", "sets Netbird log level")
 	rootCmd.PersistentFlags().StringVar(&logFile, "log-file", defaultLogFile, "sets Netbird log path. If console is specified the log will be output to stdout. If syslog is specified the log will be sent to syslog daemon.")
 	rootCmd.PersistentFlags().StringVarP(&setupKey, "setup-key", "k", "", "Setup key obtained from the Management Service Dashboard (used to register peer)")
+	rootCmd.PersistentFlags().StringVar(&setupKeyPath, "setup-key-file", "", "The path to a setup key obtained from the Management Service Dashboard (used to register peer) This is ignored if the setup-key flag is provided.")
+	rootCmd.MarkFlagsMutuallyExclusive("setup-key", "setup-key-file")
 	rootCmd.PersistentFlags().StringVar(&preSharedKey, preSharedKeyFlag, "", "Sets Wireguard PreSharedKey property. If set, then only peers that have the same key can communicate.")
 	rootCmd.PersistentFlags().StringVarP(&hostName, "hostname", "n", "", "Sets a custom hostname for the device")
 	rootCmd.PersistentFlags().BoolVarP(&anonymizeFlag, "anonymize", "A", false, "anonymize IP addresses and non-netbird.io domains in logs and status output")
@@ -253,6 +256,11 @@ var CLIBackOffSettings = &backoff.ExponentialBackOff{
 	Clock:               backoff.SystemClock,
 }
 
+func getSetupKeyFromFile(setupKeyPath string) (string, error) {
+	data, err := os.ReadFile(setupKeyPath)
+	return string(data), err
+}
+
 func handleRebrand(cmd *cobra.Command) error {
 	var err error
 	if logFile == defaultLogFile {
diff --git a/client/cmd/up.go b/client/cmd/up.go
index f69e9eb27..0eaf7bc0d 100644
--- a/client/cmd/up.go
+++ b/client/cmd/up.go
@@ -73,6 +73,13 @@ func upFunc(cmd *cobra.Command, args []string) error {
 		ctx = context.WithValue(ctx, system.DeviceNameCtxKey, hostName)
 	}
 
+	if setupKeyPath != "" && setupKey == "" {
+		setupKey, err = getSetupKeyFromFile(setupKeyPath)
+		if err != nil {
+			return err
+		}
+	}
+
 	if foregroundMode {
 		return runInForegroundMode(ctx, cmd)
 	}
diff --git a/client/cmd/up_daemon_test.go b/client/cmd/up_daemon_test.go
index 0295d2b21..daf8d0628 100644
--- a/client/cmd/up_daemon_test.go
+++ b/client/cmd/up_daemon_test.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"os"
 	"testing"
 	"time"
 
@@ -40,6 +41,36 @@ func TestUpDaemon(t *testing.T) {
 		return
 	}
 
+	// Test the setup-key-file flag.
+	tempFile, err := os.CreateTemp("", "setup-key")
+	if err != nil {
+		t.Errorf("could not create temp file, got error %v", err)
+		return
+	}
+	defer os.Remove(tempFile.Name())
+	if _, err := tempFile.Write([]byte("A2C8E62B-38F5-4553-B31E-DD66C696CEBB")); err != nil {
+		t.Errorf("could not write to temp file, got error %v", err)
+		return
+	}
+	if err := tempFile.Close(); err != nil {
+		t.Errorf("unable to close file, got error %v", err)
+	}
+	rootCmd.SetArgs([]string{
+		"login",
+		"--daemon-addr", "tcp://" + cliAddr,
+		"--setup-key-file", tempFile.Name(),
+		"--log-file", "",
+	})
+	if err := rootCmd.Execute(); err != nil {
+		t.Errorf("expected no error while running up command, got %v", err)
+		return
+	}
+	time.Sleep(time.Second * 3)
+	if status, err := state.Status(); err != nil && status != internal.StatusIdle {
+		t.Errorf("wrong status after login: %s, %v", internal.StatusIdle, err)
+		return
+	}
+
 	rootCmd.SetArgs([]string{
 		"up",
 		"--daemon-addr", "tcp://" + cliAddr,