From 56d82a99e15b0e88f4b102725b859210136d0a40 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Fri, 6 Oct 2023 15:41:29 +0200
Subject: [PATCH] Backup and restore encryption key

---
 management/cmd/management.go                  | 14 ++++++--
 management/server/activity/sqlite/crypt.go    | 32 ++++++++++++++++++-
 .../server/activity/sqlite/crypt_test.go      | 15 +++++++--
 .../server/activity/sqlite/sqlite_test.go     |  2 +-
 4 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/management/cmd/management.go b/management/cmd/management.go
index f85cf225e..7e47a315c 100644
--- a/management/cmd/management.go
+++ b/management/cmd/management.go
@@ -301,15 +301,23 @@ var (
 func initEventStore(dataDir string, key string) (activity.Store, string, error) {
 	var err error
 	if key == "" {
-		log.Debugf("generate new activity store encryption key")
-		key, err = sqlite.GenerateKey()
+		log.Debugf("restore or generate new activity store encryption key")
+		key, err = sqlite.RestoreKey(dataDir)
+		if err == nil {
+			goto CreateStore
+		} else {
+			log.Debugf("failed to restore encryption key for activity store: %s", err)
+		}
+
+		log.Infof("generate new encryption key for activity store")
+		key, err = sqlite.GenerateKey(dataDir)
 		if err != nil {
 			return nil, "", err
 		}
 	}
+CreateStore:
 	store, err := sqlite.NewSQLiteStore(dataDir, key)
 	return store, key, err
-
 }
 
 func notifyStop(msg string) {
diff --git a/management/server/activity/sqlite/crypt.go b/management/server/activity/sqlite/crypt.go
index cf4dda746..2c5efebe4 100644
--- a/management/server/activity/sqlite/crypt.go
+++ b/management/server/activity/sqlite/crypt.go
@@ -7,6 +7,12 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"fmt"
+	"os"
+	"path/filepath"
+)
+
+const (
+	backupFile = ".datastore.key"
 )
 
 var iv = []byte{10, 22, 13, 79, 05, 8, 52, 91, 87, 98, 88, 98, 35, 25, 13, 05}
@@ -15,16 +21,40 @@ type FieldEncrypt struct {
 	block cipher.Block
 }
 
-func GenerateKey() (string, error) {
+func RestoreKey(dataDir string) (string, error) {
+	fName := filepath.Join(dataDir, backupFile)
+	data, err := os.ReadFile(fName)
+	return string(data), err
+}
+
+func GenerateKey(dataDir string) (string, error) {
 	key := make([]byte, 32)
 	_, err := rand.Read(key)
 	if err != nil {
 		return "", err
 	}
 	readableKey := base64.StdEncoding.EncodeToString(key)
+
+	err = saveKey(dataDir, readableKey)
+	if err != nil {
+		return "", err
+	}
 	return readableKey, nil
 }
 
+func saveKey(dataDir, key string) error {
+	f, err := os.Create(filepath.Join(dataDir, backupFile))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = f.WriteString(key)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func NewFieldEncrypt(key string) (*FieldEncrypt, error) {
 	binKey, err := base64.StdEncoding.DecodeString(key)
 	if err != nil {
diff --git a/management/server/activity/sqlite/crypt_test.go b/management/server/activity/sqlite/crypt_test.go
index efa740921..34c5e789a 100644
--- a/management/server/activity/sqlite/crypt_test.go
+++ b/management/server/activity/sqlite/crypt_test.go
@@ -2,11 +2,20 @@ package sqlite
 
 import (
 	"testing"
+
+	log "github.com/sirupsen/logrus"
 )
 
+func TestRestoreKey(t *testing.T) {
+	_, err := RestoreKey(t.TempDir())
+	if err != nil {
+		log.Infof("err: %s", err)
+	}
+}
+
 func TestGenerateKey(t *testing.T) {
 	testData := "exampl@netbird.io"
-	key, err := GenerateKey()
+	key, err := GenerateKey(t.TempDir())
 	if err != nil {
 		t.Fatalf("failed to generate key: %s", err)
 	}
@@ -32,7 +41,7 @@ func TestGenerateKey(t *testing.T) {
 
 func TestCorruptKey(t *testing.T) {
 	testData := "exampl@netbird.io"
-	key, err := GenerateKey()
+	key, err := GenerateKey(t.TempDir())
 	if err != nil {
 		t.Fatalf("failed to generate key: %s", err)
 	}
@@ -46,7 +55,7 @@ func TestCorruptKey(t *testing.T) {
 		t.Fatalf("invalid encrypted text")
 	}
 
-	newKey, err := GenerateKey()
+	newKey, err := GenerateKey(t.TempDir())
 	if err != nil {
 		t.Fatalf("failed to generate key: %s", err)
 	}
diff --git a/management/server/activity/sqlite/sqlite_test.go b/management/server/activity/sqlite/sqlite_test.go
index f6a6f9467..c3c8e0513 100644
--- a/management/server/activity/sqlite/sqlite_test.go
+++ b/management/server/activity/sqlite/sqlite_test.go
@@ -12,7 +12,7 @@ import (
 
 func TestNewSQLiteStore(t *testing.T) {
 	dataDir := t.TempDir()
-	key, _ := GenerateKey()
+	key, _ := GenerateKey(dataDir)
 	store, err := NewSQLiteStore(dataDir, key)
 	if err != nil {
 		t.Fatal(err)