package smime

import (
	"crypto"
	"crypto/x509"
	"errors"
	"fmt"
	"mime"
	"strings"

	"go.mozilla.org/pkcs7"
)

// Encrypt wraps entity's canonical bytes as CMS EnvelopedData addressed to
// recipients, per RFC 8551 application/pkcs7-mime; smime-type=enveloped-data. Pass
// every recipient's certificate, including the sender's own, so a copy kept in Sent
// stays readable.
func Encrypt(entity Entity, recipients []*x509.Certificate) (Entity, error) {
	if len(recipients) == 0 {
		return Entity{}, errors.New("smime: no recipient certificates provided")
	}
	envelopedDER, err := pkcs7.Encrypt(entity.bytes(), recipients)
	if err != nil {
		return Entity{}, fmt.Errorf("smime: encrypt: %w", err)
	}
	return Entity{
		Headers: []string{
			`Content-Type: application/pkcs7-mime; smime-type=enveloped-data; name="smime.p7m"`,
			"Content-Transfer-Encoding: base64",
			`Content-Disposition: attachment; filename="smime.p7m"`,
		},
		Body: []byte(wrapBase64(envelopedDER)),
	}, nil
}

// Decrypt reverses Encrypt, returning the inner MIME entity that was originally
// wrapped.
func Decrypt(entity Entity, cert *x509.Certificate, key crypto.PrivateKey) (Entity, error) {
	ct := HeaderValue(entity.Headers, "Content-Type")
	mediaType, params, err := mime.ParseMediaType(ct)
	if err != nil || mediaType != "application/pkcs7-mime" {
		return Entity{}, errors.New("smime: not an application/pkcs7-mime message")
	}
	if st := params["smime-type"]; st != "" && !strings.EqualFold(st, "enveloped-data") {
		return Entity{}, fmt.Errorf("smime: unsupported smime-type %q", st)
	}

	raw := entity.Body
	if isBase64CTE(HeaderValue(entity.Headers, "Content-Transfer-Encoding")) {
		if raw, err = decodeBase64(entity.Body); err != nil {
			return Entity{}, fmt.Errorf("smime: decode envelope: %w", err)
		}
	}

	p7, err := pkcs7.Parse(raw)
	if err != nil {
		return Entity{}, fmt.Errorf("smime: parse envelope: %w", err)
	}
	plaintext, err := p7.Decrypt(cert, key)
	if err != nil {
		return Entity{}, fmt.Errorf("smime: decrypt: %w", err)
	}
	return parseEntity(plaintext)
}