package mfer

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"slices"
	"strings"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/armor"
	pgperrors "github.com/ProtonMail/go-crypto/openpgp/errors"
	"github.com/ProtonMail/go-crypto/openpgp/packet"
)

const (
	// The tags of OpenPGP signature, key and subkey packets (RFC 9580,
	// section 5). Subkeys have tags of their own, so each secret or public
	// key packet is one primary key.
	signaturePacketTag    = 2
	secretKeyPacketTag    = 5
	publicKeyPacketTag    = 6
	secretSubkeyPacketTag = 7
	publicSubkeyPacketTag = 14

	// In the body of a key or subkey packet the algorithm octet follows
	// the version octet and the four-octet creation time, and from version
	// 5 on a four-octet length as well (RFC 9580, section 5.5.2).
	keyAlgorithmOffset            = 5
	firstKeyVersionWithLength     = 5
	keyAlgorithmOffsetAfterLength = 9

	// signingKeyVersion is the only OpenPGP key version mfer signs with.
	// Its fingerprints are 40 hex characters, the length
	// --require-signature takes.
	signingKeyVersion = 4

	// armorBegin and armorEnd start the lines that begin and end an
	// armored block.
	armorBegin = "-----BEGIN "
	armorEnd   = "-----END "
)

var (
	errKeyCount    = errors.New("must hold exactly one key")
	errDSAKey      = errors.New("must not hold a DSA key")
	errNoSecretKey = errors.New("signing key file holds no secret key")
	errNotV4Key    = errors.New("signing key must be an OpenPGP version 4 key, " +
		"the only kind whose fingerprint --require-signature takes")
	errNoPassphrase = errors.New(
		"signing key is protected and no passphrase was given")
	errNotOneSignature = errors.New(
		"signature must hold exactly one signature")
	errNotOneArmoredBlock = errors.New(
		"must be exactly one armored block and nothing else")
	errMalformedArmor = errors.New("armor is malformed")
)

// SigningOptions holds the key a manifest is signed with.
type SigningOptions struct {
	// SecretKey is an OpenPGP secret key, armored or binary, as
	// gpg --export-secret-keys writes it. It must hold one primary key.
	SecretKey []byte
	// Passphrase unlocks SecretKey when it is protected.
	Passphrase []byte
}

// SecretKeyIsProtected reports whether the OpenPGP secret key secretKey,
// armored or binary, needs a passphrase to sign. It fails unless
// secretKey holds one version 4 primary key with its secret key.
func SecretKeyIsProtected(secretKey []byte) (bool, error) {
	key, err := readSecretKey(secretKey)
	if err != nil {
		return false, err
	}

	return isProtected(key), nil
}

// CheckSigningKey fails unless opts can sign now: opts.SecretKey must hold
// one version 4 primary key with a secret key that may sign and has not
// expired or been revoked, and opts.Passphrase must unlock it when it is
// protected. It lets a caller find a key that cannot sign before it builds
// a manifest.
func CheckSigningKey(opts *SigningOptions) error {
	key, err := readSigningKey(opts)
	if err != nil {
		return err
	}

	// Signing nothing fails wherever signing the manifest would.
	err = openpgp.DetachSign(io.Discard, key, bytes.NewReader(nil), nil)
	if err != nil {
		return fmt.Errorf("signing key cannot sign: %w", err)
	}

	return nil
}

// readSigningKey returns the key in opts.SecretKey, unlocked with
// opts.Passphrase if it is protected.
func readSigningKey(opts *SigningOptions) (*openpgp.Entity, error) {
	key, err := readSecretKey(opts.SecretKey)
	if err != nil {
		return nil, err
	}

	if !isProtected(key) {
		return key, nil
	}

	if len(opts.Passphrase) == 0 {
		return nil, errNoPassphrase
	}

	err = key.DecryptPrivateKeys(opts.Passphrase)
	if err != nil {
		return nil, fmt.Errorf("unlock signing key: %w", err)
	}

	return key, nil
}

// readSecretKey returns the one key in secretKey, armored or binary,
// which must be a version 4 key and include its secret key.
func readSecretKey(secretKey []byte) (*openpgp.Entity, error) {
	key, err := readOneKey(secretKey, "signing key file")
	if err != nil {
		return nil, err
	}

	if key.PrimaryKey.Version != signingKeyVersion {
		return nil, fmt.Errorf("%w; this key is version %d",
			errNotV4Key, key.PrimaryKey.Version)
	}

	if key.PrivateKey == nil {
		return nil, errNoSecretKey
	}

	return key, nil
}

// readOneKey returns the key in data, armored or binary, which must hold
// exactly one primary key and no DSA key or subkey. what names data in
// errors.
func readOneKey(data []byte, what string) (*openpgp.Entity, error) {
	packets, err := dearmor(data)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", what, err)
	}

	keys, err := countPackets(packets, secretKeyPacketTag, publicKeyPacketTag)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", what, err)
	}

	if keys != 1 {
		return nil, fmt.Errorf("%s %w, found %d", what, errKeyCount, keys)
	}

	// openpgp.ReadKeyRing checks every self-signature, and a DSA key with
	// very large numbers makes each check take seconds to minutes.
	dsa, err := holdsDSAKey(packets)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", what, err)
	}

	if dsa {
		return nil, fmt.Errorf("%s %w", what, errDSAKey)
	}

	keyring, err := openpgp.ReadKeyRing(bytes.NewReader(packets))
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", what, err)
	}

	// openpgp.ReadKeyRing also reads a subkey packet at the start as a
	// primary key.
	if len(keyring) != 1 {
		return nil, fmt.Errorf("%s %w, found %d", what, errKeyCount, len(keyring))
	}

	return keyring[0], nil
}

// isProtected reports whether any secret key in key needs a passphrase.
func isProtected(key *openpgp.Entity) bool {
	if key.PrivateKey.Encrypted {
		return true
	}

	for _, subkey := range key.Subkeys {
		if subkey.PrivateKey != nil && subkey.PrivateKey.Encrypted {
			return true
		}
	}

	return false
}

// armoredPublicKey returns the public part of key, armored.
func armoredPublicKey(key *openpgp.Entity) ([]byte, error) {
	var buf bytes.Buffer

	w, err := armor.Encode(&buf, openpgp.PublicKeyType, nil)
	if err != nil {
		return nil, err
	}

	err = key.Serialize(w)
	if err != nil {
		return nil, fmt.Errorf("write public key: %w", err)
	}

	err = w.Close()
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// fingerprint returns the fingerprint of key's primary key in upper-case
// hex, as gpg prints it.
func fingerprint(key *openpgp.Entity) string {
	return strings.ToUpper(hex.EncodeToString(key.PrimaryKey.Fingerprint))
}

// verifySignature checks that signature is one good OpenPGP signature
// over data, made by the one primary key in pubKey or one of its subkeys,
// and returns that primary key's fingerprint. signature and pubKey may each
// be armored or binary.
func verifySignature(data, signature, pubKey []byte) (string, error) {
	key, err := readOneKey(pubKey, "embedded public key block")
	if err != nil {
		return "", err
	}

	sigData, err := dearmor(signature)
	if err != nil {
		return "", fmt.Errorf("read signature: %w", err)
	}

	sigs, err := countPackets(sigData, signaturePacketTag)
	if err != nil {
		return "", fmt.Errorf("read signature: %w", err)
	}

	if sigs != 1 {
		return "", fmt.Errorf("%w, found %d", errNotOneSignature, sigs)
	}

	_, err = openpgp.CheckDetachedSignature(openpgp.EntityList{key},
		bytes.NewReader(data), bytes.NewReader(sigData), nil)
	// A manifest outlives its signing key, so a signature by a key that
	// has expired since is still good.
	if err != nil && !errors.Is(err, pgperrors.ErrKeyExpired) {
		return "", fmt.Errorf("verify signature: %w", err)
	}

	return fingerprint(key), nil
}

// dearmor returns the binary OpenPGP data in data: data itself when it is
// not armored, or else the body of its armored block. Armored data must be
// one block and nothing else: its first line is the only BEGIN line and
// its last line the only END line, white space around them aside.
// armor.Decode skips any text before a BEGIN line and reads only the first
// block, so without this a second key or signature would go unseen.
func dearmor(data []byte) ([]byte, error) {
	if !bytes.Contains(data, []byte(armorBegin)) {
		return data, nil
	}

	text := bytes.TrimSpace(data)
	lastLine := text[bytes.LastIndexByte(text, '\n')+1:]

	if !bytes.HasPrefix(text, []byte(armorBegin)) ||
		bytes.Count(text, []byte(armorBegin)) != 1 ||
		!bytes.HasPrefix(lastLine, []byte(armorEnd)) ||
		bytes.Count(text, []byte(armorEnd)) != 1 {
		return nil, errNotOneArmoredBlock
	}

	// armor.Decode passes over a block it cannot read, such as one with a
	// header line that has no colon, and returns io.EOF on finding no
	// other.
	block, err := armor.Decode(bytes.NewReader(text))
	if err != nil {
		return nil, errMalformedArmor
	}

	body, err := io.ReadAll(block.Body)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", errMalformedArmor, err)
	}

	return body, nil
}

// countPackets returns how many packets in the binary OpenPGP data have
// one of tags. It reads only each packet's header, so it also counts
// packets that openpgp.ReadKeyRing skips, such as a key with no user ID
// or of an algorithm it does not know.
func countPackets(data []byte, tags ...uint8) (int, error) {
	packets := packet.NewOpaqueReader(bytes.NewReader(data))
	count := 0

	for {
		p, err := packets.Next()
		if errors.Is(err, io.EOF) {
			return count, nil
		}

		if err != nil {
			return 0, err
		}

		if slices.Contains(tags, p.Tag) {
			count++
		}
	}
}

// holdsDSAKey reports whether any key or subkey packet in the binary
// OpenPGP data holds a DSA key. It reads each packet's algorithm octet
// rather than parsing the packet, since parsing a secret key packet checks
// its numbers, which for a DSA key with very large numbers is as slow as
// checking a self-signature.
func holdsDSAKey(data []byte) (bool, error) {
	packets := packet.NewOpaqueReader(bytes.NewReader(data))

	for {
		p, err := packets.Next()
		if errors.Is(err, io.EOF) {
			return false, nil
		}

		if err != nil {
			return false, err
		}

		if !slices.Contains([]uint8{
			secretKeyPacketTag, publicKeyPacketTag,
			secretSubkeyPacketTag, publicSubkeyPacketTag,
		}, p.Tag) {
			continue
		}

		offset := keyAlgorithmOffset
		if len(p.Contents) > 0 && p.Contents[0] >= firstKeyVersionWithLength {
			offset = keyAlgorithmOffsetAfterLength
		}

		if len(p.Contents) > offset &&
			packet.PublicKeyAlgorithm(p.Contents[offset]) == packet.PubKeyAlgoDSA {
			return true, nil
		}
	}
}