secret/internal/cli/vault.go

package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"git.eeqj.de/sneak/secret/internal/secret"
	"git.eeqj.de/sneak/secret/internal/vault"
	"git.eeqj.de/sneak/secret/pkg/agehd"
	"github.com/awnumar/memguard"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/tyler-smith/go-bip39"
)

// Sentinel errors for vault operations
var (
	errMnemonicEmpty         = errors.New("mnemonic cannot be empty")
	errInvalidMnemonicPhrase = errors.New("invalid BIP39 mnemonic phrase")
	errInvalidMnemonic       = errors.New("invalid BIP39 mnemonic")
	errVaultHasLongTermKey   = errors.New(
		"already has a long-term key configured")
	errMnemonicEnvNotSet = errors.New(
		"SB_SECRET_MNEMONIC environment variable not set")
	errPassphraseEnvNotSet = errors.New(
		"SB_UNLOCK_PASSPHRASE environment variable not set")
	errCannotRemoveLastVault = errors.New("cannot remove the last vault")
	errVaultContainsSecrets  = errors.New(
		"contains secrets; use --force to remove")
)

func newVaultCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "vault",
		Short: "Manage vaults",
		Long:  `Create, list, and select vaults for organizing secrets.`,
	}

	cmd.AddCommand(newVaultListCmd())
	cmd.AddCommand(newVaultCreateCmd())
	cmd.AddCommand(newVaultSelectCmd())
	cmd.AddCommand(newVaultImportCmd())
	cmd.AddCommand(newVaultRemoveCmd())

	return cmd
}

func newVaultListCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:     cmdUseList,
		Aliases: []string{"ls"},
		Short:   "List available vaults",
		RunE: func(cmd *cobra.Command, _ []string) error {
			jsonOutput, _ := cmd.Flags().GetBool("json")

			cli, err := NewCLIInstance()
			if err != nil {
				return fmt.Errorf("failed to initialize CLI: %w", err)
			}

			return cli.ListVaults(cmd, jsonOutput)
		},
	}

	cmd.Flags().Bool("json", false, "Output in JSON format")

	return cmd
}

func newVaultCreateCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "create <name>",
		Short: "Create a new vault",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			cli, err := NewCLIInstance()
			if err != nil {
				return fmt.Errorf("failed to initialize CLI: %w", err)
			}

			return cli.CreateVault(cmd, args[0])
		},
	}
}

func newVaultSelectCmd() *cobra.Command {
	cli, err := NewCLIInstance()
	if err != nil {
		log.Fatalf("failed to initialize CLI: %v", err)
	}

	return &cobra.Command{
		Use:               "select <name>",
		Short:             "Select a vault as current",
		Args:              cobra.ExactArgs(1),
		ValidArgsFunction: getVaultNamesCompletionFunc(cli.fs, cli.stateDir),
		RunE: func(cmd *cobra.Command, args []string) error {
			cli, err := NewCLIInstance()
			if err != nil {
				return fmt.Errorf("failed to initialize CLI: %w", err)
			}

			return cli.SelectVault(cmd, args[0])
		},
	}
}

func newVaultImportCmd() *cobra.Command {
	cli, err := NewCLIInstance()
	if err != nil {
		log.Fatalf("failed to initialize CLI: %v", err)
	}

	return &cobra.Command{
		Use:   "import <vault-name>",
		Short: "Import a mnemonic into a vault",
		Long: `Import a BIP39 mnemonic phrase into the specified vault ` +
			`(default if not specified).`,
		Args:              cobra.MaximumNArgs(1),
		ValidArgsFunction: getVaultNamesCompletionFunc(cli.fs, cli.stateDir),
		RunE: func(cmd *cobra.Command, args []string) error {
			vaultName := "default"
			if len(args) > 0 {
				vaultName = args[0]
			}

			cli, err := NewCLIInstance()
			if err != nil {
				return fmt.Errorf("failed to initialize CLI: %w", err)
			}

			return cli.VaultImport(cmd, vaultName)
		},
	}
}

func newVaultRemoveCmd() *cobra.Command {
	cli, err := NewCLIInstance()
	if err != nil {
		log.Fatalf("failed to initialize CLI: %v", err)
	}

	cmd := &cobra.Command{
		Use:     "remove <name>",
		Aliases: []string{"rm"},
		Short:   "Remove a vault",
		Long: `Remove a vault. Requires --force if the vault contains ` +
			`secrets. Will automatically switch to another vault if ` +
			`removing the currently selected one.`,
		Args:              cobra.ExactArgs(1),
		ValidArgsFunction: getVaultNamesCompletionFunc(cli.fs, cli.stateDir),
		RunE: func(cmd *cobra.Command, args []string) error {
			force, _ := cmd.Flags().GetBool("force")

			cli, err := NewCLIInstance()
			if err != nil {
				return fmt.Errorf("failed to initialize CLI: %w", err)
			}

			return cli.RemoveVault(cmd, args[0], force)
		},
	}

	cmd.Flags().BoolP("force", "f", false, "Force removal even if vault contains secrets")

	return cmd
}

// ListVaults lists all available vaults
func (cli *Instance) ListVaults(cmd *cobra.Command, jsonOutput bool) error {
	vaults, err := vault.ListVaults(cli.fs, cli.stateDir)
	if err != nil {
		return err
	}

	if jsonOutput { //nolint:nestif // Separate JSON and text output formatting logic
		// Get current vault name for context
		currentVault := ""

		currentVlt, err := vault.GetCurrentVault(cli.fs, cli.stateDir)
		if err == nil {
			currentVault = currentVlt.GetName()
		}

		result := map[string]any{
			"vaults":       vaults,
			"currentVault": currentVault,
		}

		jsonBytes, err := json.MarshalIndent(result, "", "  ")
		if err != nil {
			return err
		}

		cmd.Println(string(jsonBytes))
	} else {
		// Text output
		cmd.Println("Available vaults:")

		if len(vaults) == 0 {
			cmd.Println("  (none)")
		} else {
			// Try to get current vault for marking
			currentVault := ""

			currentVlt, err := vault.GetCurrentVault(cli.fs, cli.stateDir)
			if err == nil {
				currentVault = currentVlt.GetName()
			}

			for _, vaultName := range vaults {
				if vaultName == currentVault {
					cmd.Printf("  %s (current)\n", vaultName)
				} else {
					cmd.Printf("  %s\n", vaultName)
				}
			}
		}
	}

	return nil
}

// setMnemonicEnv sets the mnemonic environment variable and returns a
// function that restores the previous value
func setMnemonicEnv(mnemonicStr string) func() {
	originalMnemonic := os.Getenv(secret.EnvMnemonic)
	_ = os.Setenv(secret.EnvMnemonic, mnemonicStr)

	return func() {
		if originalMnemonic != "" {
			_ = os.Setenv(secret.EnvMnemonic, originalMnemonic)
		} else {
			_ = os.Unsetenv(secret.EnvMnemonic)
		}
	}
}

// resolvePassphrase returns the unlock passphrase from the environment or
// prompts the user for it with confirmation
func resolvePassphrase() (*memguard.LockedBuffer, error) {
	if envPassphrase := os.Getenv(secret.EnvUnlockPassphrase); envPassphrase != "" {
		secret.Debug("Using unlock passphrase from environment variable")

		return memguard.NewBufferFromBytes([]byte(envPassphrase)), nil
	}

	secret.Debug("Prompting user for unlock passphrase")

	// Use secure passphrase input with confirmation
	passphraseBuffer, err := readSecurePassphrase("Enter passphrase for unlocker: ")
	if err != nil {
		return nil, fmt.Errorf("failed to read passphrase: %w", err)
	}

	return passphraseBuffer, nil
}

// CreateVault creates a new vault
func (cli *Instance) CreateVault(cmd *cobra.Command, name string) error {
	secret.Debug("Creating new vault", "name", name, "state_dir", cli.stateDir)

	// Get or prompt for mnemonic
	var mnemonicStr string

	if envMnemonic := os.Getenv(secret.EnvMnemonic); envMnemonic != "" {
		secret.Debug("Using mnemonic from environment variable")

		mnemonicStr = envMnemonic
	} else {
		secret.Debug("Prompting user for mnemonic phrase")
		// Read mnemonic securely without echo
		mnemonicBuffer, err := secret.ReadPassphrase(
			"Enter your BIP39 mnemonic phrase: ")
		if err != nil {
			secret.Debug("Failed to read mnemonic from stdin", "error", err)

			return fmt.Errorf("failed to read mnemonic: %w", err)
		}
		defer mnemonicBuffer.Destroy()

		mnemonicStr = mnemonicBuffer.String()

		fmt.Fprintln(os.Stderr) // Add newline after hidden input
	}

	if mnemonicStr == "" {
		return errMnemonicEmpty
	}

	// Validate the mnemonic
	mnemonicWords := strings.Fields(mnemonicStr)
	secret.Debug("Validating BIP39 mnemonic", "word_count", len(mnemonicWords))

	if !bip39.IsMnemonicValid(mnemonicStr) {
		return errInvalidMnemonicPhrase
	}

	// Set mnemonic in environment for CreateVault to use
	restoreMnemonicEnv := setMnemonicEnv(mnemonicStr)
	defer restoreMnemonicEnv()

	// Create the vault - it will handle key derivation internally
	vlt, err := vault.CreateVault(cli.fs, cli.stateDir, name)
	if err != nil {
		return err
	}

	// Get the vault metadata to retrieve the derivation index
	vaultDir := filepath.Join(cli.stateDir, "vaults.d", name)

	metadata, err := vault.LoadVaultMetadata(cli.fs, vaultDir)
	if err != nil {
		return fmt.Errorf("failed to load vault metadata: %w", err)
	}

	// Derive the long-term key using the same index that CreateVault used
	ltIdentity, err := agehd.DeriveIdentity(mnemonicStr, metadata.DerivationIndex)
	if err != nil {
		return fmt.Errorf("failed to derive long-term key from mnemonic: %w", err)
	}

	// Unlock the vault with the derived long-term key
	vlt.Unlock(ltIdentity)

	// Get or prompt for passphrase
	passphraseBuffer, err := resolvePassphrase()
	if err != nil {
		return err
	}
	defer passphraseBuffer.Destroy()

	// Create passphrase-protected unlocker
	secret.Debug("Creating passphrase-protected unlocker")

	passphraseUnlocker, err := vlt.CreatePassphraseUnlocker(passphraseBuffer)
	if err != nil {
		return fmt.Errorf("failed to create unlocker: %w", err)
	}

	cmd.Printf("Created vault '%s'\n", vlt.GetName())
	cmd.Printf("Long-term public key: %s\n", ltIdentity.Recipient().String())
	cmd.Printf("Unlocker ID: %s\n", passphraseUnlocker.GetID())

	return nil
}

// SelectVault selects a vault as the current one
func (cli *Instance) SelectVault(cmd *cobra.Command, name string) error {
	err := vault.SelectVault(cli.fs, cli.stateDir, name)
	if err != nil {
		return err
	}

	cmd.Printf("Selected vault '%s' as current\n", name)

	return nil
}

// vaultImportPreflight verifies the vault exists without a long-term key
// and returns the vault directory, public key path, and validated mnemonic
func (cli *Instance) vaultImportPreflight(
	vlt *vault.Vault, vaultName string,
) (string, string, string, error) {
	// Check if vault exists
	vaultDir, err := vlt.GetDirectory()
	if err != nil {
		return "", "", "", err
	}

	exists, err := afero.DirExists(cli.fs, vaultDir)
	if err != nil {
		return "", "", "", fmt.Errorf("failed to check if vault exists: %w", err)
	}

	if !exists {
		return "", "", "", fmt.Errorf("vault '%s' %w",
			vaultName, errVaultDoesNotExist)
	}

	// Check if vault already has a public key
	pubKeyPath := vaultDir + "/pub.age"

	_, err = cli.fs.Stat(pubKeyPath)
	if err == nil {
		return "", "", "", fmt.Errorf("vault '%s' %w",
			vaultName, errVaultHasLongTermKey)
	}

	// Get mnemonic from environment
	mnemonic := os.Getenv(secret.EnvMnemonic)
	if mnemonic == "" {
		return "", "", "", errMnemonicEnvNotSet
	}

	// Validate the mnemonic
	mnemonicWords := strings.Fields(mnemonic)
	secret.Debug("Validating BIP39 mnemonic", "word_count", len(mnemonicWords))

	if !bip39.IsMnemonicValid(mnemonic) {
		return "", "", "", errInvalidMnemonic
	}

	return vaultDir, pubKeyPath, mnemonic, nil
}

// updateVaultImportMetadata stores the derivation info in vault metadata
func updateVaultImportMetadata(
	fs afero.Fs, vaultDir string, derivationIndex uint32,
	publicKeyHash, familyHash string,
) error {
	// Load existing metadata
	existingMetadata, err := vault.LoadVaultMetadata(fs, vaultDir)
	if err != nil {
		// If metadata doesn't exist, create new
		existingMetadata = &vault.Metadata{
			CreatedAt: time.Now(),
		}
	}

	// Update metadata with new derivation info
	existingMetadata.DerivationIndex = derivationIndex
	existingMetadata.PublicKeyHash = publicKeyHash
	existingMetadata.MnemonicFamilyHash = familyHash

	err = vault.SaveVaultMetadata(fs, vaultDir, existingMetadata)
	if err != nil {
		secret.Debug("Failed to save vault metadata", "error", err)

		return fmt.Errorf("failed to save vault metadata: %w", err)
	}

	secret.Debug("Saved vault metadata with derivation index and public key hash")

	return nil
}

// VaultImport imports a mnemonic into a specific vault
func (cli *Instance) VaultImport(cmd *cobra.Command, vaultName string) error {
	secret.Debug("Importing mnemonic into vault",
		"vault_name", vaultName, "state_dir", cli.stateDir)

	// Get the specific vault by name
	vlt := vault.NewVault(cli.fs, cli.stateDir, vaultName)

	vaultDir, pubKeyPath, mnemonic, err := cli.vaultImportPreflight(vlt, vaultName)
	if err != nil {
		return err
	}

	// Get the next available derivation index for this mnemonic
	derivationIndex, err := vault.GetNextDerivationIndex(cli.fs, cli.stateDir, mnemonic)
	if err != nil {
		secret.Debug("Failed to get next derivation index", "error", err)

		return fmt.Errorf("failed to get next derivation index: %w", err)
	}

	secret.Debug("Using derivation index", "index", derivationIndex)

	// Derive long-term key from mnemonic with the appropriate index
	secret.Debug("Deriving long-term key from mnemonic", "index", derivationIndex)

	ltIdentity, err := agehd.DeriveIdentity(mnemonic, derivationIndex)
	if err != nil {
		return fmt.Errorf("failed to derive long-term key: %w", err)
	}

	// Store long-term public key in vault
	ltPublicKey := ltIdentity.Recipient().String()
	secret.Debug("Storing long-term public key",
		"pubkey", ltPublicKey, "vault_dir", vaultDir)

	err = afero.WriteFile(cli.fs, pubKeyPath, []byte(ltPublicKey), secret.FilePerms)
	if err != nil {
		return fmt.Errorf("failed to store long-term public key: %w", err)
	}

	// Calculate public key hash from the actual derivation index being used
	// This is used to verify that the derived key matches what was stored
	publicKeyHash := vault.ComputeDoubleSHA256([]byte(ltIdentity.Recipient().String()))

	// Calculate family hash from index 0 (same for all vaults with this
	// mnemonic). This is used to identify which vaults belong to the same
	// mnemonic family.
	identity0, err := agehd.DeriveIdentity(mnemonic, 0)
	if err != nil {
		return fmt.Errorf("failed to derive identity for index 0: %w", err)
	}

	familyHash := vault.ComputeDoubleSHA256([]byte(identity0.Recipient().String()))

	err = updateVaultImportMetadata(
		cli.fs, vaultDir, derivationIndex, publicKeyHash, familyHash)
	if err != nil {
		return err
	}

	// Get passphrase from environment variable
	passphraseStr := os.Getenv(secret.EnvUnlockPassphrase)
	if passphraseStr == "" {
		return errPassphraseEnvNotSet
	}

	secret.Debug("Using unlock passphrase from environment variable")

	// Create secure buffer for passphrase
	passphraseBuffer := memguard.NewBufferFromBytes([]byte(passphraseStr))
	defer passphraseBuffer.Destroy()

	// Unlock the vault with the derived long-term key
	vlt.Unlock(ltIdentity)

	// Create passphrase-protected unlocker
	secret.Debug("Creating passphrase-protected unlocker")

	passphraseUnlocker, err := vlt.CreatePassphraseUnlocker(passphraseBuffer)
	if err != nil {
		secret.Debug("Failed to create unlocker", "error", err)

		return fmt.Errorf("failed to create unlocker: %w", err)
	}

	cmd.Printf("Successfully imported mnemonic into vault '%s'\n", vaultName)
	cmd.Printf("Long-term public key: %s\n", ltPublicKey)
	cmd.Printf("Unlocker ID: %s\n", passphraseUnlocker.GetID())

	return nil
}

// vaultHasSecrets reports whether the vault directory contains any secrets
func (cli *Instance) vaultHasSecrets(vaultDir string) bool {
	secretsDir := filepath.Join(vaultDir, "secrets.d")

	exists, _ := afero.DirExists(cli.fs, secretsDir)
	if !exists {
		return false
	}

	entries, err := afero.ReadDir(cli.fs, secretsDir)

	return err == nil && len(entries) > 0
}

// switchAwayFromVault selects another vault as current before removal
func (cli *Instance) switchAwayFromVault(
	cmd *cobra.Command, vaults []string, name string,
) error {
	// Find another vault to switch to
	var newVault string

	for _, v := range vaults {
		if v != name {
			newVault = v

			break
		}
	}

	// Switch to the new vault
	err := vault.SelectVault(cli.fs, cli.stateDir, newVault)
	if err != nil {
		return fmt.Errorf("failed to switch to vault '%s': %w", newVault, err)
	}

	cmd.Printf("Switched current vault to '%s'\n", newVault)

	return nil
}

// RemoveVault removes a vault with safety checks
func (cli *Instance) RemoveVault(cmd *cobra.Command, name string, force bool) error {
	// Get list of all vaults
	vaults, err := vault.ListVaults(cli.fs, cli.stateDir)
	if err != nil {
		return fmt.Errorf("failed to list vaults: %w", err)
	}

	// Check if vault exists
	if !slices.Contains(vaults, name) {
		return fmt.Errorf("vault '%s' %w", name, errVaultDoesNotExist)
	}

	// Don't allow removing the last vault
	if len(vaults) == 1 {
		return errCannotRemoveLastVault
	}

	// Check if this is the current vault
	currentVault, err := vault.GetCurrentVault(cli.fs, cli.stateDir)
	if err != nil {
		return fmt.Errorf("failed to get current vault: %w", err)
	}

	isCurrentVault := currentVault.GetName() == name

	// Load the vault to check for secrets
	vlt := vault.NewVault(cli.fs, cli.stateDir, name)

	vaultDir, err := vlt.GetDirectory()
	if err != nil {
		return fmt.Errorf("failed to get vault directory: %w", err)
	}

	// Check if vault has secrets
	hasSecrets := cli.vaultHasSecrets(vaultDir)

	// Require --force if vault has secrets
	if hasSecrets && !force {
		return fmt.Errorf("vault '%s' %w", name, errVaultContainsSecrets)
	}

	// If removing current vault, switch to another vault first
	if isCurrentVault {
		err = cli.switchAwayFromVault(cmd, vaults, name)
		if err != nil {
			return err
		}
	}

	// Remove the vault directory
	err = cli.fs.RemoveAll(vaultDir)
	if err != nil {
		return fmt.Errorf("failed to remove vault directory: %w", err)
	}

	cmd.Printf("Removed vault '%s'\n", name)

	if hasSecrets {
		cmd.Printf("Warning: Vault contained secrets that have been " +
			"permanently deleted\n")
	}

	return nil
}