package cms

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"errors"
	"math/big"
	"time"

	"g.activething.com/go/DateKeys/internal/der"
)

// Result is the result of checking the signature of a SignerInfo (spec
// §29.10, "Verificación").
type Result int

const (
	// Valid: the message-digest is the hash of the message and the signature
	// of the signedAttrs verifies with the key of the certificate.
	Valid Result = iota
	// Invalid: one of the two does not hold.
	Invalid
	// NotVerifiable: an algorithm, a key size or a curve outside the table.
	NotVerifiable
)

var (
	oidRSAEncryption = oid("1.2.840.113549.1.1.1")
	oidSHA256RSA     = oid("1.2.840.113549.1.1.11")
	oidSHA384RSA     = oid("1.2.840.113549.1.1.12")
	oidSHA512RSA     = oid("1.2.840.113549.1.1.13")
	oidPSS           = oid("1.2.840.113549.1.1.10")
	oidMGF1          = oid("1.2.840.113549.1.1.8")
	oidECDSA256      = oid("1.2.840.10045.4.3.2")
	oidECDSA384      = oid("1.2.840.10045.4.3.3")
	oidECDSA512      = oid("1.2.840.10045.4.3.4")
	oidECPublicKey   = oid("1.2.840.10045.2.1")
	oidP256          = oid("1.2.840.10045.3.1.7")
	oidP384          = oid("1.3.132.0.34")
	oidP521          = oid("1.3.132.0.35")
)

type scheme int

const (
	schemePKCS1 scheme = iota + 1
	schemePSS
	schemeECDSA
)

// params returns the hash, the signature scheme and the hash that the
// signature algorithm itself names, when it does, of the SignerInfo, and
// whether they are in the table.
func (s *SignerInfo) params() (crypto.Hash, scheme, bool) {
	newHash, ok := s.DigestAlg.hashOf()
	if !ok {
		return 0, 0, false
	}
	var h crypto.Hash
	switch newHash().Size() {
	case 32:
		h = crypto.SHA256
	case 48:
		h = crypto.SHA384
	default:
		h = crypto.SHA512
	}
	o, params := s.SigAlg.OID, s.SigAlg.Params
	nullOrAbsent := params == nil || bytes.Equal(params, []byte{5, 0})
	switch {
	case bytes.Equal(o, oidRSAEncryption):
		return h, schemePKCS1, nullOrAbsent
	case bytes.Equal(o, oidSHA256RSA):
		return h, schemePKCS1, nullOrAbsent && h == crypto.SHA256
	case bytes.Equal(o, oidSHA384RSA):
		return h, schemePKCS1, nullOrAbsent && h == crypto.SHA384
	case bytes.Equal(o, oidSHA512RSA):
		return h, schemePKCS1, nullOrAbsent && h == crypto.SHA512
	case bytes.Equal(o, oidECDSA256):
		return h, schemeECDSA, params == nil && h == crypto.SHA256
	case bytes.Equal(o, oidECDSA384):
		return h, schemeECDSA, params == nil && h == crypto.SHA384
	case bytes.Equal(o, oidECDSA512):
		return h, schemeECDSA, params == nil && h == crypto.SHA512
	case bytes.Equal(o, oidPSS):
		return h, schemePSS, pssParamsOK(params, newHash().Size(), s.DigestAlg)
	}
	return 0, 0, false
}

// pssParamsOK checks RSASSA-PSS-params (RFC 4055): the hash of digestAlgorithm,
// MGF1 with that hash, a salt of its length and trailerField 1.
func pssParamsOK(params []byte, hashLen int, digest algID) bool {
	if params == nil {
		return false
	}
	id, f, err := der.Split(params)
	if err != nil || id != 0x30 {
		return false
	}
	var hashOK, mgfOK, saltOK bool
	var last byte
	for _, e := range f {
		_, in, err := der.Split(e)
		if err != nil || len(in) != 1 || e[0] <= last {
			return false // the fields come in order of tag, each once
		}
		last = e[0]
		switch e[0] {
		case 0xa0:
			a, err := parseAlgID(in[0])
			hashOK = err == nil && bytes.Equal(a.OID, digest.OID) && (a.Params == nil || bytes.Equal(a.Params, []byte{5, 0}))
		case 0xa1:
			a, err := parseAlgID(in[0])
			if err != nil || !bytes.Equal(a.OID, oidMGF1) || a.Params == nil {
				return false
			}
			inner, err := parseAlgID(a.Params)
			mgfOK = err == nil && bytes.Equal(inner.OID, digest.OID) && (inner.Params == nil || bytes.Equal(inner.Params, []byte{5, 0}))
		case 0xa2:
			n, ok := smallInt(in[0])
			saltOK = ok && n == hashLen
		default: // [3] trailerField is 1, its DEFAULT: DER does not write it
			return false
		}
	}
	// hashAlgorithm and maskGenAlgorithm default to SHA-1, and the salt to 20
	// bytes: none of them is in the table, so each must be present.
	return hashOK && mgfOK && saltOK
}

// smallInt reads an INTEGER element of at most 4 bytes that is not negative.
func smallInt(b []byte) (int, bool) {
	if len(b) == 0 || b[0] != 0x02 {
		return 0, false
	}
	c, err := der.Content(b)
	if err != nil || len(c) == 0 || len(c) > 4 || c[0]&0x80 != 0 {
		return 0, false
	}
	n := 0
	for _, x := range c {
		n = n<<8 | int(x)
	}
	return n, true
}

// publicKey returns the key of the certificate when the table has it (spec
// §29.10): a SubjectPublicKeyInfo of rsaEncryption with NULL parameters and
// an RSAPublicKey of exactly a modulus and an exponent, the modulus odd and
// of 2048 to 4096 bits and the exponent odd from 3 to 2^31 - 1; or of
// id-ecPublicKey with the named curve P-256, P-384 or P-521 and the
// uncompressed form of a point of it. Anything else is not usable.
func (c *Cert) publicKey() (any, scheme, bool) {
	_, f, err := der.Split(c.SPKI)
	if err != nil || len(f) != 2 || f[0][0] != 0x30 || f[1][0] != 0x03 {
		return nil, 0, false
	}
	alg, err := parseAlgID(f[0])
	if err != nil {
		return nil, 0, false
	}
	bits, err := der.Content(f[1])
	if err != nil || len(bits) < 2 || bits[0] != 0 {
		return nil, 0, false
	}
	key := bits[1:]
	switch {
	case bytes.Equal(alg.OID, oidRSAEncryption) && bytes.Equal(alg.Params, []byte{5, 0}):
		if der.Check(key) != nil {
			return nil, 0, false
		}
		id, ne, err := der.Split(key)
		if err != nil || id != 0x30 || len(ne) != 2 || ne[0][0] != 0x02 || ne[1][0] != 0x02 {
			return nil, 0, false
		}
		nb, _ := der.Content(ne[0])
		eb, _ := der.Content(ne[1])
		if nb[0]&0x80 != 0 || eb[0]&0x80 != 0 || len(eb) > 4 {
			return nil, 0, false
		}
		n := new(big.Int).SetBytes(nb)
		e := new(big.Int).SetBytes(eb).Int64()
		if b := n.BitLen(); b < 2048 || b > 4096 || n.Bit(0) == 0 || e < 3 || e%2 == 0 || e > 1<<31-1 {
			return nil, 0, false
		}
		return &rsa.PublicKey{N: n, E: int(e)}, schemePKCS1, true
	case bytes.Equal(alg.OID, oidECPublicKey):
		curveOID, ok := oidOf(alg.Params)
		if !ok {
			return nil, 0, false
		}
		var curve elliptic.Curve
		switch {
		case bytes.Equal(curveOID, oidP256):
			curve = elliptic.P256()
		case bytes.Equal(curveOID, oidP384):
			curve = elliptic.P384()
		case bytes.Equal(curveOID, oidP521):
			curve = elliptic.P521()
		default:
			return nil, 0, false
		}
		k, err := ecdsa.ParseUncompressedPublicKey(curve, key)
		if err != nil {
			return nil, 0, false
		}
		return k, schemeECDSA, true
	}
	return nil, 0, false
}

// Check checks the signature of the SignerInfo over message, the bytes that
// the signature is detached from (spec §29.10), in the order of the spec: not
// verifiable for an algorithm, a key or a curve outside the table; invalid
// when the message-digest is not the hash of message, or the signature of the
// signedAttrs does not verify with the key of the certificate, which is the
// case of a key of a scheme other than the one of the algorithm.
func (s *SignerInfo) Check(message []byte) Result {
	h, sch, ok := s.params()
	if !ok {
		return NotVerifiable
	}
	key, ksch, ok := s.Cert.publicKey()
	if !ok {
		return NotVerifiable
	}
	if sum := hashBytes(h, message); !bytes.Equal(sum, s.MessageDigest) {
		return Invalid
	}
	if sch != ksch && !(sch == schemePSS && ksch == schemePKCS1) {
		return Invalid
	}
	// The signature covers the signedAttrs with the tag of a SET.
	attrs := bytes.Clone(s.SignedAttrs)
	attrs[0] = 0x31
	digest := hashBytes(h, attrs)
	var valid bool
	switch k := key.(type) {
	case *rsa.PublicKey:
		if sch == schemePSS {
			valid = rsa.VerifyPSS(k, h, digest, s.Signature, &rsa.PSSOptions{SaltLength: h.Size(), Hash: h}) == nil
		} else {
			valid = rsa.VerifyPKCS1v15(k, h, digest, s.Signature) == nil
		}
	case *ecdsa.PublicKey:
		valid = ecdsa.VerifyASN1(k, digest, s.Signature)
	}
	if !valid {
		return Invalid
	}
	return Valid
}

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

// Token is a time-stamp token of RFC 3161 read with the profile of spec
// §29.11.
type Token struct {
	// GenTime is t, and Accuracy the precision of the token, zero in the
	// fields it does not carry. HasAccuracy reports whether it carries the
	// field at all: without it, the token does not say its precision (spec
	// v0.16, §29.11).
	GenTime     time.Time
	Accuracy    time.Duration
	HasAccuracy bool
	// Policy is the content of the object identifier of its policy.
	Policy []byte
	// ImprintAlg is the hash of the messageImprint, and Imprint the hash.
	ImprintAlg algID
	Imprint    []byte
	// TSA is the certificate of the time-stamping authority.
	TSA *Cert

	data *SignedData
}

// maxAccuracy bounds the seconds of accuracy (spec §29.11): far above any
// real one, and far below what would overflow a Duration.
const maxAccuracy = 1<<31 - 1

// ParseToken reads a time-stamp token. It fails with ErrForm when the form
// breaks the profile, and with ErrAlgorithm when an algorithm is outside the
// table, in that order (spec §29.11): the verdicts S2 and S1.
func ParseToken(b []byte) (*Token, error) {
	sd, err := parse(b, true)
	if err != nil {
		return nil, err
	}
	// The TSTInfo is an OCTET STRING inside the token, so the check of the
	// token did not reach it: it is read here, field by field.
	t, imprintAlg, hash, err := parseTSTInfo(sd.EContent)
	if err != nil {
		return nil, err
	}
	t.TSA, t.data = sd.Signers[0].Cert, sd
	if t.ImprintAlg, err = parseAlgID(imprintAlg); err != nil {
		return nil, err
	}
	t.Imprint = hash
	if _, ok := t.ImprintAlg.hashOf(); !ok {
		return nil, ErrAlgorithm
	}
	if _, _, ok := sd.Signers[0].params(); !ok {
		return nil, ErrAlgorithm
	}
	if _, _, ok := t.TSA.publicKey(); !ok {
		return nil, ErrAlgorithm
	}
	return t, nil
}

// parseTSTInfo reads the TSTInfo of RFC 3161 3.2.1 in DER: the fields in
// order, each once, and nothing after the last. It returns the messageImprint
// algorithm, as the DER of its AlgorithmIdentifier, and the hash.
func parseTSTInfo(b []byte) (*Token, []byte, []byte, error) {
	bad := func(what string) (*Token, []byte, []byte, error) {
		return nil, nil, nil, formErr("the TSTInfo: %s", what)
	}
	if err := der.Check(b); err != nil {
		return bad(err.Error())
	}
	id, f, err := der.Split(b)
	if err != nil || id != 0x30 || len(f) < 5 {
		return bad("not a SEQUENCE of at least five fields")
	}
	if v, ok := smallInt(f[0]); !ok || v != 1 {
		return bad("the version is not 1")
	}
	if f[1][0] != 0x06 || f[3][0] != 0x02 || f[4][0] != 0x18 {
		return bad("policy, serialNumber or genTime")
	}
	_, mi, err := der.Split(f[2])
	if f[2][0] != 0x30 || err != nil || len(mi) != 2 || mi[0][0] != 0x30 || mi[1][0] != 0x04 {
		return bad("the messageImprint")
	}
	hash, err := der.Content(mi[1])
	if err != nil {
		return bad("the messageImprint")
	}
	gen, _, err := der.ParseTime(f[4])
	if err != nil {
		return bad("genTime: " + err.Error())
	}
	policy, err := der.Content(f[1])
	if err != nil {
		return bad("policy")
	}
	t := &Token{GenTime: gen, Policy: policy}
	rest := f[5:]
	if len(rest) > 0 && rest[0][0] == 0x30 {
		if t.Accuracy, err = parseAccuracy(rest[0]); err != nil {
			return bad(err.Error())
		}
		t.HasAccuracy = true
		rest = rest[1:]
	}
	if len(rest) > 0 && rest[0][0] == 0x01 {
		if c, _ := der.Content(rest[0]); len(c) != 1 || c[0] != 0xff {
			return bad("ordering FALSE is its default and DER does not write it")
		}
		rest = rest[1:]
	}
	if len(rest) > 0 && rest[0][0] == 0x02 { // nonce
		rest = rest[1:]
	}
	if len(rest) > 0 && rest[0][0] == 0xa0 { // tsa
		rest = rest[1:]
	}
	if len(rest) > 0 && rest[0][0] == 0xa1 { // extensions
		rest = rest[1:]
	}
	if len(rest) != 0 {
		return bad("a field out of its place, or one that does not exist")
	}
	return t, mi[0], hash, nil
}

// parseAccuracy reads Accuracy: seconds from 0 to 2^31 - 1, and millis and
// micros from 1 to 999, in that order, each optional (RFC 3161 2.4.2, spec
// §29.11), each a minimal INTEGER. A negative number would make a seal after
// the opening date look before it.
func parseAccuracy(b []byte) (time.Duration, error) {
	_, f, err := der.Split(b)
	if err != nil {
		return 0, errors.New("accuracy")
	}
	var total time.Duration
	if len(f) > 0 && f[0][0] == 0x02 {
		secs, ok := smallInt(f[0])
		if !ok || secs > maxAccuracy {
			return 0, errors.New("accuracy seconds outside 0 to 2^31 - 1")
		}
		total += time.Duration(secs) * time.Second
		f = f[1:]
	}
	for _, part := range []struct {
		tag  byte
		unit time.Duration
	}{{0x80, time.Millisecond}, {0x81, time.Microsecond}} {
		if len(f) > 0 && f[0][0] == part.tag {
			c, err := der.Content(f[0])
			if err != nil || len(c) < 1 || len(c) > 2 || c[0]&0x80 != 0 || len(c) == 2 && c[0] == 0 && c[1]&0x80 == 0 {
				return 0, errors.New("accuracy millis or micros that are not a minimal INTEGER")
			}
			n := 0
			for _, x := range c {
				n = n<<8 | int(x)
			}
			if n < 1 || n > 999 {
				return 0, errors.New("accuracy millis or micros outside 1 to 999")
			}
			total += time.Duration(n) * part.unit
			f = f[1:]
		}
	}
	if len(f) != 0 {
		return 0, errors.New("accuracy has a field out of its place")
	}
	return total, nil
}

// ImprintIsSHA256 reports whether the messageImprint uses SHA-256, which a
// seal of seal_type 2 requires (spec §29.11).
func (t *Token) ImprintIsSHA256() bool { return bytes.Equal(t.ImprintAlg.OID, oidSHA256) }

// BTSP reports whether the policy of the token is the best practices
// time-stamp policy of ETSI EN 319 421 (0.4.0.2023.1.1), compared by the
// bytes of its DER, which requires accuracy in every token (spec v0.16,
// §29.11).
func (t *Token) BTSP() bool { return bytes.Equal(t.Policy, oidBTSP) }

// Check verifies the token over subject, the bytes that it seals: the
// message-digest is the hash of the TSTInfo, the signature of the TSA
// verifies, the messageImprint is the hash of subject, of any length, and the
// certificate of the TSA is valid at genTime. It returns false for the
// verdict S3.
func (t *Token) Check(subject []byte) bool {
	s := t.data.Signers[0]
	if s.Check(t.data.EContent) != Valid {
		return false
	}
	newHash, _ := t.ImprintAlg.hashOf()
	h := newHash()
	h.Write(subject)
	return bytes.Equal(h.Sum(nil), t.Imprint) && t.TSA.ValidAt(t.GenTime)
}