tyto/backend/internal/pki/cert.go

package pki

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"time"
)

// CertificateBundle contains a certificate and its private key.
type CertificateBundle struct {
	Certificate    *x509.Certificate
	PrivateKey     *rsa.PrivateKey
	CertificatePEM []byte
	PrivateKeyPEM  []byte
}

// ServerCertConfig contains options for generating a server certificate.
type ServerCertConfig struct {
	CommonName   string
	Organization string
	DNSNames     []string
	IPAddresses  []net.IP
	Validity     time.Duration
	KeySize      int
}

// AgentCertConfig contains options for generating an agent certificate.
type AgentCertConfig struct {
	AgentID      string // Used as CommonName
	Organization string
	Validity     time.Duration
	KeySize      int
}

// GenerateServerCert creates a new server certificate signed by the CA.
func (ca *CA) GenerateServerCert(cfg ServerCertConfig) (*CertificateBundle, error) {
	if cfg.CommonName == "" {
		return nil, fmt.Errorf("common name is required")
	}
	if cfg.KeySize == 0 {
		cfg.KeySize = DefaultKeySize
	}
	if cfg.Validity == 0 {
		cfg.Validity = DefaultCertValidity
	}

	// Generate private key
	key, err := rsa.GenerateKey(rand.Reader, cfg.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate server key: %w", err)
	}

	// Generate serial number
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   cfg.CommonName,
			Organization: []string{cfg.Organization},
		},
		NotBefore:             now,
		NotAfter:              now.Add(cfg.Validity),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		DNSNames:              cfg.DNSNames,
		IPAddresses:           cfg.IPAddresses,
	}

	// Sign with CA
	certDER, err := x509.CreateCertificate(rand.Reader, template, ca.cert, &key.PublicKey, ca.key)
	if err != nil {
		return nil, fmt.Errorf("failed to create server certificate: %w", err)
	}

	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse server certificate: %w", err)
	}

	// Encode to PEM
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	})

	bundle := &CertificateBundle{
		Certificate:    cert,
		PrivateKey:     key,
		CertificatePEM: certPEM,
		PrivateKeyPEM:  keyPEM,
	}

	// Store certificate info
	if ca.store != nil {
		info := CertInfoFromX509(cert, CertTypeServer)
		ca.store.AddCertificate(info)
	}

	return bundle, nil
}

// GenerateAgentCert creates a new agent certificate signed by the CA.
// The AgentID is used as the CommonName and is extracted during mTLS verification.
func (ca *CA) GenerateAgentCert(cfg AgentCertConfig) (*CertificateBundle, error) {
	if cfg.AgentID == "" {
		return nil, fmt.Errorf("agent ID is required")
	}
	if cfg.KeySize == 0 {
		cfg.KeySize = DefaultKeySize
	}
	if cfg.Validity == 0 {
		cfg.Validity = DefaultCertValidity
	}

	// Generate private key
	key, err := rsa.GenerateKey(rand.Reader, cfg.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate agent key: %w", err)
	}

	// Generate serial number
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   cfg.AgentID,
			Organization: []string{cfg.Organization},
		},
		NotBefore:             now,
		NotAfter:              now.Add(cfg.Validity),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	// Sign with CA
	certDER, err := x509.CreateCertificate(rand.Reader, template, ca.cert, &key.PublicKey, ca.key)
	if err != nil {
		return nil, fmt.Errorf("failed to create agent certificate: %w", err)
	}

	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse agent certificate: %w", err)
	}

	// Encode to PEM
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	})

	bundle := &CertificateBundle{
		Certificate:    cert,
		PrivateKey:     key,
		CertificatePEM: certPEM,
		PrivateKeyPEM:  keyPEM,
	}

	// Store certificate info
	if ca.store != nil {
		info := CertInfoFromX509(cert, CertTypeAgent)
		ca.store.AddCertificate(info)
	}

	return bundle, nil
}

// Info returns metadata about this certificate bundle.
func (b *CertificateBundle) Info(certType CertType) CertInfo {
	return CertInfoFromX509(b.Certificate, certType)
}

// SaveToFiles writes the certificate and key to separate files.
func (b *CertificateBundle) SaveToFiles(certPath, keyPath string) error {
	if err := writeFile(certPath, b.CertificatePEM, 0644); err != nil {
		return fmt.Errorf("failed to write certificate: %w", err)
	}

	if err := writeFile(keyPath, b.PrivateKeyPEM, 0600); err != nil {
		return fmt.Errorf("failed to write private key: %w", err)
	}

	return nil
}

// LoadCertificateBundle loads a certificate and key from PEM files.
func LoadCertificateBundle(certPath, keyPath string) (*CertificateBundle, error) {
	certPEM, err := readFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read certificate: %w", err)
	}

	keyPEM, err := readFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read private key: %w", err)
	}

	// Parse certificate
	certBlock, _ := pem.Decode(certPEM)
	if certBlock == nil {
		return nil, fmt.Errorf("failed to decode certificate PEM")
	}

	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse certificate: %w", err)
	}

	// Parse private key
	keyBlock, _ := pem.Decode(keyPEM)
	if keyBlock == nil {
		return nil, fmt.Errorf("failed to decode private key PEM")
	}

	var key *rsa.PrivateKey
	switch keyBlock.Type {
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	case "PRIVATE KEY":
		parsedKey, parseErr := x509.ParsePKCS8PrivateKey(keyBlock.Bytes)
		if parseErr != nil {
			return nil, fmt.Errorf("failed to parse private key: %w", parseErr)
		}
		var ok bool
		key, ok = parsedKey.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("private key is not RSA")
		}
	default:
		return nil, fmt.Errorf("unsupported key type: %s", keyBlock.Type)
	}
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	return &CertificateBundle{
		Certificate:    cert,
		PrivateKey:     key,
		CertificatePEM: certPEM,
		PrivateKeyPEM:  keyPEM,
	}, nil
}