keyfunc/internal/cli/age/age.go

// Package age groups the commands that derive age identities and
// encrypt and decrypt with them.
package age

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"
	"sneak.berlin/go/keyfunc/internal/agekey"
	"sneak.berlin/go/keyfunc/internal/cli/options"
	"sneak.berlin/go/keyfunc/internal/cli/signals"
	"sneak.berlin/go/keyfunc/internal/derive"
)

// ErrInterrupted is returned when SIGINT, SIGTERM or SIGHUP came before
// the file --output names was put in place.
var ErrInterrupted = errors.New(
	"interrupted by a signal; the output file was left as it was",
)

// signalWait is how long after the work has ended a signal still keeps
// the new file from being put in place. Ctrl-C on "producer | keyfunc
// age encrypt -o file" ends the producer as well, and the end of the
// input can reach the work a moment before the signal reaches the tool.
const signalWait = 100 * time.Millisecond

// Command returns the age command and everything under it.
func Command() *cobra.Command {
	group := &cobra.Command{
		Use:   "age",
		Short: "derive age identities and encrypt and decrypt with them",
	}

	group.AddCommand(public(), private(), encrypt(), decrypt())

	return group
}

// public returns the command that prints the recipient.
func public() *cobra.Command {
	return &cobra.Command{
		Use:   "pub",
		Short: "print the recipient, the age1... public key",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			key, err := derived(cmd)
			if err != nil {
				return err
			}

			return write(cmd, key.Recipient())
		},
	}
}

// private returns the command that prints the identity.
func private() *cobra.Command {
	return &cobra.Command{
		Use:   "priv",
		Short: "print the identity, the AGE-SECRET-KEY-1... line",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			key, err := derived(cmd)
			if err != nil {
				return err
			}

			return write(cmd, key.Identity())
		},
	}
}

// encrypt returns the command that encrypts a file or standard input.
func encrypt() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "encrypt [file]",
		Short: "encrypt to the derived recipient and any others given",
		Args:  cobra.MaximumNArgs(1),
		RunE:  runEncrypt,
	}

	cmd.Flags().StringArray(
		"to", nil,
		"another recipient to encrypt to, as well as the derived one",
	)
	cmd.Flags().Bool(
		"armor", false,
		"write the text form instead of the binary one",
	)
	addOutput(cmd)

	return cmd
}

// decrypt returns the command that decrypts a file or standard input.
func decrypt() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "decrypt [file]",
		Short: "decrypt with the derived identity",
		Args:  cobra.MaximumNArgs(1),
		RunE:  runDecrypt,
	}

	addOutput(cmd)

	return cmd
}

// runEncrypt encrypts to the derived recipient and any others given.
func runEncrypt(cmd *cobra.Command, args []string) error {
	key, err := derived(cmd)
	if err != nil {
		return err
	}

	to, err := cmd.Flags().GetStringArray("to")
	if err != nil {
		return fmt.Errorf("reading the recipients: %w", err)
	}

	armored, err := cmd.Flags().GetBool("armor")
	if err != nil {
		return fmt.Errorf("reading the armor flag: %w", err)
	}

	return through(cmd, args, func(dst io.Writer, src io.Reader) error {
		return key.Encrypt(dst, src, to, armored)
	})
}

// runDecrypt decrypts with the derived identity.
func runDecrypt(cmd *cobra.Command, args []string) error {
	key, err := derived(cmd)
	if err != nil {
		return err
	}

	return through(cmd, args, key.Decrypt)
}

// through opens the input the arguments ask for and hands it to the
// work, with the file --output names to write to, or the command's own
// output when it names none.
func through(
	cmd *cobra.Command, args []string,
	work func(io.Writer, io.Reader) error,
) error {
	src, closeSrc, err := input(cmd, args)
	if err != nil {
		return err
	}

	defer closeSrc()

	name, err := cmd.Flags().GetString("output")
	if err != nil {
		return fmt.Errorf("reading the output file: %w", err)
	}

	if name == "" {
		return work(cmd.OutOrStdout(), src)
	}

	return output(name, src, work)
}

// input returns what to read from: the named file, or the command's
// own input when no file is named. The second result closes a file
// that was opened and does nothing otherwise.
func input(cmd *cobra.Command, args []string) (io.Reader, func(), error) {
	if len(args) == 0 {
		return cmd.InOrStdin(), func() {}, nil
	}

	file, err := os.Open(args[0])
	if err != nil {
		return nil, nil, fmt.Errorf("opening %s: %w", args[0], err)
	}

	return file, func() { _ = file.Close() }, nil
}

// output has the work write a new file beside the named one, and puts
// the new file in the named file's place only when the work succeeded,
// so a file that is already there survives a run that failed.
//
// Meanwhile SIGINT, SIGTERM and SIGHUP are caught, as signals.Context
// does. One that comes while the work runs, or within signalWait after
// it has ended, wins: the new file is removed and ErrInterrupted
// returned at once, without waiting for the work, which may be blocked
// reading its input.
func output(
	name string, src io.Reader, work func(io.Writer, io.Reader) error,
) error {
	interrupted, stop := signals.Context(context.Background())
	defer stop()

	// The file is made in the same directory so that putting it in
	// place is a rename and never a copy, and it is readable only by
	// its owner, which is the mode it keeps once renamed.
	file, err := os.CreateTemp(filepath.Dir(name), filepath.Base(name)+".")
	if err != nil {
		return fmt.Errorf("creating a file beside %s: %w", name, err)
	}

	worked := make(chan error, 1)

	go func() { worked <- work(file, src) }()

	select {
	case failed := <-worked:
		select {
		case <-interrupted.Done():
		case <-time.After(signalWait):
			return finish(file, name, failed)
		}
	case <-interrupted.Done():
	}

	return finish(file, name, ErrInterrupted)
}

// finish closes the new file and puts it in the named file's place, or
// throws it away when the work failed. It returns the error the caller
// should report.
func finish(file *os.File, name string, failed error) error {
	closeErr := file.Close()

	if failed != nil || closeErr != nil {
		_ = os.Remove(file.Name())

		if failed != nil {
			return failed
		}

		return fmt.Errorf("finishing %s: %w", name, closeErr)
	}

	err := os.Rename(file.Name(), name)
	if err != nil {
		_ = os.Remove(file.Name())

		return fmt.Errorf("putting %s in place: %w", name, err)
	}

	return nil
}

// addOutput gives a command its output file flag.
func addOutput(cmd *cobra.Command) {
	cmd.Flags().StringP(
		"output", "o", "",
		"write to this file instead of standard output",
	)
}

// write sends one line to wherever the command's output goes.
func write(cmd *cobra.Command, line string) error {
	_, err := fmt.Fprintln(cmd.OutOrStdout(), line)
	if err != nil {
		return fmt.Errorf("writing the key: %w", err)
	}

	return nil
}

// derived returns the age key for this run.
func derived(cmd *cobra.Command) (*agekey.Key, error) {
	index, err := options.Index(cmd)
	if err != nil {
		return nil, err
	}

	words, err := options.Mnemonic(cmd)
	if err != nil {
		return nil, err
	}

	material, err := derive.Bytes(words, agekey.Application, index)
	if err != nil {
		return nil, err
	}

	return agekey.New(material)
}