clean up code some

cmd/envvault/main.go

@@ -46,78 +46,6 @@ func init() {
 	localDbFilePath = ".envvault.json.enc"
 }
 
-type EnvStore struct {
-	Vars map[string]string   `json:"vars"`
-	Tags map[string][]string `json:"tags,omitempty"`
-}
-
-// GetTags returns all tags for a given variable name
-func (s *EnvStore) GetTags(varName string) []string {
-	if s.Tags == nil {
-		return []string{}
-	}
-	return s.Tags[varName]
-}
-
-// SetTags sets the tags for a given variable name
-func (s *EnvStore) SetTags(varName string, tags []string) {
-	if s.Tags == nil {
-		s.Tags = make(map[string][]string)
-	}
-
-	// Create a union of existing and new tags
-	existingTags := s.Tags[varName]
-	unionTags := make([]string, 0, len(existingTags)+len(tags))
-
-	// Add existing tags to the union
-	for _, tag := range existingTags {
-		unionTags = append(unionTags, tag)
-	}
-
-	// Add new tags if they don't already exist
-	for _, tag := range tags {
-		if !slices.Contains(unionTags, tag) {
-			unionTags = append(unionTags, tag)
-		}
-	}
-
-	s.Tags[varName] = unionTags
-}
-
-// RemoveTags removes the specified tags from a variable
-func (s *EnvStore) RemoveTags(varName string, tagsToRemove []string) {
-	if s.Tags == nil || len(s.Tags[varName]) == 0 {
-		return
-	}
-
-	existingTags := s.Tags[varName]
-	remainingTags := make([]string, 0, len(existingTags))
-
-	// Keep only tags that are not in tagsToRemove
-	for _, tag := range existingTags {
-		if !slices.Contains(tagsToRemove, tag) {
-			remainingTags = append(remainingTags, tag)
-		}
-	}
-
-	s.Tags[varName] = remainingTags
-}
-
-// HasAnyTag returns true if the variable has any of the specified tags
-func (s *EnvStore) HasAnyTag(varName string, tags []string) bool {
-	if len(tags) == 0 {
-		return true // No tags specified, include all variables
-	}
-
-	varTags := s.GetTags(varName)
-	for _, tag := range tags {
-		if slices.Contains(varTags, tag) {
-			return true
-		}
-	}
-	return false
-}
-
 type CLI struct {
 	Init   InitCmd   `cmd:"" help:"Initialize the vault."`
 	List   ListCmd   `cmd:"" help:"List all environment variables."`
@@ -208,11 +136,13 @@ func main() {
 }
 
 func subcommandInitVault() { // {{{
+	if exists(keyFilePath) {
+		internal.Log.Println("Vault already initialized")
+		fmt.Fprintln(os.Stderr, "Vault already initialized.")
+		return
+	}
+
 	internal.Log.Println("Initializing vault")
-
-	if _, err := os.Stat(keyFilePath); os.IsNotExist(err) {
-		internal.Log.Println("Key file does not exist, creating new vault")
-
 	fmt.Fprint(os.Stderr, "Enter a new master password: ")
 	password, err := term.ReadPassword(int(os.Stdin.Fd()))
 	fmt.Fprintln(os.Stderr)
@@ -262,27 +192,28 @@ func subcommandInitVault() { // {{{
 
 	// Initialize the db file
 	internal.Log.Println("Initializing empty environment store")
-		emptyStore := &EnvStore{Vars: make(map[string]string)}
+	emptyStore := &EnvStore{filePath: dbFilePath, Vars: make(map[string]string)}
-		saveGlobalEnvStore(emptyStore)
+	if err := emptyStore.Save(); err != nil {
+		internal.Log.Printf("Failed to initialize global environment store: %v", err)
+		log.Fatal(err)
+	}
 
 	internal.Log.Println("Vault initialized successfully")
 	fmt.Fprintln(os.Stderr, "Vault initialized successfully.")
-	} else {
-		internal.Log.Println("Vault already initialized")
-		fmt.Fprintln(os.Stderr, "Vault already initialized.")
-	}
 }
 
 // }}}
 
 func subcommandAddEnvVar(cmd AddCmd) { // {{{
-	var store *EnvStore
+	var storeKind StoreKind
 	if cmd.Local {
-		store = loadLocalEnvStore()
+		storeKind = StoreKindLocal
-		internal.Log.Printf("Adding environment variable '%s' to local vault", cmd.Name)
 	} else {
-		store = loadGlobalEnvStore()
+		storeKind = StoreKindGlobal
-		internal.Log.Printf("Adding environment variable '%s' to global vault", cmd.Name)
+	}
+	store, err := NewEnvStore(storeKind)
+	if err != nil {
+		log.Fatal(err)
 	}
 
 	if cmd.Value == "" {
@@ -296,28 +227,23 @@ func subcommandAddEnvVar(cmd AddCmd) { // {{{
 	}
 	store.Vars[cmd.Name] = cmd.Value
 
-	if cmd.Local {
+	if err := store.Save(); err != nil {
-		saveLocalEnvStore(store)
+		log.Fatal(err)
-		fmt.Fprintf(os.Stderr, "Environment variable '%s' added to local vault.\n", cmd.Name)
-	} else {
-		saveGlobalEnvStore(store)
-		fmt.Fprintf(os.Stderr, "Environment variable '%s' added to global vault.\n", cmd.Name)
 	}
+	fmt.Fprintf(os.Stderr, "Environment variable '%s' added to vault.\n", cmd.Name)
 } // }}}
 
 func subcommandRmEnvVar(name string, localOnly bool) { // {{{
-	localExists := false
-
 	// Check local store if it exists
-	if _, err := os.Stat(localDbFilePath); err == nil {
+	if exists(localDbFilePath) {
-		localStore := loadLocalEnvStore()
+		localStore, err := NewEnvStoreFromPath(localDbFilePath)
-		if _, exists := localStore.Vars[name]; exists {
+		if err != nil {
-			delete(localStore.Vars, name)
+			log.Fatal(err)
-			saveLocalEnvStore(localStore)
+		}
-			localExists = true
+
-			fmt.Fprintf(os.Stderr, "Environment variable '%s' removed from local vault.\n", name)
+		localStore.RemoveEnvVar(name)
-		} else if localOnly {
+		if err := localStore.Save(); err != nil {
-			fmt.Fprintf(os.Stderr, "Environment variable '%s' not found in local vault.\n", name)
+			log.Fatal(err)
 		}
 	} else if localOnly {
 		fmt.Fprintf(os.Stderr, "Local vault does not exist.\n")
@@ -325,31 +251,28 @@ func subcommandRmEnvVar(name string, localOnly bool) { // {{{
 
 	// Check global store if not local-only
 	if !localOnly {
-		globalStore := loadGlobalEnvStore()
+		globalStore, err := NewEnvStoreFromPath(dbFilePath)
-		if _, exists := globalStore.Vars[name]; exists {
+		if err != nil {
-			delete(globalStore.Vars, name)
+			log.Fatal(err)
-			saveGlobalEnvStore(globalStore)
+		}
-			fmt.Fprintf(os.Stderr, "Environment variable '%s' removed from global vault.\n", name)
+
-		} else if !localExists {
+		globalStore.RemoveEnvVar(name)
-			fmt.Fprintf(os.Stderr, "Environment variable '%s' not found in global vault.\n", name)
+		if err := globalStore.Save(); err != nil {
+			log.Fatal(err)
 		}
 	}
 } // }}}
 
 func subcommandTagEnvVar(cmd TagCmd) { // {{{
-	var store *EnvStore
+	var storeKind StoreKind
-
-	// Determine which store to use
 	if cmd.Local {
-		if _, err := os.Stat(localDbFilePath); os.IsNotExist(err) {
+		storeKind = StoreKindLocal
-			fmt.Fprintf(os.Stderr, "Local vault does not exist.\n")
-			os.Exit(1)
-		}
-		store = loadLocalEnvStore()
-		internal.Log.Printf("Using local vault")
 	} else {
-		store = loadGlobalEnvStore()
+		storeKind = StoreKindGlobal
-		internal.Log.Printf("Using global vault")
+	}
+	store, err := NewEnvStore(storeKind)
+	if err != nil {
+		log.Fatal(err)
 	}
 
 	// Check if the variable exists
@@ -387,12 +310,6 @@ func subcommandTagEnvVar(cmd TagCmd) { // {{{
 	if cmd.Remove {
 		internal.Log.Printf("Removing tags %v from variable '%s'", cmd.Tags, cmd.Name)
 		store.RemoveTags(cmd.Name, cmd.Tags)
-		fmt.Fprintf(
-			os.Stderr,
-			"Removed tags from '%s': %s\n",
-			cmd.Name,
-			strings.Join(cmd.Tags, ", "),
-		)
 	} else {
 		internal.Log.Printf("Adding tags %v to variable '%s'", cmd.Tags, cmd.Name)
 		store.SetTags(cmd.Name, cmd.Tags)
@@ -400,26 +317,21 @@ func subcommandTagEnvVar(cmd TagCmd) { // {{{
 	}
 
 	// Save the store
-	if cmd.Local {
+	if err := store.Save(); err != nil {
-		saveLocalEnvStore(store)
+		log.Fatal(err)
-	} else {
-		saveGlobalEnvStore(store)
 	}
 } // }}}
 
 func subcommandListEnvVars(cmdArgs ListCmd) { // {{{
-	var store *EnvStore
+	var storeKind StoreKind
-
-	// Determine which store to use
 	if cmdArgs.Local {
-		if _, err := os.Stat(localDbFilePath); os.IsNotExist(err) {
+		storeKind = StoreKindLocal
-			fmt.Fprintf(os.Stderr, "Local vault does not exist.\n")
-			return
-		}
-		store = loadLocalEnvStore()
 	} else {
-		// Default: load and merge stores
+		storeKind = StoreKindGlobal
-		store = loadEnvStore()
+	}
+	store, err := NewEnvStore(storeKind)
+	if err != nil {
+		log.Fatal(err)
 	}
 
 	// Sort keys for consistent output
@@ -452,21 +364,17 @@ func subcommandListEnvVars(cmdArgs ListCmd) { // {{{
 } // }}}
 
 func subcommandExecCommand(cmdArgs ExecCmd) { // {{{
-	var store *EnvStore
+	var storeKind StoreKind
-
-	// Determine which store to use
 	if cmdArgs.Local {
-		if _, err := os.Stat(localDbFilePath); os.IsNotExist(err) {
+		storeKind = StoreKindLocal
-			fmt.Fprintf(os.Stderr, "Local vault does not exist.\n")
-			os.Exit(1)
-		}
-		store = loadLocalEnvStore()
-		internal.Log.Println("Using local vault only for execution")
 	} else {
-		// Default: load and merge stores
+		storeKind = StoreKindGlobal
-		store = loadEnvStore()
-		internal.Log.Println("Using merged vaults for execution")
 	}
+	store, err := NewEnvStore(storeKind)
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	envVars := os.Environ()
 	for candidateEnvName, candidateEnvValue := range store.Vars {
 		// Skip if not matching tag filter
@@ -628,6 +536,191 @@ func subcommandRekeyVault() { // {{{
 //
 ////////////////////////////////////////////////////////////////////////////////
 
+func exists(path string) bool {
+	_, err := os.Stat(path)
+	if err != nil && os.IsNotExist(err) {
+		return false
+	}
+	return err == nil
+}
+
+type EnvStore struct { // {{{
+	filePath string
+
+	Vars map[string]string   `json:"vars"`
+	Tags map[string][]string `json:"tags,omitempty"`
+}
+
+func NewEnvStoreFromPath(filePath string) (*EnvStore, error) {
+	key := loadKey()
+
+	store := EnvStore{
+		filePath: filePath,
+		Vars:     make(map[string]string),
+		Tags:     make(map[string][]string),
+	}
+
+	data, err := os.ReadFile(filePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return &store, nil
+		}
+		return nil, err
+	}
+
+	if len(data) == 0 {
+		return &store, nil
+	}
+
+	var nonce [24]byte
+	copy(nonce[:], data[:24])
+	decrypted, ok := secretbox.Open(nil, data[24:], &nonce, &key)
+	if !ok {
+		log.Fatal("Decryption failed")
+	}
+	if err := json.Unmarshal(decrypted, &store); err != nil {
+		log.Fatal(err)
+	}
+	// Initialize Tags map if it's nil
+	if store.Tags == nil {
+		store.Tags = make(map[string][]string)
+	}
+	return &store, nil
+}
+
+type StoreKind int
+
+const (
+	StoreKindGlobal = iota
+	StoreKindLocal
+	StoreKindMerged
+)
+
+func NewEnvStore(kind StoreKind) (*EnvStore, error) {
+	globalStore, globalErr := NewEnvStoreFromPath(dbFilePath)
+	localStore, localErr := NewEnvStoreFromPath(localDbFilePath)
+
+	if kind == StoreKindGlobal {
+		if globalErr != nil {
+			return nil, globalErr
+		}
+		return globalStore, nil
+	}
+	if kind == StoreKindLocal {
+		if localErr != nil {
+			return nil, localErr
+		}
+		return localStore, nil
+	}
+
+	if globalErr != nil {
+		return nil, globalErr
+	}
+	if localErr != nil {
+		return nil, localErr
+	}
+
+	// Override/add local vars and tags
+	maps.Copy(globalStore.Vars, localStore.Vars)
+	maps.Copy(globalStore.Tags, localStore.Tags)
+	return globalStore, nil
+}
+
+func (s *EnvStore) Save() error {
+	key := loadKey()
+	data, err := json.Marshal(s)
+	if err != nil {
+		return err
+	}
+
+	var nonce [24]byte
+	if _, err := rand.Read(nonce[:]); err != nil {
+		return err
+	}
+	encrypted := secretbox.Seal(nonce[:], data, &nonce, &key)
+
+	if err := os.MkdirAll(filepath.Dir(s.filePath), 0700); err != nil {
+		return err
+	}
+	if err := os.WriteFile(s.filePath, encrypted, 0600); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// GetTags returns all tags for a given variable name
+func (s *EnvStore) GetTags(varName string) []string {
+	if s.Tags == nil {
+		return []string{}
+	}
+	return s.Tags[varName]
+}
+
+// SetTags sets the tags for a given variable name
+func (s *EnvStore) SetTags(varName string, tags []string) {
+	if s.Tags == nil {
+		s.Tags = make(map[string][]string)
+	}
+
+	// Create a union of existing and new tags
+	existingTags := s.Tags[varName]
+	unionTags := make([]string, 0, len(existingTags)+len(tags))
+
+	// Add existing tags to the union
+	for _, tag := range existingTags {
+		unionTags = append(unionTags, tag)
+	}
+
+	// Add new tags if they don't already exist
+	for _, tag := range tags {
+		if !slices.Contains(unionTags, tag) {
+			unionTags = append(unionTags, tag)
+		}
+	}
+
+	s.Tags[varName] = unionTags
+}
+
+// RemoveTags removes the specified tags from a variable
+func (s *EnvStore) RemoveTags(varName string, tagsToRemove []string) {
+	if s.Tags == nil || len(s.Tags[varName]) == 0 {
+		return
+	}
+
+	existingTags := s.Tags[varName]
+	remainingTags := make([]string, 0, len(existingTags))
+
+	// Keep only tags that are not in tagsToRemove
+	for _, tag := range existingTags {
+		if !slices.Contains(tagsToRemove, tag) {
+			remainingTags = append(remainingTags, tag)
+		}
+	}
+
+	s.Tags[varName] = remainingTags
+}
+
+// HasAnyTag returns true if the variable has any of the specified tags
+func (s *EnvStore) HasAnyTag(varName string, tags []string) bool {
+	if len(tags) == 0 {
+		return true // No tags specified, include all variables
+	}
+
+	varTags := s.GetTags(varName)
+	for _, tag := range tags {
+		if slices.Contains(varTags, tag) {
+			return true
+		}
+	}
+	return false
+}
+
+func (s *EnvStore) RemoveEnvVar(name string) {
+	delete(s.Vars, name)
+	delete(s.Tags, name)
+} // }}}
+
 func loadPassword() []byte { // {{{
 	internal.Log.Println("Getting password without caching")
 
@@ -673,14 +766,14 @@ func loadPassword() []byte { // {{{
 } // }}}
 
 func loadKey() [32]byte { // {{{
-	// Get password but don't cache it yet
-	password := loadPassword()
-
 	encryptedKey, err := os.ReadFile(keyFilePath)
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	// Get password but don't cache it yet
+	password := loadPassword()
+
 	decryptedKey, err := decryptKeyWithPassword(encryptedKey, password)
 	if err != nil {
 		log.Fatal("Invalid password or corrupted key file")
@@ -753,118 +846,4 @@ func decryptKeyWithPassword(encryptedKey, password []byte) ([]byte, error) { //
 
 	return decrypted, nil
 } // }}}
-
+// }}}
-func loadStoreFile(filePath string) *EnvStore { // {{{
-	key := loadKey()
-
-	data, err := os.ReadFile(filePath)
-	if err != nil {
-		if os.IsNotExist(err) {
-			return &EnvStore{
-				Vars: make(map[string]string),
-				Tags: make(map[string][]string),
-			}
-		}
-		log.Fatal(err)
-	}
-
-	var store EnvStore
-	if len(data) > 0 {
-		var nonce [24]byte
-		copy(nonce[:], data[:24])
-		decrypted, ok := secretbox.Open(nil, data[24:], &nonce, &key)
-		if !ok {
-			log.Fatal("Decryption failed")
-		}
-		if err := json.Unmarshal(decrypted, &store); err != nil {
-			log.Fatal(err)
-		}
-		// Initialize Tags map if it's nil
-		if store.Tags == nil {
-			store.Tags = make(map[string][]string)
-		}
-	} else {
-		return &EnvStore{
-			Vars: make(map[string]string),
-			Tags: make(map[string][]string),
-		}
-	}
-	return &store
-} // }}}
-
-func loadGlobalEnvStore() *EnvStore { // {{{
-	if err := os.MkdirAll(filepath.Dir(dbFilePath), 0700); err != nil {
-		log.Fatal(err)
-	}
-	return loadStoreFile(dbFilePath)
-} // }}}
-
-func loadLocalEnvStore() *EnvStore { // {{{
-	return loadStoreFile(localDbFilePath)
-} // }}}
-
-func loadEnvStore() *EnvStore { // {{{
-	// Load global store first
-	globalStore := loadGlobalEnvStore()
-
-	// Check if local store exists
-	if _, err := os.Stat(localDbFilePath); os.IsNotExist(err) {
-		// No local store, just return global
-		return globalStore
-	}
-
-	// Load local store
-	localStore := loadLocalEnvStore()
-
-	// Merge stores (local takes precedence)
-	mergedStore := &EnvStore{
-		Vars: make(map[string]string),
-		Tags: make(map[string][]string),
-	}
-
-	// Copy all global vars and tags
-	maps.Copy(mergedStore.Vars, globalStore.Vars)
-	maps.Copy(mergedStore.Tags, globalStore.Tags)
-
-	// Override/add local vars and tags
-	maps.Copy(mergedStore.Vars, localStore.Vars)
-	maps.Copy(mergedStore.Tags, localStore.Tags)
-
-	return mergedStore
-} // }}}
-
-func saveGlobalEnvStore(store *EnvStore) { // {{{
-	key := loadKey()
-	data, err := json.Marshal(store)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	var nonce [24]byte
-	if _, err := rand.Read(nonce[:]); err != nil {
-		log.Fatal(err)
-	}
-	encrypted := secretbox.Seal(nonce[:], data, &nonce, &key)
-
-	if err := os.WriteFile(dbFilePath, encrypted, 0600); err != nil {
-		log.Fatal(err)
-	}
-} // }}}
-
-func saveLocalEnvStore(store *EnvStore) { // {{{
-	key := loadKey()
-	data, err := json.Marshal(store)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	var nonce [24]byte
-	if _, err := rand.Read(nonce[:]); err != nil {
-		log.Fatal(err)
-	}
-	encrypted := secretbox.Seal(nonce[:], data, &nonce, &key)
-
-	if err := os.WriteFile(localDbFilePath, encrypted, 0600); err != nil {
-		log.Fatal(err)
-	}
-} // }}}