// Package cmstest builds the CMS signatures and the RFC 3161 tokens that the
// tests of internal/cms and of capsule read: certificates of test keys, a
// detached signature of a message with its signedAttrs and, optionally, a
// time-stamp token of its signature. It is the encoder that a signing
// application has; nothing outside tests uses it.
package cmstest

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"math/big"
	"slices"
	"time"
)

// Signer is a certificate with its private key.
type Signer struct {
	Cert *x509.Certificate
	Key  crypto.Signer
}

// NewRSA returns a signer with an RSA key of bits bits, valid in
// [notBefore, notAfter], named cn.
func NewRSA(cn string, bits int, notBefore, notAfter time.Time) Signer {
	k, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		panic(err)
	}
	return newSigner(cn, k, notBefore, notAfter)
}

// NewECDSA returns a signer with an ECDSA key on curve.
func NewECDSA(cn string, curve elliptic.Curve, notBefore, notAfter time.Time) Signer {
	k, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		panic(err)
	}
	return newSigner(cn, k, notBefore, notAfter)
}

func newSigner(cn string, k crypto.Signer, notBefore, notAfter time.Time) Signer {
	serial, _ := rand.Int(rand.Reader, big.NewInt(1<<62))
	t := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: cn, Organization: []string{"DateKeys test"}},
		NotBefore:    notBefore, NotAfter: notAfter,
		KeyUsage:     x509.KeyUsageDigitalSignature,
		SubjectKeyId: []byte(cn),
	}
	raw, err := x509.CreateCertificate(rand.Reader, t, t, k.Public(), k)
	if err != nil {
		panic(err)
	}
	c, err := x509.ParseCertificate(raw)
	if err != nil {
		panic(err)
	}
	return Signer{Cert: c, Key: k}
}

// The DER building blocks.

func tlv(tag byte, content ...[]byte) []byte {
	c := bytes.Join(content, nil)
	out := []byte{tag}
	switch n := len(c); {
	case n < 0x80:
		out = append(out, byte(n))
	case n < 0x100:
		out = append(out, 0x81, byte(n))
	case n < 0x10000:
		out = append(out, 0x82, byte(n>>8), byte(n))
	default:
		out = append(out, 0x83, byte(n>>16), byte(n>>8), byte(n))
	}
	return append(out, c...)
}

// TLV builds an element of any tag from the encodings of its content.
func TLV(tag byte, content ...[]byte) []byte { return tlv(tag, content...) }

// TokenRaw returns a token that tsa signs over the given TSTInfo, as it is:
// for tests that need a TSTInfo that Token would not write.
func TokenRaw(tstInfo []byte, tsa Signer) []byte {
	return build(tstInfo, Options{Hash: crypto.SHA256}, true, []Signer{tsa})
}

// GeneralizedTime builds a GeneralizedTime with the text s.
func GeneralizedTime(s string) []byte { return tlv(0x18, []byte(s)) }

// Seq is a SEQUENCE.
func Seq(content ...[]byte) []byte { return tlv(0x30, content...) }

// Set is a SET OF, in DER order.
func Set(tag byte, elems ...[]byte) []byte {
	e := slices.Clone(elems)
	slices.SortFunc(e, bytes.Compare)
	return tlv(tag, e...)
}

// OID is an OBJECT IDENTIFIER.
func OID(oid asn1.ObjectIdentifier) []byte {
	b, err := asn1.Marshal(oid)
	if err != nil {
		panic(err)
	}
	return b
}

// Octets is an OCTET STRING.
func Octets(b []byte) []byte { return tlv(0x04, b) }

// Int is an INTEGER.
func Int(n int64) []byte {
	b, err := asn1.Marshal(n)
	if err != nil {
		panic(err)
	}
	return b
}

var (
	OIDData       = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 1}
	OIDSignedData = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 2}
	oidContent    = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 3}
	oidDigest     = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 4}
	OIDSigCertV2  = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 16, 2, 47}
	oidSigCertV1  = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 16, 2, 12}
	OIDTimeStamp  = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 16, 2, 14}
	oidTSTInfo    = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 9, 16, 1, 4}
	oidOCSP       = asn1.ObjectIdentifier{1, 3, 6, 1, 5, 5, 7, 16, 2}
	oidSHA256     = asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 1}
	oidSHA384     = asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 2}
	oidSHA512     = asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 3}
	oidRSA        = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1}
	oidPSS        = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 10}
	oidMGF1       = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 8}
	oidECDSA256   = asn1.ObjectIdentifier{1, 2, 840, 10045, 4, 3, 2}
	oidECDSA384   = asn1.ObjectIdentifier{1, 2, 840, 10045, 4, 3, 3}
	oidECDSA512   = asn1.ObjectIdentifier{1, 2, 840, 10045, 4, 3, 4}
)

func hashAlg(h crypto.Hash) []byte {
	switch h {
	case crypto.SHA384:
		return Seq(OID(oidSHA384))
	case crypto.SHA512:
		return Seq(OID(oidSHA512))
	}
	return Seq(OID(oidSHA256))
}

func hashOID(h crypto.Hash) asn1.ObjectIdentifier {
	switch h {
	case crypto.SHA384:
		return oidSHA384
	case crypto.SHA512:
		return oidSHA512
	}
	return oidSHA256
}

func sum(h crypto.Hash, b []byte) []byte {
	x := h.New()
	x.Write(b)
	return x.Sum(nil)
}

// Options changes what Build writes, to make signatures that the profile
// rejects or that verify to something else.
type Options struct {
	// Hash is the digest of the signature, SHA-256 by default.
	Hash crypto.Hash
	// PSS signs with RSASSA-PSS instead of PKCS #1 v1.5.
	PSS bool
	// SKI names the signer by subjectKeyIdentifier instead of by issuer
	// and serial.
	SKI bool
	// Token, when not nil, is the time-stamp token of the signature that
	// Build wants as an unsigned attribute: it receives the signature value.
	Token func(signature []byte) []byte
	// Message is what the message-digest covers, when not the signed message.
	Message []byte
	// Mutate edits the signedAttrs, as a list of the DER of each attribute,
	// before they are signed.
	Mutate func(attrs [][]byte) [][]byte
	// Token2 puts two signature-time-stamp attributes, to break the profile.
	Token2 bool
	// OCSP adds this response in crls.
	OCSP []byte
	// OmitCert leaves the certificate of the signer out of certificates.
	OmitCert bool
	// PSSTrailer writes trailerField [3] 1 in the PSS parameters, which is
	// its default and DER does not write it.
	PSSTrailer bool
	// Junk adds an unsigned attribute of this many bytes, which decides
	// nothing, to make a signature as large as a chain of certificates.
	Junk int
	// SigCertV1 adds a signing-certificate attribute beside the v2.
	SigCertV1 bool
	// ExtraAttrs are added to the signed attributes, as the DER of each.
	ExtraAttrs [][]byte
}

func attr(oid asn1.ObjectIdentifier, values ...[]byte) []byte {
	return Seq(OID(oid), Set(0x31, values...))
}

// Signature returns the detached CMS signature of message by the signers, in
// DER, as AutoFirma writes it: one SignerInfo each, with content-type,
// message-digest and signing-certificate-v2 as signed attributes.
func Signature(message []byte, opts Options, signers ...Signer) []byte {
	return build(message, opts, false, signers)
}

func build(message []byte, o Options, token bool, signers []Signer) []byte {
	if o.Hash == 0 {
		o.Hash = crypto.SHA256
	}
	var certs, infos [][]byte
	for _, s := range signers {
		if !o.OmitCert {
			certs = append(certs, s.Cert.Raw)
		}
		infos = append(infos, signerInfo(message, o, token, s))
	}
	var body []byte
	body = append(body, Int(1)...)
	body = append(body, Set(0x31, hashAlg(o.Hash))...)
	body = append(body, encap(message, token)...)
	if len(certs) > 0 {
		body = append(body, Set(0xa0, certs...)...)
	}
	if o.OCSP != nil {
		body = append(body, Set(0xa1, tlv(0xa1, OID(oidOCSP), o.OCSP))...)
	}
	body = append(body, Set(0x31, infos...)...)
	return Seq(OID(OIDSignedData), tlv(0xa0, Seq(body)))
}

func encap(content []byte, token bool) []byte {
	if !token {
		return Seq(OID(OIDData))
	}
	return Seq(OID(oidTSTInfo), tlv(0xa0, Octets(content)))
}

func signerInfo(message []byte, o Options, token bool, s Signer) []byte {
	var sid []byte
	if o.SKI {
		sid = tlv(0x80, s.Cert.SubjectKeyId)
	} else {
		sid = Seq(s.Cert.RawIssuer, mustMarshal(s.Cert.SerialNumber))
	}
	contentType := OIDData
	if token {
		contentType = oidTSTInfo
	}
	md := message
	if o.Message != nil {
		md = o.Message
	}
	essHash := sha256.Sum256(s.Cert.Raw)
	attrs := [][]byte{
		attr(oidContent, OID(contentType)),
		attr(oidDigest, Octets(sum(o.Hash, md))),
	}
	if token {
		h := sha1.Sum(s.Cert.Raw)
		attrs = append(attrs, attr(oidSigCertV1, Seq(Seq(Seq(Octets(h[:]))))))
	} else {
		attrs = append(attrs, attr(OIDSigCertV2, Seq(Seq(Seq(Octets(essHash[:]))))))
	}
	if o.SigCertV1 {
		h := sha1.Sum(s.Cert.Raw)
		attrs = append(attrs, attr(oidSigCertV1, Seq(Seq(Seq(Octets(h[:]))))))
	}
	attrs = append(attrs, o.ExtraAttrs...)
	if o.Mutate != nil {
		attrs = o.Mutate(attrs)
	}
	signed := Set(0xa0, attrs...)
	forSig := append([]byte{0x31}, signed[1:]...)
	digest := sum(o.Hash, forSig)

	var sigAlg, sig []byte
	switch k := s.Key.(type) {
	case *rsa.PrivateKey:
		if o.PSS {
			params := Seq(tlv(0xa0, hashAlg(o.Hash)), tlv(0xa1, Seq(OID(oidMGF1), hashAlg(o.Hash))), tlv(0xa2, Int(int64(o.Hash.Size()))))
			if o.PSSTrailer {
				params = Seq(tlv(0xa0, hashAlg(o.Hash)), tlv(0xa1, Seq(OID(oidMGF1), hashAlg(o.Hash))), tlv(0xa2, Int(int64(o.Hash.Size()))), tlv(0xa3, Int(1)))
			}
			sigAlg = Seq(OID(oidPSS), params)
			sig, _ = rsa.SignPSS(rand.Reader, k, o.Hash, digest, &rsa.PSSOptions{SaltLength: o.Hash.Size(), Hash: o.Hash})
		} else {
			sigAlg = Seq(OID(oidRSA), []byte{5, 0})
			sig, _ = rsa.SignPKCS1v15(rand.Reader, k, o.Hash, digest)
		}
	case *ecdsa.PrivateKey:
		oid := oidECDSA256
		switch o.Hash {
		case crypto.SHA384:
			oid = oidECDSA384
		case crypto.SHA512:
			oid = oidECDSA512
		}
		sigAlg = Seq(OID(oid))
		sig, _ = ecdsa.SignASN1(rand.Reader, k, digest)
	}
	version := int64(1)
	if o.SKI {
		version = 3
	}
	f := [][]byte{Int(version), sid, hashAlg(o.Hash), signed, sigAlg, Octets(sig)}
	var unsigned [][]byte
	if o.Junk > 0 {
		unsigned = append(unsigned, attr(asn1.ObjectIdentifier{1, 2, 3, 4, 5}, Octets(make([]byte, o.Junk))))
	}
	if o.Token != nil {
		t := o.Token(sig)
		v := attr(OIDTimeStamp, t)
		if o.Token2 {
			v = Seq(OID(OIDTimeStamp), Set(0x31, t, Seq(OID(OIDData))))
		}
		unsigned = append(unsigned, v)
	}
	if len(unsigned) > 0 {
		f = append(f, Set(0xa1, unsigned...))
	}
	return Seq(f...)
}

func mustMarshal(v any) []byte {
	b, err := asn1.Marshal(v)
	if err != nil {
		panic(err)
	}
	return b
}

// TokenOptions changes what Token writes.
type TokenOptions struct {
	Hash     crypto.Hash   // the hash of the messageImprint, SHA-256 by default
	Accuracy time.Duration // whole seconds; zero for none
	Version  int           // the version of the TSTInfo, 1 by default
	// Imprint, when not nil, is written as the hashed message instead of
	// the hash of the subject.
	Imprint []byte
}

// Token returns the RFC 3161 time-stamp token that tsa issues over subject
// at genTime.
func Token(subject []byte, genTime time.Time, o TokenOptions, tsa Signer) []byte {
	if o.Hash == 0 {
		o.Hash = crypto.SHA256
	}
	if o.Version == 0 {
		o.Version = 1
	}
	imprint := sum(o.Hash, subject)
	if o.Imprint != nil {
		imprint = o.Imprint
	}
	gt, err := asn1.MarshalWithParams(genTime.UTC(), "generalized")
	if err != nil {
		panic(err)
	}
	info := []byte{}
	info = append(info, Int(int64(o.Version))...)
	info = append(info, OID(asn1.ObjectIdentifier{1, 2, 3, 4})...)
	info = append(info, Seq(hashAlg(o.Hash), Octets(imprint))...)
	info = append(info, Int(42)...)
	info = append(info, gt...)
	if o.Accuracy != 0 {
		info = append(info, Seq(Int(int64(o.Accuracy/time.Second)))...)
	}
	tst := Seq(info)
	return build(tst, Options{Hash: crypto.SHA256}, true, []Signer{tsa})
}