From c02b8f3dff4d00a23ee483de71468dbcb227b07c Mon Sep 17 00:00:00 2001
From: Brian Buller <brian@bullercodeworks.com>
Date: Fri, 17 May 2024 07:50:03 -0500
Subject: [PATCH] Reprompt For Password

---
 cmd/rofi.go   | 54 ++++++++++++++++++++++++---------------------------
 util/files.go |  2 ++
 2 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/cmd/rofi.go b/cmd/rofi.go
index f525a6f..1aff750 100644
--- a/cmd/rofi.go
+++ b/cmd/rofi.go
@@ -26,45 +26,41 @@ func init() {
 }
 
 func runRofi(cmd *cobra.Command, args []string) error {
+	var db *models.KeePassDB
+
+	var argIsPw bool
+	pwFile := fmt.Sprintf("%s%s", ConfigDir, "pw")
+	pass, err := util.ReadPWFile(pwFile)
+	if err != nil {
+		// See if we have a password...
+		if len(args) == 1 {
+			pass = args[0]
+			argIsPw = true
+		} else {
+			fmt.Println("Enter Master Password")
+			return nil
+		}
+	}
+
+	db, err = models.NewKeePassDB(viper.GetString("database"), pass)
+	if err != nil {
+		util.RemovePWFile(pwFile)
+		return fmt.Errorf("Error opening DB. Re-enter Master Password:\n%v", err)
+	} else if err = util.WritePWFile(pwFile, pass); err != nil {
+		return fmt.Errorf("DB Opened, but failed to persist password: %v", err)
+	}
+
 	var path []string
 	for i := range args {
 		path = append(path, strings.Split(args[i], "/")...)
 	}
 
-	var db *models.KeePassDB
-	var havePassword bool
-
-	pwFile := fmt.Sprintf("%s%s", ConfigDir, "pw")
-	pass, err := util.ReadPWFile(pwFile)
-	if err != nil {
-		havePassword = false
-		// See if we have a password...
-		if len(args) == 1 {
-			pass = args[0]
-		} else {
-			fmt.Println("Enter Master Password")
-			return nil
-		}
-	} else {
-		havePassword = true
-	}
-
-	db, err = models.NewKeePassDB(viper.GetString("database"), pass)
-	if err != nil {
-		return err
-	} else if !havePassword {
-		if err = util.WritePWFile(pwFile, pass); err != nil {
-			return err
-		}
-		path = []string{}
-	}
-	if len(path) == 0 {
+	if len(path) == 0 || argIsPw {
 		list := db.GetAllEntriesFromRoot()
 		for i := range list {
 			fmt.Println(strings.Join(list[i], "/"))
 		}
 	} else {
-		// Only one... Could be a single group nested, or it's an entry
 		if entry, entryErr := db.FindEntryFromRoot(path); entryErr != nil {
 			return entryErr
 		} else {
diff --git a/util/files.go b/util/files.go
index 5353a86..acc116a 100644
--- a/util/files.go
+++ b/util/files.go
@@ -5,6 +5,7 @@ package util
 
 import (
 	"io/ioutil"
+	"os"
 )
 
 func ReadFile(path string) (string, error) {
@@ -35,3 +36,4 @@ func WritePWFile(path, pw string) error {
 	}
 	return ioutil.WriteFile(path, []byte(pass), 0600)
 }
+func RemovePWFile(path string) error { return os.Remove(path) }