// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Package encryption handles encryption of the bundle using age.
package encryption

//go:generate go run default_recipients.go

import (
	"bufio"
	_ "embed"
	"errors"
	"fmt"
	"io"
	"strings"

	"filippo.io/age"
	"filippo.io/age/agessh"
)

// defaultRecipientsData is the embedded recipients file generated by
// default_recipients.go from the public SSH keys of the public members of the
// siderolabs GitHub organization.
//
//go:embed recipients.txt
var defaultRecipientsData string

// Recipient is one or more age recipients sharing a comment, together with the
// raw public keys they were parsed from, so callers can display who a bundle is
// going to be encrypted to.
//
// A single Recipient groups all the keys found under one comment in the
// recipients file (e.g. all SSH keys belonging to a single GitHub user).
type Recipient struct {
	recipients []age.Recipient

	// Comment is the annotation from the recipients file, e.g. the GitHub
	// username and real name ("alongwill (Andy Longwill)"). It may be empty.
	Comment string

	// PublicKeys are the raw recipient lines (SSH public keys or age "age1..."
	// recipients) this recipient was parsed from.
	PublicKeys []string
}

// String returns a human-readable representation of the recipient.
func (r Recipient) String() string {
	if r.Comment != "" {
		return r.Comment
	}

	if len(r.PublicKeys) > 0 {
		return r.PublicKeys[0]
	}

	return ""
}

// DefaultRecipients returns the list of recipients the bundle will be encrypted
// to by default, parsed from the embedded recipients file, grouped per comment.
//
// It is useful to display to the user what the bundle is going to be encrypted to.
func DefaultRecipients() ([]Recipient, error) {
	return parseRecipients(defaultRecipientsData)
}

// Option configures Encrypt.
type Option func(*config)

type config struct {
	recipients []string
	additional []string
	override   bool
}

// WithRecipients overrides the default recipients entirely: the bundle is
// encrypted only to the given recipients, not to the embedded defaults.
//
// Each recipient is a single recipient line: an SSH public key (ssh-rsa or
// ssh-ed25519) or an age "age1..." recipient.
func WithRecipients(recipients ...string) Option {
	return func(c *config) {
		c.override = true
		c.recipients = append(c.recipients, recipients...)
	}
}

// WithAdditionalRecipients adds recipients on top of the current set (the
// embedded defaults, or the ones set via WithRecipients).
//
// Each recipient is a single recipient line: an SSH public key (ssh-rsa or
// ssh-ed25519) or an age "age1..." recipient.
func WithAdditionalRecipients(recipients ...string) Option {
	return func(c *config) {
		c.additional = append(c.additional, recipients...)
	}
}

// Encrypt wraps dst with an age encryption layer and returns an io.WriteCloser
// that encrypts everything written to it.
//
// With no options, the bundle is encrypted to the default recipients (embedded
// from recipients.txt). Use WithRecipients to replace the defaults entirely, and
// WithAdditionalRecipients to add extra recipients.
//
// The returned io.WriteCloser must be closed to flush the encrypted output; the
// underlying dst is not closed.
func Encrypt(dst io.Writer, opts ...Option) (io.WriteCloser, error) {
	var cfg config

	for _, o := range opts {
		o(&cfg)
	}

	var recipients []age.Recipient

	if cfg.override {
		parsed, err := parseRecipientLines(cfg.recipients)
		if err != nil {
			return nil, err
		}

		recipients = append(recipients, parsed...)
	} else {
		defaults, err := DefaultRecipients()
		if err != nil {
			return nil, fmt.Errorf("failed to load default recipients: %w", err)
		}

		for _, r := range defaults {
			recipients = append(recipients, r.recipients...)
		}
	}

	additional, err := parseRecipientLines(cfg.additional)
	if err != nil {
		return nil, err
	}

	recipients = append(recipients, additional...)

	if len(recipients) == 0 {
		return nil, errors.New("no recipients to encrypt to")
	}

	return age.Encrypt(dst, recipients...)
}

// parseRecipientLines parses a list of recipient lines into age recipients.
func parseRecipientLines(lines []string) ([]age.Recipient, error) {
	recipients := make([]age.Recipient, 0, len(lines))

	for _, line := range lines {
		r, err := parseRecipient(line)
		if err != nil {
			return nil, fmt.Errorf("failed to parse recipient %q: %w", line, err)
		}

		recipients = append(recipients, r)
	}

	return recipients, nil
}

// parseRecipients parses an age recipients file: blank lines are ignored, lines
// starting with "#" are comments, and every other line is a recipient. All the
// recipients following a comment, up to the next comment, are grouped into a
// single Recipient carrying that comment.
func parseRecipients(data string) ([]Recipient, error) {
	var (
		recipients  []Recipient
		comment     string
		current     Recipient
		haveCurrent bool
	)

	flush := func() {
		if haveCurrent && len(current.recipients) > 0 {
			recipients = append(recipients, current)
		}

		haveCurrent = false
	}

	scanner := bufio.NewScanner(strings.NewReader(data))
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())

		switch {
		case line == "":
			continue
		case strings.HasPrefix(line, "#"):
			flush()

			comment = strings.TrimSpace(strings.TrimPrefix(line, "#"))

			continue
		}

		r, err := parseRecipient(line)
		if err != nil {
			return nil, fmt.Errorf("failed to parse recipient %q: %w", line, err)
		}

		if !haveCurrent {
			current = Recipient{Comment: comment}
			haveCurrent = true
		}

		current.recipients = append(current.recipients, r)
		current.PublicKeys = append(current.PublicKeys, line)
	}

	flush()

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("failed to read recipients: %w", err)
	}

	return recipients, nil
}

// parseRecipient parses a single recipient line: an age "age1..." recipient or
// an SSH public key.
func parseRecipient(line string) (age.Recipient, error) {
	if strings.HasPrefix(line, "age1") {
		return age.ParseX25519Recipient(line)
	}

	return agessh.ParseRecipient(line)
}
