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
+}