// Package mtls implements mutual-TLS at the edge: a CA + leaf issuance
// with SPIFFE identities, a hot-rotating server TLS config, identity
// extraction from the verified client cert, and per-route authorization.
// Stdlib crypto only.
package mtls

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"net"
	"net/url"
	"time"
)

// CA is a tiny certificate authority for issuing short-lived workload
// certs (SVIDs). In production this is SPIRE / Netflix Metatron.
type CA struct {
	Cert    *x509.Certificate
	Key     *ecdsa.PrivateKey
	CertPEM []byte
}

// NewCA creates a self-signed CA.
func NewCA() (*CA, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	tmpl := &x509.Certificate{
		SerialNumber:          serial(),
		Subject:               pkix.Name{CommonName: "Edge CA"},
		NotBefore:             time.Now().Add(-time.Minute),
		NotAfter:              time.Now().Add(24 * time.Hour),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	return &CA{Cert: cert, Key: key, CertPEM: certPEM}, nil
}

// Pool returns an x509 pool trusting this CA (for ClientCAs / RootCAs).
func (ca *CA) Pool() *x509.CertPool {
	p := x509.NewCertPool()
	p.AppendCertsFromPEM(ca.CertPEM)
	return p
}

// Issue creates a leaf certificate carrying a SPIFFE ID in its URI SAN.
// `server` controls ExtKeyUsage (ServerAuth vs ClientAuth) and, for
// servers, adds loopback SANs so a local mTLS demo verifies. `ttl` is
// short on purpose — SVIDs are reissued continuously.
func (ca *CA) Issue(commonName, spiffeID string, server bool, ttl time.Duration) (tls.Certificate, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return tls.Certificate{}, err
	}
	spiffeURL, err := url.Parse(spiffeID)
	if err != nil {
		return tls.Certificate{}, err
	}
	usage := x509.ExtKeyUsageClientAuth
	tmpl := &x509.Certificate{
		SerialNumber: serial(),
		Subject:      pkix.Name{CommonName: commonName},
		NotBefore:    time.Now().Add(-time.Minute),
		NotAfter:     time.Now().Add(ttl),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		URIs:         []*url.URL{spiffeURL}, // SPIFFE ID lives in the URI SAN
	}
	if server {
		usage = x509.ExtKeyUsageServerAuth
		tmpl.DNSNames = []string{"localhost", "gateway.local"}
		tmpl.IPAddresses = []net.IP{net.ParseIP("127.0.0.1"), net.IPv6loopback}
	}
	tmpl.ExtKeyUsage = []x509.ExtKeyUsage{usage}

	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.Cert, &key.PublicKey, ca.Key)
	if err != nil {
		return tls.Certificate{}, err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return tls.Certificate{}, err
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
	pair, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return tls.Certificate{}, err
	}
	pair.Leaf, _ = x509.ParseCertificate(der)
	return pair, nil
}

func serial() *big.Int {
	n, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	return n
}

// EncodePEM re-encodes a tls.Certificate's leaf cert and EC private key
// to PEM (e.g. to write files for curl in the demo).
func EncodePEM(cert tls.Certificate) (certPEM, keyPEM []byte, err error) {
	if len(cert.Certificate) == 0 {
		return nil, nil, errNoCert
	}
	certPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: cert.Certificate[0]})
	key, ok := cert.PrivateKey.(*ecdsa.PrivateKey)
	if !ok {
		return nil, nil, errNoCert
	}
	der, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, nil, err
	}
	keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: der})
	return certPEM, keyPEM, nil
}

var errNoCert = errorString("mtls: certificate has no leaf/key")

type errorString string

func (e errorString) Error() string { return string(e) }