// Copyright 2019 The age Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"iter"
	"os"
	"path/filepath"
	"regexp"
	"runtime/debug"
	"slices"
	"strings"
	"unicode"
	"unicode/utf8"

	"filippo.io/age"
	"filippo.io/age/agessh"
	"filippo.io/age/armor"
	"filippo.io/age/internal/term"
	"filippo.io/age/plugin"
)

const usage = `Usage:
    age [--encrypt] (-r RECIPIENT | -R PATH)... [--armor] [-o OUTPUT] [INPUT]
    age [--encrypt] --passphrase [--armor] [-o OUTPUT] [INPUT]
    age --decrypt [-i PATH]... [-o OUTPUT] [INPUT]

Options:
    -e, --encrypt               Encrypt the input to the output. Default if omitted.
    -d, --decrypt               Decrypt the input to the output.
    -o, --output OUTPUT         Write the result to the file at path OUTPUT.
    -a, --armor                 Encrypt to a PEM encoded format.
    -p, --passphrase            Encrypt with a passphrase.
    -r, --recipient RECIPIENT   Encrypt to the specified RECIPIENT. Can be repeated.
    -R, --recipients-file PATH  Encrypt to recipients listed at PATH. Can be repeated.
    -i, --identity PATH         Use the identity file at PATH. Can be repeated.
    --version                   Print the version.

INPUT defaults to standard input, and OUTPUT defaults to standard output.
If OUTPUT exists, it will be overwritten.

RECIPIENT can be an age public key generated by age-keygen ("age1...")
or an SSH public key ("ssh-ed25519 AAAA...", "ssh-rsa AAAA...").

Recipient files contain one or more recipients, one per line. Empty lines
and lines starting with "#" are ignored as comments. "-" may be used to
read recipients from standard input.

Identity files contain one or more secret keys ("AGE-SECRET-KEY-1..."),
one per line, or an SSH key. Empty lines and lines starting with "#" are
ignored as comments. Passphrase encrypted age files can be used as
identity files. Multiple key files can be provided, and any unused ones
will be ignored. "-" may be used to read identities from standard input.

When --encrypt is specified explicitly, -i can also be used to encrypt to an
identity file symmetrically, instead or in addition to normal recipients.

Example:
    $ age-keygen -o key.txt
    Public key: age1ql3z7hjy54pw3hyww5ayyfg7zqgvc7w3j2elw8zmrj2kg5sfn9aqmcac8p
    $ tar cvz ~/data | age -r age1ql3z7hjy54pw3hyww5ayyfg7zqgvc7w3j2elw8zmrj2kg5sfn9aqmcac8p > data.tar.gz.age
    $ age --decrypt -i key.txt -o data.tar.gz data.tar.gz.age`

// stdinInUse is used to ensure only one of input, recipients, or identities
// file is read from stdin. It's a singleton like os.Stdin.
var stdinInUse bool

type multiFlag []string

func (f *multiFlag) String() string { return fmt.Sprint(*f) }

func (f *multiFlag) Set(value string) error {
	*f = append(*f, value)
	return nil
}

type identityFlag struct {
	Type, Value string
}

// identityFlags tracks -i and -j flags, preserving their relative order, so
// that "age -d -j agent -i encrypted-fallback-keys.age" behaves as expected.
type identityFlags []identityFlag

func (f *identityFlags) addIdentityFlag(value string) error {
	*f = append(*f, identityFlag{Type: "i", Value: value})
	return nil
}

func (f *identityFlags) addPluginFlag(value string) error {
	*f = append(*f, identityFlag{Type: "j", Value: value})
	return nil
}

// Version can be set at link time to override debug.BuildInfo.Main.Version when
// building manually without git history. It should look like "v1.2.3".
var Version string

func main() {
	flag.Usage = func() { fmt.Fprintf(os.Stderr, "%s\n", usage) }

	if len(os.Args) == 1 {
		flag.Usage()
		os.Exit(1)
	}

	var (
		outFlag                          string
		decryptFlag, encryptFlag         bool
		passFlag, versionFlag, armorFlag bool
		recipientFlags                   multiFlag
		recipientsFileFlags              multiFlag
		identityFlags                    identityFlags
	)

	flag.BoolVar(&versionFlag, "version", false, "print the version")
	flag.BoolVar(&decryptFlag, "d", false, "decrypt the input")
	flag.BoolVar(&decryptFlag, "decrypt", false, "decrypt the input")
	flag.BoolVar(&encryptFlag, "e", false, "encrypt the input")
	flag.BoolVar(&encryptFlag, "encrypt", false, "encrypt the input")
	flag.BoolVar(&passFlag, "p", false, "use a passphrase")
	flag.BoolVar(&passFlag, "passphrase", false, "use a passphrase")
	flag.StringVar(&outFlag, "o", "", "output to `FILE` (default stdout)")
	flag.StringVar(&outFlag, "output", "", "output to `FILE` (default stdout)")
	flag.BoolVar(&armorFlag, "a", false, "generate an armored file")
	flag.BoolVar(&armorFlag, "armor", false, "generate an armored file")
	flag.Var(&recipientFlags, "r", "recipient (can be repeated)")
	flag.Var(&recipientFlags, "recipient", "recipient (can be repeated)")
	flag.Var(&recipientsFileFlags, "R", "recipients file (can be repeated)")
	flag.Var(&recipientsFileFlags, "recipients-file", "recipients file (can be repeated)")
	flag.Func("i", "identity (can be repeated)", identityFlags.addIdentityFlag)
	flag.Func("identity", "identity (can be repeated)", identityFlags.addIdentityFlag)
	flag.Func("j", "data-less plugin (can be repeated)", identityFlags.addPluginFlag)
	flag.Parse()

	if versionFlag {
		if buildInfo, ok := debug.ReadBuildInfo(); ok && Version == "" {
			Version = buildInfo.Main.Version
		}
		fmt.Println(Version)
		return
	}

	if flag.NArg() > 1 {
		var hints []string
		quotedArgs := strings.Trim(fmt.Sprintf("%q", flag.Args()), "[]")

		// If the second argument looks like a flag, suggest moving the first
		// argument to the back (as long as the arguments don't need quoting).
		if strings.HasPrefix(flag.Arg(1), "-") {
			hints = append(hints, "the input file must be specified after all flags")

			safe := true
			unsafeShell := regexp.MustCompile(`[^\w@%+=:,./-]`)
			if slices.ContainsFunc(os.Args, unsafeShell.MatchString) {
				safe = false
			}
			if safe {
				i := len(os.Args) - flag.NArg()
				newArgs := append([]string{}, os.Args[:i]...)
				newArgs = append(newArgs, os.Args[i+1:]...)
				newArgs = append(newArgs, os.Args[i])
				hints = append(hints, "did you mean:")
				hints = append(hints, "    "+strings.Join(newArgs, " "))
			}
		} else {
			hints = append(hints, "only a single input file may be specified at a time")
		}

		errorWithHint("too many INPUT arguments: "+quotedArgs, hints...)
	}

	switch {
	case decryptFlag:
		if encryptFlag {
			errorf("-e/--encrypt can't be used with -d/--decrypt")
		}
		if armorFlag {
			errorWithHint("-a/--armor can't be used with -d/--decrypt",
				"note that armored files are detected automatically, try again without -a/--armor")
		}
		if passFlag {
			errorWithHint("-p/--passphrase can't be used with -d/--decrypt",
				"note that password protected files are detected automatically")
		}
		if len(recipientFlags) > 0 {
			errorWithHint("-r/--recipient can't be used with -d/--decrypt",
				"did you mean to use -i/--identity to specify a private key?")
		}
		if len(recipientsFileFlags) > 0 {
			errorWithHint("-R/--recipients-file can't be used with -d/--decrypt",
				"did you mean to use -i/--identity to specify a private key?")
		}
	default: // encrypt
		if len(identityFlags) > 0 && !encryptFlag {
			errorWithHint("-i/--identity and -j can't be used in encryption mode unless symmetric encryption is explicitly selected with -e/--encrypt",
				"did you forget to specify -d/--decrypt?")
		}
		if len(recipientFlags)+len(recipientsFileFlags)+len(identityFlags) == 0 && !passFlag {
			errorWithHint("missing recipients",
				"did you forget to specify -r/--recipient, -R/--recipients-file or -p/--passphrase?")
		}
		if len(recipientFlags) > 0 && passFlag {
			errorf("-p/--passphrase can't be combined with -r/--recipient")
		}
		if len(recipientsFileFlags) > 0 && passFlag {
			errorf("-p/--passphrase can't be combined with -R/--recipients-file")
		}
		if len(identityFlags) > 0 && passFlag {
			errorf("-p/--passphrase can't be combined with -i/--identity and -j")
		}
	}

	warnDuplicates(slices.Values(recipientFlags), "recipient")
	warnDuplicates(slices.Values(recipientsFileFlags), "recipients file")
	warnDuplicates(func(yield func(string) bool) {
		for _, f := range identityFlags {
			if f.Type == "i" && !yield(f.Value) {
				return
			}
		}
	}, "identity file")

	var inUseFiles []string
	for _, i := range identityFlags {
		if i.Type != "i" {
			continue
		}
		inUseFiles = append(inUseFiles, absPath(i.Value))
	}
	for _, f := range recipientsFileFlags {
		inUseFiles = append(inUseFiles, absPath(f))
	}

	var in io.Reader = os.Stdin
	var out io.Writer = os.Stdout
	if name := flag.Arg(0); name != "" && name != "-" {
		inUseFiles = append(inUseFiles, absPath(name))
		f, err := os.Open(name)
		if err != nil {
			errorf("failed to open input file %q: %v", name, err)
		}
		defer f.Close()
		in = f
	} else {
		stdinInUse = true
		if decryptFlag && term.IsTerminal(os.Stdin) {
			// If the input comes from a TTY, assume it's armored, and buffer up
			// to the END line (or EOF/EOT) so that a password prompt or the
			// output don't get in the way of typing the input. See Issue 364.
			buf, err := bufferTerminalInput(in)
			if err != nil {
				errorf("failed to buffer terminal input: %v", err)
			}
			in = buf
		}
	}
	if name := outFlag; name != "" && name != "-" {
		outFI, outErr := os.Stat(name)
		for _, f := range inUseFiles {
			same := f == absPath(name)
			if !same && outErr == nil {
				if fi, err := os.Stat(f); err == nil {
					// Catch symlinks and hard links.
					same = os.SameFile(fi, outFI)
				}
			}
			if same {
				errorf("input and output file are the same: %q", name)
			}
		}
		f := newLazyOpener(name)
		defer func() {
			if err := f.Close(); err != nil {
				errorf("failed to close output file %q: %v", name, err)
			}
		}()
		out = f
	} else if term.IsTerminal(os.Stdout) {
		buf := &bytes.Buffer{}
		defer func() {
			if out == buf {
				io.Copy(os.Stdout, buf)
			}
		}()
		if name != "-" {
			if decryptFlag {
				// Buffer the output to check it's printable.
				out = buf
				defer func() {
					if !utf8.Valid(buf.Bytes()) || bytes.ContainsFunc(buf.Bytes(), func(r rune) bool {
						return r != '\n' && r != '\r' && r != '\t' && unicode.IsControl(r)
					}) {
						errorWithHint("refusing to output binary to the terminal",
							`force anyway with "-o -"`)
					}
				}()
			} else if !armorFlag {
				// If the output wouldn't be armored, refuse to send binary to
				// the terminal unless explicitly requested with "-o -".
				errorWithHint("refusing to output binary to the terminal",
					"did you mean to use -a/--armor?",
					`force anyway with "-o -"`)
			}
		}
		if in == os.Stdin && term.IsTerminal(os.Stdin) {
			// If the input comes from a TTY and output will go to a TTY,
			// buffer it up so it doesn't get in the way of typing the input.
			out = buf
		}
	}

	switch {
	case decryptFlag && len(identityFlags) == 0:
		decryptPass(in, out)
	case decryptFlag:
		decryptNotPass(identityFlags, in, out)
	case passFlag:
		encryptPass(in, out, armorFlag)
	default:
		encryptNotPass(recipientFlags, recipientsFileFlags, identityFlags, in, out, armorFlag)
	}
}

func passphrasePromptForEncryption() (string, error) {
	pass, err := term.ReadSecret("Enter passphrase (leave empty to autogenerate a secure one):")
	if err != nil {
		return "", fmt.Errorf("could not read passphrase: %v", err)
	}
	p := string(pass)
	if p == "" {
		var words []string
		for range 10 {
			words = append(words, randomWord())
		}
		p = strings.Join(words, "-")
		err := printfToTerminal("using autogenerated passphrase %q", p)
		if err != nil {
			return "", fmt.Errorf("could not print passphrase: %v", err)
		}
	} else {
		confirm, err := term.ReadSecret("Confirm passphrase:")
		if err != nil {
			return "", fmt.Errorf("could not read passphrase: %v", err)
		}
		if string(confirm) != p {
			return "", fmt.Errorf("passphrases didn't match")
		}
	}
	return p, nil
}

func encryptNotPass(recs, files []string, identities identityFlags, in io.Reader, out io.Writer, armor bool) {
	var recipients []age.Recipient
	for _, arg := range recs {
		r, err := parseRecipient(arg)
		if err, ok := err.(gitHubRecipientError); ok {
			errorWithHint(err.Error(), "instead, use recipient files like",
				"    curl -O https://github.com/"+err.username+".keys",
				"    age -R "+err.username+".keys")
		}
		if err != nil {
			errorf("%v", err)
		}
		recipients = append(recipients, r)
	}
	for _, name := range files {
		recs, err := parseRecipientsFile(name)
		if err != nil {
			errorf("failed to parse recipient file %q: %v", name, err)
		}
		recipients = append(recipients, recs...)
	}
	for _, f := range identities {
		switch f.Type {
		case "i":
			ids, err := parseIdentitiesFile(f.Value)
			if err != nil {
				errorf("reading %q: %v", f.Value, err)
			}
			r, err := identitiesToRecipients(ids)
			if err != nil {
				errorf("internal error processing %q: %v", f.Value, err)
			}
			recipients = append(recipients, r...)
		case "j":
			id, err := plugin.NewIdentityWithoutData(f.Value, plugin.NewTerminalUI(printf, warningf))
			if err != nil {
				errorf("initializing %q: %v", f.Value, err)
			}
			recipients = append(recipients, id.Recipient())
		}
	}
	encrypt(recipients, in, out, armor)
}

func encryptPass(in io.Reader, out io.Writer, armor bool) {
	pass, err := passphrasePromptForEncryption()
	if err != nil {
		errorf("%v", err)
	}

	r, err := age.NewScryptRecipient(pass)
	if err != nil {
		errorf("%v", err)
	}
	testOnlyConfigureScryptIdentity(r)
	encrypt([]age.Recipient{r}, in, out, armor)
}

var testOnlyConfigureScryptIdentity = func(*age.ScryptRecipient) {}

func encrypt(recipients []age.Recipient, in io.Reader, out io.Writer, withArmor bool) {
	if withArmor {
		a := armor.NewWriter(out)
		defer func() {
			if err := a.Close(); err != nil {
				errorf("%v", err)
			}
		}()
		out = a
	}
	w, err := age.Encrypt(out, recipients...)
	if e := new(plugin.NotFoundError); errors.As(err, &e) {
		errorWithHint(err.Error(),
			fmt.Sprintf("you might want to install the %q plugin", e.Name),
			"visit https://age-encryption.org/awesome#plugins for a list of available plugins")
	} else if err != nil {
		errorf("%v", err)
	}
	if _, err := io.Copy(w, in); err != nil {
		errorf("%v", err)
	}
	if err := w.Close(); err != nil {
		errorf("%v", err)
	}
}

// crlfMangledIntro and utf16MangledIntro are the intro lines of the age format
// after mangling by various versions of PowerShell redirection, truncated to
// the length of the correct intro line. See issue 290.
const crlfMangledIntro = "age-encryption.org/v1" + "\r"
const utf16MangledIntro = "\xff\xfe" + "a\x00g\x00e\x00-\x00e\x00n\x00c\x00r\x00y\x00p\x00"

type rejectScryptIdentity struct{}

func (rejectScryptIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if len(stanzas) != 1 || stanzas[0].Type != "scrypt" {
		return nil, age.ErrIncorrectIdentity
	}
	errorWithHint("file is passphrase-encrypted but identities were specified with -i/--identity or -j",
		"remove all -i/--identity/-j flags to decrypt passphrase-encrypted files")
	panic("unreachable")
}

func decryptNotPass(flags identityFlags, in io.Reader, out io.Writer) {
	var identities []age.Identity
	for _, f := range flags {
		switch f.Type {
		case "i":
			ids, err := parseIdentitiesFile(f.Value)
			if err != nil {
				errorf("reading %q: %v", f.Value, err)
			}
			identities = append(identities, ids...)
		case "j":
			id, err := plugin.NewIdentityWithoutData(f.Value, plugin.NewTerminalUI(printf, warningf))
			if err != nil {
				errorf("initializing %q: %v", f.Value, err)
			}
			identities = append(identities, id)
		}
	}
	identities = append(identities, rejectScryptIdentity{})
	decrypt(identities, in, out)
}

func decryptPass(in io.Reader, out io.Writer) {
	identities := []age.Identity{
		// If there is an scrypt recipient (it will have to be the only one and)
		// this identity will be invoked.
		lazyScryptIdentity,
	}

	decrypt(identities, in, out)
}

func decrypt(identities []age.Identity, in io.Reader, out io.Writer) {
	rr := bufio.NewReader(in)
	if intro, _ := rr.Peek(len(crlfMangledIntro)); string(intro) == crlfMangledIntro ||
		string(intro) == utf16MangledIntro {
		errorWithHint("invalid header intro",
			"it looks like this file was corrupted by PowerShell redirection",
			"consider using -o or -a to encrypt files in PowerShell")
	}

	const maxWhitespace = 1024
	start, _ := rr.Peek(maxWhitespace + len(armor.Header))
	if strings.HasPrefix(string(bytes.TrimSpace(start)), armor.Header) {
		in = armor.NewReader(rr)
	} else {
		in = rr
	}

	r, err := age.Decrypt(in, identities...)
	if e := new(plugin.NotFoundError); errors.As(err, &e) {
		errorWithHint(err.Error(),
			fmt.Sprintf("you might want to install the %q plugin", e.Name),
			"visit https://age-encryption.org/awesome#plugins for a list of available plugins")
	} else if errors.As(err, new(*age.NoIdentityMatchError)) &&
		len(identities) == 1 && identities[0] == lazyScryptIdentity {
		errorWithHint("the file is not passphrase-encrypted, identities are required",
			"specify identities with -i/--identity or -j to decrypt this file")
	} else if err != nil {
		errorf("%v", err)
	}
	if _, err := io.Copy(out, r); err != nil {
		errorf("%v", err)
	}
	// Trigger the lazyOpener even if r is empty, if Copy succeeded.
	if _, err := out.Write(nil); err != nil {
		errorf("%v", err)
	}
}

var lazyScryptIdentity = &LazyScryptIdentity{passphrasePromptForDecryption}

func passphrasePromptForDecryption() (string, error) {
	pass, err := term.ReadSecret("Enter passphrase:")
	if err != nil {
		return "", fmt.Errorf("could not read passphrase: %v", err)
	}
	return string(pass), nil
}

func identitiesToRecipients(ids []age.Identity) ([]age.Recipient, error) {
	var recipients []age.Recipient
	for _, id := range ids {
		switch id := id.(type) {
		case *age.X25519Identity:
			recipients = append(recipients, id.Recipient())
		case *age.HybridIdentity:
			recipients = append(recipients, id.Recipient())
		case *plugin.Identity:
			recipients = append(recipients, id.Recipient())
		case *agessh.RSAIdentity:
			recipients = append(recipients, id.Recipient())
		case *agessh.Ed25519Identity:
			recipients = append(recipients, id.Recipient())
		case *agessh.EncryptedSSHIdentity:
			recipients = append(recipients, id.Recipient())
		case *EncryptedIdentity:
			r, err := id.Recipients()
			if err != nil {
				return nil, err
			}
			recipients = append(recipients, r...)
		default:
			return nil, fmt.Errorf("unexpected identity type: %T", id)
		}
	}
	return recipients, nil
}

type lazyOpener struct {
	name string
	f    *os.File
	err  error
}

func newLazyOpener(name string) io.WriteCloser {
	return &lazyOpener{name: name}
}

func (l *lazyOpener) Write(p []byte) (n int, err error) {
	if l.f == nil && l.err == nil {
		l.f, l.err = os.Create(l.name)
	}
	if l.err != nil {
		return 0, l.err
	}
	return l.f.Write(p)
}

func (l *lazyOpener) Close() error {
	if l.f != nil {
		return l.f.Close()
	}
	return l.err
}

func absPath(name string) string {
	if abs, err := filepath.Abs(name); err == nil {
		return abs
	}
	return name
}

func warnDuplicates(s iter.Seq[string], name string) {
	seen := make(map[string]bool)
	warned := make(map[string]bool)
	for e := range s {
		if seen[e] && !warned[e] {
			warningf("duplicate %s %q", name, e)
			warned[e] = true
		}
		seen[e] = true
	}
}
