vaultik/internal/cli/snapshot_restore.go

package cli

import (
	"context"
	"errors"
	"os"

	"github.com/spf13/cobra"
	"go.uber.org/fx"
	"sneak.berlin/go/vaultik/internal/config"
	"sneak.berlin/go/vaultik/internal/globals"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/storage"
	"sneak.berlin/go/vaultik/internal/vaultik"
)

// restoreMinArgs is the minimum positional argument count of
// `snapshot restore <snapshot-id> <target-dir> [paths...]`.
const restoreMinArgs = 2

// RestoreOptions contains options for the restore command
type RestoreOptions struct {
	TargetDir string
	Paths     []string // Optional paths to restore (empty = all)
	Verify    bool     // Verify restored files after restore
}

// RestoreApp contains all dependencies needed for restore
type RestoreApp struct {
	Globals    *globals.Globals
	Config     *config.Config
	Storage    storage.Storer
	Vaultik    *vaultik.Vaultik
	Shutdowner fx.Shutdowner
}

// newSnapshotRestoreCommand creates the 'snapshot restore' subcommand
func newSnapshotRestoreCommand() *cobra.Command {
	opts := &RestoreOptions{}

	cmd := &cobra.Command{
		Use:   "restore <snapshot-id> <target-dir> [paths...]",
		Short: "Restore files from a snapshot",
		Long: `Download and decrypt files from a backup snapshot.

This command will restore files from the specified snapshot to the
target directory.
If no paths are specified, all files are restored.
If paths are specified, only matching files/directories are restored.

Requires the VAULTIK_AGE_SECRET_KEY environment variable to be set with
the age private key.

Examples:
  # Restore entire snapshot
  vaultik snapshot restore myhost_docs_2025-01-01T12:00:00Z /restore

  # Restore specific file
  vaultik snapshot restore myhost_docs_2025-01-01T12:00:00Z /restore \
    /home/user/important.txt

  # Restore specific directory
  vaultik snapshot restore myhost_docs_2025-01-01T12:00:00Z /restore \
    /home/user/documents/

  # Restore and verify all files
  vaultik snapshot restore --verify myhost_docs_2025-01-01T12:00:00Z /restore`,
		Args: cobra.MinimumNArgs(restoreMinArgs),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runRestore(cmd, args, opts)
		},
	}

	cmd.Flags().BoolVar(&opts.Verify, "verify", false,
		"Verify restored files by checking chunk hashes")

	return cmd
}

// runRestore parses arguments and runs the restore operation through the app framework
func runRestore(cmd *cobra.Command, args []string, opts *RestoreOptions) error {
	snapshotID := args[0]

	opts.TargetDir = args[1]
	if len(args) > restoreMinArgs {
		opts.Paths = args[restoreMinArgs:]
	}

	// Use unified config resolution
	configPath, err := ResolveConfigPath()
	if err != nil {
		return err
	}

	// Use the app framework like other commands
	rootFlags := GetRootFlags()

	return RunWithApp(cmd.Context(), AppOptions{
		ConfigPath: configPath,
		LogOptions: log.Options{
			Verbose: rootFlags.Verbose,
			Debug:   rootFlags.Debug,
			Quiet:   rootFlags.Quiet,
		},
		Modules: buildRestoreModules(),
		Invokes: buildRestoreInvokes(snapshotID, opts),
	})
}

// buildRestoreModules returns the fx.Options for dependency injection in restore
func buildRestoreModules() []fx.Option {
	return []fx.Option{
		fx.Provide(fx.Annotate(
			func(g *globals.Globals, cfg *config.Config,
				storer storage.Storer, v *vaultik.Vaultik, shutdowner fx.Shutdowner) *RestoreApp {
				return &RestoreApp{
					Globals:    g,
					Config:     cfg,
					Storage:    storer,
					Vaultik:    v,
					Shutdowner: shutdowner,
				}
			},
		)),
	}
}

// buildRestoreInvokes returns the fx.Options that wire up the restore lifecycle
func buildRestoreInvokes(snapshotID string, opts *RestoreOptions) []fx.Option {
	return []fx.Option{
		fx.Invoke(func(app *RestoreApp, lc fx.Lifecycle) {
			lc.Append(fx.Hook{
				OnStart: func(_ context.Context) error {
					// Start the restore operation in a goroutine
					go func() {
						// Run the restore operation
						restoreOpts := &vaultik.RestoreOptions{
							SnapshotID: snapshotID,
							TargetDir:  opts.TargetDir,
							Paths:      opts.Paths,
							Verify:     opts.Verify,
							SkipErrors: GetRootFlags().SkipErrors,
						}

						err := app.Vaultik.Restore(restoreOpts)
						if err != nil {
							if !errors.Is(err, context.Canceled) {
								log.Error("Restore operation failed", "error", err)
								ReportErrorf("Restore failed: %v", err)
								os.Exit(1)
							}
						}

						// Shutdown the app when restore completes
						err = app.Shutdowner.Shutdown()
						if err != nil {
							log.Error("Failed to shutdown", "error", err)
						}
					}()

					return nil
				},
				OnStop: func(_ context.Context) error {
					log.Debug("Stopping restore operation")
					app.Vaultik.Cancel()

					return nil
				},
			})
		}),
	}
}