Add UI master key setup and change flows

session/session.go

@@ -99,17 +99,17 @@ func (m *Manager) SaveRemote() error {
 		return ErrNoPath
 	}
 
-	var encoded bytes.Buffer
-	if err := vault.SaveKDBXWithConfigAndKey(&encoded, m.model, m.key, m.config); err != nil {
-		return fmt.Errorf("encode vault: %w", err)
+	encoded, err := m.persistableBytes()
+	if err != nil {
+		return err
 	}
 
-	version, err := m.remoteClient.Save(m.remotePath, bytes.NewReader(encoded.Bytes()), m.remoteVersion)
+	version, err := m.remoteClient.Save(m.remotePath, bytes.NewReader(encoded), m.remoteVersion)
 	if err != nil {
 		return fmt.Errorf("save remote %s: %w", m.remotePath, err)
 	}
 
-	m.encoded = encoded.Bytes()
+	m.encoded = encoded
 	m.remoteVersion = version
 	return nil
 }
@@ -165,16 +165,60 @@ func (m *Manager) Unlock(key vault.MasterKey) error {
 	return nil
 }
 
-func (m *Manager) saveToPath(path string) error {
-	var encoded bytes.Buffer
-	if err := vault.SaveKDBXWithConfigAndKey(&encoded, m.model, m.key, m.config); err != nil {
-		return fmt.Errorf("encode vault: %w", err)
+func (m *Manager) ChangeMasterKey(key vault.MasterKey) error {
+	var (
+		model  vault.Model
+		config *vault.KDBXConfig
+		err    error
+	)
+
+	if m.locked {
+		model, config, err = vault.LoadKDBXWithConfig(bytes.NewReader(m.encoded), m.key)
+		if err != nil {
+			return fmt.Errorf("decode locked vault: %w", err)
+		}
+	} else {
+		model = m.model
+		config = m.config
 	}
 
-	if err := os.WriteFile(path, encoded.Bytes(), 0o600); err != nil {
+	var encoded bytes.Buffer
+	if err := vault.SaveKDBXWithConfigAndKey(&encoded, model, key, config); err != nil {
+		return fmt.Errorf("encode vault with updated master key: %w", err)
+	}
+
+	m.key = key
+	m.config = config
+	m.encoded = encoded.Bytes()
+	if !m.locked {
+		m.model = model
+	}
+
+	return nil
+}
+
+func (m *Manager) saveToPath(path string) error {
+	encoded, err := m.persistableBytes()
+	if err != nil {
+		return err
+	}
+
+	if err := os.WriteFile(path, encoded, 0o600); err != nil {
 		return fmt.Errorf("write %s: %w", path, err)
 	}
 
-	m.encoded = encoded.Bytes()
+	m.encoded = encoded
 	return nil
 }
+
+func (m *Manager) persistableBytes() ([]byte, error) {
+	if m.locked {
+		return append([]byte(nil), m.encoded...), nil
+	}
+
+	var encoded bytes.Buffer
+	if err := vault.SaveKDBXWithConfigAndKey(&encoded, m.model, m.key, m.config); err != nil {
+		return nil, fmt.Errorf("encode vault: %w", err)
+	}
+	return encoded.Bytes(), nil
+}

session/session_test.go

@@ -391,6 +391,68 @@ func TestSaveUsesRemoteTargetWhenVaultWasOpenedFromWebDAV(t *testing.T) {
 	}
 }
 
+func TestChangeMasterKeyReencryptsSavedAndLockedVault(t *testing.T) {
+	t.Parallel()
+
+	originalKey := vault.MasterKey{Password: "old-password"}
+	updatedKey := vault.MasterKey{
+		Password:    "new-password",
+		KeyFileData: []byte("updated-key-file"),
+	}
+	model := vault.Model{
+		Entries: []vault.Entry{
+			{
+				ID:       "entry-1",
+				Title:    "Vault Console",
+				Username: "dannyocean",
+				Password: "token-1",
+				URL:      "https://vault.crew.example.invalid",
+				Path:     []string{"Root", "Internet"},
+			},
+		},
+	}
+
+	path := filepath.Join(t.TempDir(), "keepassgo.kdbx")
+
+	var sess Manager
+	if err := sess.Create(model, originalKey); err != nil {
+		t.Fatalf("Create() error = %v", err)
+	}
+	if err := sess.SaveAs(path); err != nil {
+		t.Fatalf("SaveAs() error = %v", err)
+	}
+	if err := sess.Lock(); err != nil {
+		t.Fatalf("Lock() error = %v", err)
+	}
+	if err := sess.ChangeMasterKey(updatedKey); err != nil {
+		t.Fatalf("ChangeMasterKey() error = %v", err)
+	}
+	if err := sess.Save(); err != nil {
+		t.Fatalf("Save() error = %v", err)
+	}
+
+	if err := sess.Unlock(updatedKey); err != nil {
+		t.Fatalf("Unlock(updatedKey) error = %v", err)
+	}
+
+	current, err := sess.Current()
+	if err != nil {
+		t.Fatalf("Current() error = %v", err)
+	}
+	got := current.EntriesInPath([]string{"Root", "Internet"})
+	if len(got) != 1 || got[0].Title != "Vault Console" {
+		t.Fatalf("Current() entries = %#v, want Vault Console entry after ChangeMasterKey", got)
+	}
+
+	var reopened Manager
+	if err := reopened.Open(path, updatedKey); err != nil {
+		t.Fatalf("Open(updatedKey) error = %v", err)
+	}
+	if err := reopened.Open(path, originalKey); !errors.Is(err, vault.ErrInvalidMasterKey) {
+		t.Fatalf("Open(originalKey) error = %v, want ErrInvalidMasterKey", err)
+	}
+}
+
 func TestSavePreservesOpenedKDBXSecuritySettings(t *testing.T) {
 	t.Parallel()