yao/openapi/oauth/signing.go

package oauth

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/yaoapp/yao/openapi/oauth/types"
	"github.com/yaoapp/yao/share"
)

// SigningCertificates holds the loaded signing certificates and keys
type SigningCertificates struct {
	// Primary signing certificate and private key
	SigningCert    *x509.Certificate `json:"-"`
	SigningKey     interface{}       `json:"-"` // *rsa.PrivateKey, *ecdsa.PrivateKey, etc.
	SigningKeyPair *tls.Certificate  `json:"-"`

	// Verification certificates for token validation
	VerificationCerts []*x509.Certificate `json:"-"`

	// mTLS CA certificate for client validation
	MTLSClientCACert *x509.Certificate `json:"-"`

	// Signing algorithm
	Algorithm string `json:"algorithm"`

	// Certificate paths for reference
	SigningCertPath string `json:"signing_cert_path"`
	SigningKeyPath  string `json:"signing_key_path"`

	// Auto-generated flag
	IsAutoGenerated bool `json:"is_auto_generated"`
}

// LoadSigningCertificates loads or generates signing certificates based on configuration
func LoadSigningCertificates(config *types.SigningConfig) (*SigningCertificates, error) {
	certs := &SigningCertificates{
		Algorithm: config.SigningAlgorithm,
	}

	// Check if certificate files exist
	signingCertExists := fileExists(config.SigningCertPath)
	signingKeyExists := fileExists(config.SigningKeyPath)

	// If both files exist, try to load them
	if signingCertExists && signingKeyExists {
		err := loadExistingCertificates(certs, config)
		if err != nil {
			// If loading fails, log warning and generate new certificates
			fmt.Printf("Warning: Failed to load existing certificates (%v), generating new temporary certificates\n", err)
			return generateTemporaryCertificates(config)
		}
		return certs, nil
	}

	// If certificates don't exist, generate temporary ones
	return generateTemporaryCertificates(config)
}

// loadExistingCertificates loads certificates from the configured paths
func loadExistingCertificates(certs *SigningCertificates, config *types.SigningConfig) error {
	// Load signing certificate
	certPEM, err := os.ReadFile(config.SigningCertPath)
	if err != nil {
		return fmt.Errorf("failed to read signing certificate: %w", err)
	}

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

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

	// Load signing key
	keyPEM, err := os.ReadFile(config.SigningKeyPath)
	if err != nil {
		return fmt.Errorf("failed to read signing key: %w", err)
	}

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

	var signingKey interface{}
	if config.SigningKeyPassword != "" {
		// Decrypt encrypted key
		keyBytes, err := x509.DecryptPEMBlock(keyBlock, []byte(config.SigningKeyPassword))
		if err != nil {
			return fmt.Errorf("failed to decrypt signing key: %w", err)
		}
		signingKey, err = parsePrivateKey(keyBytes)
		if err != nil {
			return fmt.Errorf("failed to parse decrypted signing key: %w", err)
		}
	} else {
		// Parse unencrypted key
		var err error
		signingKey, err = parsePrivateKey(keyBlock.Bytes)
		if err != nil {
			return fmt.Errorf("failed to parse signing key: %w", err)
		}
	}

	// Create TLS certificate pair
	keyPair, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return fmt.Errorf("failed to create key pair: %w", err)
	}

	certs.SigningCert = signingCert
	certs.SigningKey = signingKey
	certs.SigningKeyPair = &keyPair
	certs.SigningCertPath = config.SigningCertPath
	certs.SigningKeyPath = config.SigningKeyPath
	certs.IsAutoGenerated = false

	// Load verification certificates if configured
	if len(config.VerificationCerts) > 0 {
		verificationCerts, err := loadVerificationCertificates(config.VerificationCerts)
		if err != nil {
			return fmt.Errorf("failed to load verification certificates: %w", err)
		}
		certs.VerificationCerts = verificationCerts
	}

	// Load mTLS CA certificate if configured
	if config.MTLSClientCACertPath != "" {
		mtlsCACert, err := loadCertificateFromFile(config.MTLSClientCACertPath)
		if err != nil {
			return fmt.Errorf("failed to load mTLS CA certificate: %w", err)
		}
		certs.MTLSClientCACert = mtlsCACert
	}

	return nil
}

// generateTemporaryCertificates generates temporary self-signed certificates
func generateTemporaryCertificates(config *types.SigningConfig) (*SigningCertificates, error) {
	// Generate RSA private key
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}

	// Create certificate template
	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization:  []string{share.App.Name},
			Country:       []string{"US"},
			Province:      []string{""},
			Locality:      []string{""},
			StreetAddress: []string{""},
			PostalCode:    []string{""},
			CommonName:    fmt.Sprintf("%s OAuth Signing Certificate", share.App.Name),
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(365 * 24 * time.Hour), // Valid for 1 year
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	// Generate certificate
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate: %w", err)
	}

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

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

	keyDER, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal private key: %w", err)
	}

	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: keyDER,
	})

	// Create TLS certificate pair
	keyPair, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("failed to create key pair: %w", err)
	}

	// Determine system directory for storing temporary certificates
	systemDir := getSystemCertificateDirectory()

	// Create directory if it doesn't exist
	if err := os.MkdirAll(systemDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create system certificate directory: %w", err)
	}

	// Generate unique filenames
	timestamp := time.Now().Format("20060102150405")
	certPath := filepath.Join(systemDir, fmt.Sprintf("oauth_signing_cert_%s.pem", timestamp))
	keyPath := filepath.Join(systemDir, fmt.Sprintf("oauth_signing_key_%s.pem", timestamp))

	// Save certificate and key to system directory
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return nil, fmt.Errorf("failed to write certificate file: %w", err)
	}

	if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
		return nil, fmt.Errorf("failed to write key file: %w", err)
	}

	fmt.Printf("Generated temporary OAuth signing certificate at: %s\n", certPath)
	fmt.Printf("Generated temporary OAuth signing key at: %s\n", keyPath)

	return &SigningCertificates{
		SigningCert:     cert,
		SigningKey:      privateKey,
		SigningKeyPair:  &keyPair,
		Algorithm:       config.SigningAlgorithm,
		SigningCertPath: certPath,
		SigningKeyPath:  keyPath,
		IsAutoGenerated: true,
	}, nil
}

// loadVerificationCertificates loads additional verification certificates
func loadVerificationCertificates(certPaths []string) ([]*x509.Certificate, error) {
	var certs []*x509.Certificate

	for _, certPath := range certPaths {
		cert, err := loadCertificateFromFile(certPath)
		if err != nil {
			return nil, fmt.Errorf("failed to load verification certificate %s: %w", certPath, err)
		}
		certs = append(certs, cert)
	}

	return certs, nil
}

// loadCertificateFromFile loads a certificate from a PEM file
func loadCertificateFromFile(certPath string) (*x509.Certificate, error) {
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read certificate file: %w", err)
	}

	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)
	}

	return cert, nil
}

// parsePrivateKey parses a private key from DER bytes
func parsePrivateKey(der []byte) (interface{}, error) {
	// Try PKCS#8 first
	if key, err := x509.ParsePKCS8PrivateKey(der); err == nil {
		return key, nil
	}

	// Try PKCS#1 RSA
	if key, err := x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}

	// Try EC private key
	if key, err := x509.ParseECPrivateKey(der); err == nil {
		return key, nil
	}

	return nil, fmt.Errorf("unable to parse private key")
}

// fileExists checks if a file exists
func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return !os.IsNotExist(err)
}

// getSystemCertificateDirectory returns the appropriate system directory for storing certificates
func getSystemCertificateDirectory() string {
	// Use different directories based on the operating system
	homeDir, err := os.UserHomeDir()
	if err != nil {
		// Fallback to temporary directory
		return filepath.Join(os.TempDir(), "yao-oauth-certs")
	}

	// Create a hidden directory in user's home
	return filepath.Join(homeDir, ".yao", "oauth", "certs")
}

// ValidateCertificate validates a certificate for OAuth signing
func (c *SigningCertificates) ValidateCertificate() error {
	if c.SigningCert == nil {
		return fmt.Errorf("signing certificate is nil")
	}

	// Check if certificate is expired
	now := time.Now()
	if now.Before(c.SigningCert.NotBefore) {
		return fmt.Errorf("signing certificate is not yet valid")
	}

	if now.After(c.SigningCert.NotAfter) {
		return fmt.Errorf("signing certificate has expired")
	}

	// Check if certificate has appropriate key usage
	if c.SigningCert.KeyUsage&x509.KeyUsageDigitalSignature == 0 {
		return fmt.Errorf("signing certificate does not have digital signature key usage")
	}

	return nil
}

// GetPublicKey returns the public key from the signing certificate
func (c *SigningCertificates) GetPublicKey() interface{} {
	if c.SigningCert == nil {
		return nil
	}
	return c.SigningCert.PublicKey
}

// GetKeyID returns a key identifier for the signing certificate
func (c *SigningCertificates) GetKeyID() string {
	if c.SigningCert == nil {
		return ""
	}

	// Use the certificate's serial number as key ID
	return c.SigningCert.SerialNumber.String()
}

// CleanupTemporaryCertificates removes auto-generated temporary certificates
func (c *SigningCertificates) CleanupTemporaryCertificates() error {
	if !c.IsAutoGenerated {
		return nil // Don't delete user-provided certificates
	}

	var errs []error

	if c.SigningCertPath != "" && fileExists(c.SigningCertPath) {
		if err := os.Remove(c.SigningCertPath); err != nil {
			errs = append(errs, fmt.Errorf("failed to remove certificate file %s: %w", c.SigningCertPath, err))
		}
	}

	if c.SigningKeyPath != "" && fileExists(c.SigningKeyPath) {
		if err := os.Remove(c.SigningKeyPath); err != nil {
			errs = append(errs, fmt.Errorf("failed to remove key file %s: %w", c.SigningKeyPath, err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("cleanup errors: %v", errs)
	}

	return nil
}

// Service signing certificate methods

// GetSigningCertificates returns the signing certificates for the service
func (s *Service) GetSigningCertificates() *SigningCertificates {
	return s.signingCerts
}

// GetSigningKey returns the signing private key
func (s *Service) GetSigningKey() interface{} {
	if s.signingCerts == nil {
		return nil
	}
	return s.signingCerts.SigningKey
}

// GetSigningCertificate returns the signing certificate
func (s *Service) GetSigningCertificate() interface{} {
	if s.signingCerts == nil {
		return nil
	}
	return s.signingCerts.SigningCert
}

// GetSigningAlgorithm returns the signing algorithm
func (s *Service) GetSigningAlgorithm() string {
	if s.signingCerts == nil {
		return "RS256" // default
	}
	return s.signingCerts.Algorithm
}

// GetKeyID returns the key identifier for JWT token signing
func (s *Service) GetKeyID() string {
	if s.signingCerts == nil {
		return ""
	}
	return s.signingCerts.GetKeyID()
}

// SignToken signs a token based on the configured format (jwt or opaque)
func (s *Service) SignToken(tokenType, clientID, scope, subject string, expiresIn int) (string, error) {
	switch s.config.Token.AccessTokenFormat {
	case "jwt":
		return s.signJWTToken(tokenType, clientID, scope, subject, expiresIn)
	case "opaque":
		return s.signOpaqueToken(tokenType, clientID, scope, subject)
	default:
		// Default to JWT if format is not specified or unknown
		return s.signJWTToken(tokenType, clientID, scope, subject, expiresIn)
	}
}

// VerifyToken verifies a token based on its format and returns token claims
func (s *Service) VerifyToken(token string) (*types.TokenClaims, error) {
	// First try to verify as JWT (JWT tokens contain dots)
	if strings.Contains(token, ".") {
		return s.verifyJWTToken(token)
	}

	// Otherwise, verify as opaque token
	return s.verifyOpaqueToken(token)
}

// signJWTToken signs a JWT token using the configured signing algorithm
func (s *Service) signJWTToken(tokenType, clientID, scope, subject string, expiresIn int) (string, error) {
	if s.signingCerts == nil || s.signingCerts.SigningKey == nil {
		return "", fmt.Errorf("signing certificates not initialized")
	}

	now := time.Now()
	claims := &types.JWTClaims{
		StandardClaims: jwt.StandardClaims{
			Issuer:    s.config.IssuerURL,
			Subject:   subject,
			Audience:  clientID,
			ExpiresAt: now.Add(time.Duration(expiresIn) * time.Second).Unix(),
			NotBefore: now.Unix(),
			IssuedAt:  now.Unix(),
			Id:        generateJTI(),
		},
		ClientID:  clientID,
		Scope:     scope,
		TokenType: tokenType,
	}

	// Create token with claims
	token := jwt.NewWithClaims(getSigningMethod(s.config.Token.AccessTokenSigningAlg), claims)

	// Set key ID in header
	token.Header["kid"] = s.GetKeyID()

	// Sign token with private key
	return token.SignedString(s.signingCerts.SigningKey)
}

// verifyJWTToken verifies a JWT token and returns its claims
func (s *Service) verifyJWTToken(tokenString string) (*types.TokenClaims, error) {
	if s.signingCerts == nil || s.signingCerts.SigningCert == nil {
		return nil, fmt.Errorf("signing certificates not initialized")
	}

	// Parse token with claims
	token, err := jwt.ParseWithClaims(tokenString, &types.JWTClaims{}, func(token *jwt.Token) (interface{}, error) {
		// Validate signing method
		expectedMethod := getSigningMethod(s.config.Token.AccessTokenSigningAlg)
		if token.Method != expectedMethod {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}

		// Return public key for verification
		return s.signingCerts.GetPublicKey(), nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to parse JWT token: %w", err)
	}

	if !token.Valid {
		return nil, fmt.Errorf("invalid JWT token")
	}

	// Extract claims
	jwtClaims, ok := token.Claims.(*types.JWTClaims)
	if !ok {
		return nil, fmt.Errorf("invalid JWT claims type")
	}

	// Convert to TokenClaims
	tokenClaims := &types.TokenClaims{
		Subject:   jwtClaims.Subject,
		ClientID:  jwtClaims.ClientID,
		Scope:     jwtClaims.Scope,
		TokenType: jwtClaims.TokenType,
		ExpiresAt: time.Unix(jwtClaims.ExpiresAt, 0),
		IssuedAt:  time.Unix(jwtClaims.IssuedAt, 0),
		Issuer:    jwtClaims.Issuer,
		Audience:  []string{jwtClaims.Audience},
		JTI:       jwtClaims.Id,
	}

	return tokenClaims, nil
}

// signOpaqueToken signs an opaque token using HMAC or RSA signature
func (s *Service) signOpaqueToken(tokenType, clientID, scope, subject string) (string, error) {
	// Generate base opaque token
	baseToken, err := s.generateOpaqueTokenBase(tokenType, clientID)
	if err != nil {
		return "", fmt.Errorf("failed to generate base opaque token: %w", err)
	}

	// Create token metadata for signature
	tokenData := fmt.Sprintf("%s.%s.%s.%s.%d", baseToken, clientID, scope, subject, time.Now().Unix())

	// Sign the token data
	signature, err := s.signData([]byte(tokenData))
	if err != nil {
		return "", fmt.Errorf("failed to sign opaque token: %w", err)
	}

	// Combine base token with signature
	signedToken := fmt.Sprintf("%s.%s", baseToken, base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(signature))

	return signedToken, nil
}

// verifyOpaqueToken verifies an opaque token signature and returns token claims
func (s *Service) verifyOpaqueToken(token string) (*types.TokenClaims, error) {
	parts := strings.Split(token, ".")
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid opaque token format")
	}

	baseToken := parts[0]
	signaturePart := parts[len(parts)-1]

	// Decode signature
	signature, err := base64.URLEncoding.WithPadding(base64.NoPadding).DecodeString(signaturePart)
	if err != nil {
		return nil, fmt.Errorf("failed to decode token signature: %w", err)
	}

	// Extract token information from store
	tokenInfo, err := s.getAccessTokenData(token)
	if err != nil {
		return nil, fmt.Errorf("token not found or invalid: %w", err)
	}

	// Reconstruct token data for verification
	clientID := tokenInfo["client_id"].(string)
	scope := ""
	if scopeVal, ok := tokenInfo["scope"].(string); ok {
		scope = scopeVal
	}
	subject := ""
	if subjectVal, ok := tokenInfo["subject"].(string); ok {
		subject = subjectVal
	}
	issuedAt := tokenInfo["issued_at"].(int64)

	tokenData := fmt.Sprintf("%s.%s.%s.%s.%d", baseToken, clientID, scope, subject, issuedAt)

	// Verify signature
	if err := s.verifySignature([]byte(tokenData), signature); err != nil {
		return nil, fmt.Errorf("invalid token signature: %w", err)
	}

	// Build token claims
	tokenClaims := &types.TokenClaims{
		Subject:   subject,
		ClientID:  clientID,
		Scope:     scope,
		TokenType: "access_token",
		IssuedAt:  time.Unix(issuedAt, 0),
		Issuer:    s.config.IssuerURL,
	}

	if expiresAt, ok := tokenInfo["expires_at"].(int64); ok {
		tokenClaims.ExpiresAt = time.Unix(expiresAt, 0)
	}

	return tokenClaims, nil
}

// signData signs data using the configured signing key
func (s *Service) signData(data []byte) ([]byte, error) {
	if s.signingCerts == nil || s.signingCerts.SigningKey == nil {
		return nil, fmt.Errorf("signing key not available")
	}

	switch key := s.signingCerts.SigningKey.(type) {
	case *rsa.PrivateKey:
		// Use RSA-PSS for signing
		hash := sha256.Sum256(data)
		signature, err := rsa.SignPSS(rand.Reader, key, crypto.SHA256, hash[:], nil)
		if err != nil {
			return nil, fmt.Errorf("failed to sign with RSA key: %w", err)
		}
		return signature, nil
	default:
		return nil, fmt.Errorf("unsupported signing key type: %T", key)
	}
}

// verifySignature verifies a signature using the configured public key
func (s *Service) verifySignature(data []byte, signature []byte) error {
	if s.signingCerts == nil || s.signingCerts.SigningCert == nil {
		return fmt.Errorf("signing certificate not available")
	}

	switch pubKey := s.signingCerts.GetPublicKey().(type) {
	case *rsa.PublicKey:
		// Use RSA-PSS for verification
		hash := sha256.Sum256(data)
		err := rsa.VerifyPSS(pubKey, crypto.SHA256, hash[:], signature, nil)
		if err != nil {
			return fmt.Errorf("failed to verify RSA signature: %w", err)
		}
		return nil
	default:
		return fmt.Errorf("unsupported public key type: %T", pubKey)
	}
}

// generateOpaqueTokenBase generates the base part of an opaque token
func (s *Service) generateOpaqueTokenBase(tokenType, clientID string) (string, error) {
	// Generate random bytes for token
	randomBytes := make([]byte, 32)
	if _, err := rand.Read(randomBytes); err != nil {
		return "", fmt.Errorf("failed to generate random bytes: %w", err)
	}

	// Create base token with type, client ID, timestamp, and random component
	randomPart := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(randomBytes)
	timestamp := time.Now().Format("20060102150405")

	return fmt.Sprintf("%s_%s_%s_%s", tokenType, clientID, timestamp, randomPart), nil
}

// getSigningMethod returns the JWT signing method for the given algorithm
func getSigningMethod(algorithm string) jwt.SigningMethod {
	switch algorithm {
	case "RS256":
		return jwt.SigningMethodRS256
	case "RS384":
		return jwt.SigningMethodRS384
	case "RS512":
		return jwt.SigningMethodRS512
	case "PS256":
		return jwt.SigningMethodPS256
	case "PS384":
		return jwt.SigningMethodPS384
	case "PS512":
		return jwt.SigningMethodPS512
	default:
		return jwt.SigningMethodRS256 // Default
	}
}

// generateJTI generates a unique JWT ID
func generateJTI() string {
	randomBytes := make([]byte, 16)
	rand.Read(randomBytes)
	return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(randomBytes)
}