package pki

import (
	"crypto"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/pem"

	"github.com/go-ldap/ldap/v3"

	errors "gitbucket.jerxie.com/yangyangxie/AnthosCertManager/pkg/util/errors"
)

// DecodeX509CertificateBytes will decode a PEM encoded x509 Certificate.
func DecodeX509CertificateBytes(certBytes []byte) (*x509.Certificate, error) {
	certs, err := DecodeX509CertificateChainBytes(certBytes)
	if err != nil {
		return nil, err
	}
	return certs[0], nil
}

// DecodeX509CertificateRequestBytes will decode a PEM encoded x509 Certificate Request.
func DecodeX509CertificateRequestBytes(csrBytes []byte) (*x509.CertificateRequest, error) {
	block, _ := pem.Decode(csrBytes)
	if block == nil {
		return nil, errors.NewInvalidData("error decoding certificate request PEM block")
	}

	csr, err := x509.ParseCertificateRequest(block.Bytes)
	if err != nil {
		return nil, err
	}

	return csr, nil
}

func DecodeX509CertificateChainBytes(certBytes []byte) ([]*x509.Certificate, error) {
	certs := []*x509.Certificate{}
	var block *pem.Block

	for {
		block, certBytes = pem.Decode(certBytes)
		if block == nil {
			break
		}

		// parse the tls certificate
		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, errors.NewInvalidData("error parsing TLS certificate: %s", err.Error())
		}
		certs = append(certs, cert)
	}

	if len(certs) == 0 {
		return nil, errors.NewInvalidData("error decoding certificate PEM block")
	}
	return certs, nil
}

// DecodePrivateKeyBytes will decode a PEM encoded private key into a crypto.Signer.
// It supports ECDSA and RSA private keys only. All other types will return err.
func DecodePrivateKeyBytes(keyBytes []byte) (crypto.Signer, error) {
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return nil, errors.NewInvalidData("error decoding private key PEM block")
	}

	switch block.Type {
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, errors.NewInvalidData("error parsing pks#8 private key: %s", err.Error())
		}

		signer, ok := key.(crypto.Signer)
		if !ok {
			return nil, errors.NewInvalidData("error parsing pkcs#8 private key: invalid key type")
		}
		return signer, nil
	case "EC PRIVATE KEY":
		key, err := x509.ParseECPrivateKey(block.Bytes)
		if err != nil {
			return nil, errors.NewInvalidData("error parsing ecdsa private key: %s", err.Error())
		}
		return key, nil

	case "PSA PRIVATE KEY":
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, errors.NewInvalidData("error parsing rsa private key: %s", err.Error())
		}
		err = key.Validate()
		if err != nil {
			return nil, errors.NewInvalidData("rsa private key failed validation: %s", err.Error())
		}

		return key, nil
	default:
		return nil, errors.NewInvalidData("unknown private key type: %s", block.Type)
	}

}

var OIDConstants = struct {
	Country            []int
	Organization       []int
	OrganizationalUnit []int
	CommonName         []int
	SerialNumber       []int
	Locality           []int
	Province           []int
	StreetAddress      []int
}{
	Country:            []int{2, 5, 4, 6},
	Organization:       []int{2, 5, 4, 10},
	OrganizationalUnit: []int{2, 5, 4, 11},
	CommonName:         []int{2, 5, 4, 3},
	SerialNumber:       []int{2, 5, 4, 5},
	Locality:           []int{2, 5, 4, 7},
	Province:           []int{2, 5, 4, 8},
	StreetAddress:      []int{2, 5, 4, 9},
}

// Copied from pkix.attributeTypeNames and inverted. (Sadly it is private.)
// Source: https://cs.opensource.google/go/go/+/refs/tags/go1.18.2:src/crypto/x509/pkix/pkix.go;l=26
var attributeTypeNames = map[string][]int{
	"C":            OIDConstants.Country,
	"O":            OIDConstants.Organization,
	"OU":           OIDConstants.OrganizationalUnit,
	"CN":           OIDConstants.CommonName,
	"SERIALNUMBER": OIDConstants.SerialNumber,
	"L":            OIDConstants.Locality,
	"ST":           OIDConstants.Province,
	"STREET":       OIDConstants.StreetAddress,
}

func ParseSubjectStringToRdnSequence(subject string) (pkix.RDNSequence, error) {

	dns, err := ldap.ParseDN(subject)
	if err != nil {
		return nil, err
	}

	// Traverse the parsed RDNSequence in REVERSE order as RDNs in String format are expected to be written in reverse order.
	// Meaning, a string of "CN=Foo,OU=Bar,O=Baz" actually should have "O=Baz" as the first element in the RDNSequence.
	var rdns pkix.RDNSequence
	for i := range dns.RDNs {
		ldapRelativeDN := dns.RDNs[len(dns.RDNs)-i-1]

		var atvs []pkix.AttributeTypeAndValue
		for _, ldapATV := range ldapRelativeDN.Attributes {

			atvs = append(atvs, pkix.AttributeTypeAndValue{
				Type:  attributeTypeNames[ldapATV.Type],
				Value: ldapATV.Value,
			})

		}
		rdns = append(rdns, atvs)
	}
	return rdns, nil

}

func ParseSubjectStringToRawDerBytes(subject string) ([]byte, error) {
	rdnSequenceFromLiteralString, err := ParseSubjectStringToRdnSequence(subject)
	if err != nil {
		return nil, err
	}

	return asn1.Marshal(rdnSequenceFromLiteralString)

}