URBM/internal/restic/restic.go

package restic

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"git.casaderoll.de/michael/urbm/internal/model"
)

type SecretGetter interface{ Get(string) (string, error) }

type Runner struct {
	Binary     string
	RuntimeDir string
	Secrets    SecretGetter
	Log        *slog.Logger
	mu         sync.Mutex
	processes  map[string]*os.Process
}

type runIDKey struct{}

func WithRunID(ctx context.Context, runID string) context.Context {
	return context.WithValue(ctx, runIDKey{}, runID)
}

func (r *Runner) Pause(runID string) error  { return r.signal(runID, syscall.SIGSTOP) }
func (r *Runner) Resume(runID string) error { return r.signal(runID, syscall.SIGCONT) }

func (r *Runner) signal(runID string, signal syscall.Signal) error {
	r.mu.Lock()
	process := r.processes[runID]
	r.mu.Unlock()
	if process == nil {
		return errors.New("restic process is not active yet")
	}
	if err := syscall.Kill(-process.Pid, signal); err != nil {
		return fmt.Errorf("signal restic process group: %w", err)
	}
	return nil
}

func (r *Runner) registerProcess(runID string, process *os.Process) func() {
	if runID == "" {
		return func() {}
	}
	r.mu.Lock()
	if r.processes == nil {
		r.processes = make(map[string]*os.Process)
	}
	r.processes[runID] = process
	r.mu.Unlock()
	return func() {
		r.mu.Lock()
		delete(r.processes, runID)
		r.mu.Unlock()
	}
}

type Summary struct {
	MessageType         string `json:"message_type"`
	SnapshotID          string `json:"snapshot_id"`
	DataAdded           int64  `json:"data_added"`
	FilesNew            int64  `json:"files_new"`
	FilesChanged        int64  `json:"files_changed"`
	TotalBytesProcessed int64  `json:"total_bytes_processed"`
	TotalFilesProcessed int64  `json:"total_files_processed"`
}

type Progress struct {
	MessageType      string   `json:"message_type"`
	PercentDone      float64  `json:"percent_done"`
	TotalFiles       int64    `json:"total_files"`
	FilesDone        int64    `json:"files_done"`
	TotalBytes       int64    `json:"total_bytes"`
	BytesDone        int64    `json:"bytes_done"`
	SecondsRemaining int64    `json:"seconds_remaining"`
	CurrentFiles     []string `json:"current_files"`
}

type ProgressCallback func(Progress)

type Stats struct {
	TotalSize      int64 `json:"total_size"`
	TotalFileCount int64 `json:"total_file_count"`
	TotalBlobCount int64 `json:"total_blob_count"`
}

func (r *Runner) Init(ctx context.Context, repo model.Repository) error {
	return r.run(ctx, repo, []string{"init"}, nil, nil)
}

func (r *Runner) Unlock(ctx context.Context, repo model.Repository) error {
	return r.run(ctx, repo, []string{"unlock"}, nil, nil)
}

func (r *Runner) ChangePassword(ctx context.Context, repo model.Repository, currentPassword, newPassword string) error {
	newPasswordPath, cleanup, err := r.writePasswordFile("restic-new-password-*", newPassword)
	if err != nil {
		return err
	}
	defer cleanup()
	return r.runWithInputEnvPassword(ctx, repo, []string{"key", "passwd", "--new-password-file", newPasswordPath}, nil, nil, nil, nil, currentPassword)
}

func (r *Runner) TestPassword(ctx context.Context, repo model.Repository, password string) error {
	return r.runWithInputEnvPassword(ctx, repo, []string{"snapshots", "--json"}, nil, nil, nil, nil, password)
}

func (r *Runner) Backup(ctx context.Context, repo model.Repository, job model.Job, sources []string, progress ProgressCallback) (Summary, error) {
	args := []string{"backup", "--json", "--compression", job.Compression}
	for _, tag := range append([]string{"urbm", "job:" + job.ID}, job.Tags...) {
		args = append(args, "--tag", tag)
	}
	for _, exclude := range job.Excludes {
		args = append(args, "--exclude", exclude)
	}
	args = append(args, sources...)
	var summary Summary
	err := r.runWithEnv(ctx, repo, args, func(line []byte) {
		var status Progress
		if json.Unmarshal(line, &status) == nil && status.MessageType == "status" {
			if progress != nil {
				progress(status)
			}
			return
		}
		var candidate Summary
		if json.Unmarshal(line, &candidate) == nil && candidate.MessageType == "summary" {
			summary = candidate
		}
	}, nil, backupEnvironment(job))
	return summary, err
}

func (r *Runner) BackupImage(ctx context.Context, repo model.Repository, job model.Job, image io.Reader, progress ProgressCallback) (Summary, error) {
	args := []string{"backup", "--json", "--stdin", "--stdin-filename", "urbm-usb-flash.img", "--compression", job.Compression}
	for _, tag := range append([]string{"urbm", "job:" + job.ID, "usb-image"}, job.Tags...) {
		args = append(args, "--tag", tag)
	}
	var summary Summary
	err := r.runWithInputAndEnv(ctx, repo, args, func(line []byte) {
		var status Progress
		if json.Unmarshal(line, &status) == nil && status.MessageType == "status" {
			if progress != nil {
				progress(status)
			}
			return
		}
		var candidate Summary
		if json.Unmarshal(line, &candidate) == nil && candidate.MessageType == "summary" {
			summary = candidate
		}
	}, nil, image, backupEnvironment(job))
	return summary, err
}

func (r *Runner) Forget(ctx context.Context, repo model.Repository, job model.Job) error {
	args := []string{"forget", "--tag", "job:" + job.ID}
	if job.Retention.KeepWithinDays > 0 {
		args = append(args, "--keep-within", fmt.Sprintf("%dd", job.Retention.KeepWithinDays))
	}
	if job.Retention.Daily > 0 {
		args = append(args, "--keep-daily", fmt.Sprint(job.Retention.Daily))
	}
	if job.Retention.Weekly > 0 {
		args = append(args, "--keep-weekly", fmt.Sprint(job.Retention.Weekly))
	}
	if job.Retention.Monthly > 0 {
		args = append(args, "--keep-monthly", fmt.Sprint(job.Retention.Monthly))
	}
	err := r.run(ctx, repo, args, nil, nil)
	if err == nil || !isRepositoryLocked(err) {
		return err
	}
	if unlockErr := r.Unlock(ctx, repo); unlockErr != nil {
		return fmt.Errorf("retention repository locked; remove stale lock: %w", unlockErr)
	}
	return r.run(ctx, repo, args, nil, nil)
}

func (r *Runner) Check(ctx context.Context, repo model.Repository) error {
	return r.run(ctx, repo, []string{"check"}, nil, nil)
}
func (r *Runner) Prune(ctx context.Context, repo model.Repository) error {
	return r.run(ctx, repo, []string{"prune"}, nil, nil)
}

func (r *Runner) Snapshots(ctx context.Context, repo model.Repository) ([]model.Snapshot, error) {
	var output []byte
	err := r.run(ctx, repo, []string{"snapshots", "--json"}, nil, &output)
	if err != nil {
		return nil, err
	}
	var snapshots []model.Snapshot
	if err := json.Unmarshal(output, &snapshots); err != nil {
		return nil, fmt.Errorf("decode snapshots: %w", err)
	}
	return snapshots, nil
}

func (r *Runner) Stats(ctx context.Context, repo model.Repository, mode string) (Stats, error) {
	args := []string{"stats", "--json"}
	if mode != "" {
		args = append(args, "--mode", mode)
	}
	stats, err := r.stats(ctx, repo, args)
	if err == nil || !isRepositoryLocked(err) {
		return stats, err
	}
	if unlockErr := r.Unlock(ctx, repo); unlockErr != nil {
		return Stats{}, fmt.Errorf("repository statistics locked; remove stale lock: %w", unlockErr)
	}
	return r.stats(ctx, repo, args)
}

func (r *Runner) stats(ctx context.Context, repo model.Repository, args []string) (Stats, error) {
	var output []byte
	if err := r.run(ctx, repo, args, nil, &output); err != nil {
		return Stats{}, err
	}
	stats, err := decodeStats(output)
	if err != nil {
		return Stats{}, fmt.Errorf("decode repository stats: %w", err)
	}
	return stats, nil
}

func decodeStats(output []byte) (Stats, error) {
	for index := len(output) - 1; index >= 0; index-- {
		if output[index] != '{' {
			continue
		}
		var stats Stats
		decoder := json.NewDecoder(bytes.NewReader(output[index:]))
		if err := decoder.Decode(&stats); err == nil {
			return stats, nil
		}
	}
	return Stats{}, errors.New("no valid statistics object in Restic output")
}

func (r *Runner) List(ctx context.Context, repo model.Repository, snapshot, path string) ([]map[string]any, error) {
	args := []string{"ls", snapshot, "--json"}
	if path != "" {
		args = append(args, path)
	}
	var output []byte
	if err := r.run(ctx, repo, args, nil, &output); err != nil {
		return nil, err
	}
	items := []map[string]any{}
	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	for scanner.Scan() {
		var item map[string]any
		if json.Unmarshal(scanner.Bytes(), &item) == nil && item["struct_type"] == "node" {
			items = append(items, item)
		}
		if len(items) >= 5000 {
			break
		}
	}
	return items, scanner.Err()
}

func (r *Runner) Restore(ctx context.Context, repo model.Repository, task model.RestoreTask) error {
	args := []string{"restore", task.SnapshotID, "--target", task.Target}
	for _, include := range task.Includes {
		args = append(args, "--include", include)
	}
	return r.run(ctx, repo, args, nil, nil)
}

func (r *Runner) run(ctx context.Context, repo model.Repository, args []string, onLine func([]byte), capture *[]byte) error {
	return r.runWithInputAndEnv(ctx, repo, args, onLine, capture, nil, nil)
}

func (r *Runner) runWithEnv(ctx context.Context, repo model.Repository, args []string, onLine func([]byte), capture *[]byte, environment []string) error {
	return r.runWithInputAndEnv(ctx, repo, args, onLine, capture, nil, environment)
}

func (r *Runner) runWithInputAndEnv(ctx context.Context, repo model.Repository, args []string, onLine func([]byte), capture *[]byte, input io.Reader, environment []string) error {
	password, err := r.Secrets.Get(repo.PasswordRef)
	if err != nil {
		return fmt.Errorf("authentication: load repository password: %w", err)
	}
	return r.runWithInputEnvPassword(ctx, repo, args, onLine, capture, input, environment, password)
}

func (r *Runner) runWithInputEnvPassword(ctx context.Context, repo model.Repository, args []string, onLine func([]byte), capture *[]byte, input io.Reader, environment []string, password string) error {
	passwordPath, cleanup, err := r.writePasswordFile("restic-password-*", password)
	if err != nil {
		return err
	}
	defer cleanup()

	secretDir := filepath.Join(r.RuntimeDir, "secrets")
	globalArgs := []string{"-r", repositoryLocation(repo)}
	credentialPath := ""
	if repo.Type == model.RepositorySFTP && repo.CredentialRef != "" {
		privateKey, err := r.Secrets.Get(repo.CredentialRef)
		if err != nil {
			return fmt.Errorf("authentication: load SFTP private key: %w", err)
		}
		keyFile, err := os.CreateTemp(secretDir, "sftp-key-*")
		if err != nil {
			return err
		}
		credentialPath = keyFile.Name()
		defer os.Remove(credentialPath)
		if err := keyFile.Chmod(0600); err != nil {
			keyFile.Close()
			return err
		}
		if _, err := io.WriteString(keyFile, privateKey); err != nil {
			keyFile.Close()
			return err
		}
		if err := keyFile.Close(); err != nil {
			return err
		}
		knownHosts := repo.Options["knownHostsPath"]
		command := "ssh -i " + credentialPath + " -o IdentitiesOnly=yes -o UserKnownHostsFile=" + knownHosts + " -o StrictHostKeyChecking=yes"
		globalArgs = append(globalArgs, "-o", "sftp.command="+command)
	}
	fullArgs := append(globalArgs, args...)
	cmd := exec.CommandContext(ctx, r.Binary, fullArgs...)
	cmd.Stdin = input
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		if cmd.Process == nil {
			return os.ErrProcessDone
		}
		if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
			if errors.Is(err, syscall.ESRCH) {
				return os.ErrProcessDone
			}
			return err
		}
		return nil
	}
	cmd.Env = append([]string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "HOME=/root", "RESTIC_PASSWORD_FILE=" + passwordPath}, environment...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return classify(err)
	}
	runID, _ := ctx.Value(runIDKey{}).(string)
	unregister := r.registerProcess(runID, cmd.Process)
	defer unregister()
	var captured []byte
	done := make(chan struct{})
	go func() {
		scanner := bufio.NewScanner(stdout)
		scanner.Buffer(make([]byte, 64*1024), 2*1024*1024)
		for scanner.Scan() {
			line := append([]byte(nil), scanner.Bytes()...)
			if capture != nil {
				captured = append(captured, line...)
				captured = append(captured, '\n')
			}
			if onLine != nil {
				onLine(line)
			}
		}
		close(done)
	}()
	errBytes, _ := io.ReadAll(io.LimitReader(stderr, 1024*1024))
	err = cmd.Wait()
	<-done
	if capture != nil {
		*capture = captured
	}
	if err != nil {
		message := strings.TrimSpace(redact(string(errBytes), repo))
		if errors.Is(ctx.Err(), context.Canceled) {
			return fmt.Errorf("cancelled: %w", ctx.Err())
		}
		return resticError(message, err)
	}
	return nil
}

func (r *Runner) writePasswordFile(pattern, password string) (string, func(), error) {
	secretDir := filepath.Join(r.RuntimeDir, "secrets")
	if err := os.MkdirAll(secretDir, 0700); err != nil {
		return "", nil, err
	}
	f, err := os.CreateTemp(secretDir, pattern)
	if err != nil {
		return "", nil, err
	}
	passwordPath := f.Name()
	cleanup := func() { _ = os.Remove(passwordPath) }
	if err := f.Chmod(0600); err != nil {
		f.Close()
		cleanup()
		return "", nil, err
	}
	if _, err := io.WriteString(f, password); err != nil {
		f.Close()
		cleanup()
		return "", nil, err
	}
	if err := f.Close(); err != nil {
		cleanup()
		return "", nil, err
	}
	return passwordPath, cleanup, nil
}

func backupEnvironment(job model.Job) []string {
	if job.CPUCores <= 0 {
		return nil
	}
	return []string{"GOMAXPROCS=" + fmt.Sprint(job.CPUCores)}
}

func resticError(message string, commandErr error) error {
	lower := strings.ToLower(message)
	switch {
	case strings.Contains(lower, "config file already exists") || strings.Contains(lower, "repository master key and config already initialized"):
		return errors.New("repository: destination is already initialized; do not select Initialize again. Enter the existing repository password and select Test")
	case strings.Contains(lower, "ciphertext verification failed"):
		return errors.New("authentication: the saved repository password does not match this existing repository, or its Restic key is damaged. Enter the original repository password under Edit and select Test")
	case strings.Contains(lower, "wrong password") || strings.Contains(lower, "no key found"):
		return errors.New("authentication: wrong repository password. Enter the original Restic password under Edit and select Test")
	case strings.Contains(lower, "repository does not exist") || strings.Contains(lower, "is there a repository at the following location?"):
		return errors.New("repository: Am konfigurierten Ziel wurde keine Restic-config gefunden. Für ein vorhandenes Repository den exakten Repository-Ordner eintragen; nur ein wirklich neues, leeres Ziel initialisieren")
	default:
		message = strings.TrimSpace(strings.TrimPrefix(message, "Fatal:"))
		return fmt.Errorf("restic: %s: %w", message, commandErr)
	}
}

func isRepositoryLocked(err error) bool {
	if err == nil {
		return false
	}
	lower := strings.ToLower(err.Error())
	return strings.Contains(lower, "repository is already locked") ||
		strings.Contains(lower, "unable to create lock in backend")
}

func repositoryLocation(repo model.Repository) string {
	if repo.Type == model.RepositorySFTP && !strings.HasPrefix(repo.Location, "sftp:") {
		return "sftp:" + repo.Location
	}
	return repo.Location
}

func redact(value string, repo model.Repository) string {
	value = strings.ReplaceAll(value, repo.Location, "[repository]")
	return value
}

func classify(err error) error {
	if errors.Is(err, exec.ErrNotFound) {
		return fmt.Errorf("environment: restic binary not found: %w", err)
	}
	return err
}