config.go

package config

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
	"slices"
	"strings"

	"github.com/spf13/viper"
	"go.uber.org/zap"
)

var (
	driverOptions = []string{"postgres"}
)

type Config struct {
	AppPort                     string
	Environment                 string
	DBDriver                    string
	DBConnectionString          string
	DBDebug                     bool
	JWTSecret                   string
	JWTPrivateKey               *rsa.PrivateKey
	JWTPublicKey                *rsa.PublicKey
	APIAllowedOrigins           []string
	MetricsEnabled              bool
	MetricsPort                 string
	WebBaseURL                  string
	SSO                         *SSOConfig
	Email                       *EmailConfig
	Worker                      *WorkerConfig
	EvidenceDefaultExpiryMonths int    // Default expiration in months for evidence without explicit expiry
	DigestEnabled               bool   // Enable or disable the digest scheduler
	DigestSchedule              string // Cron schedule for digest emails
	Workflow                    *WorkflowConfig
	Risk                        *RiskConfig
}

func NewConfig(logger *zap.SugaredLogger) *Config {
	if !viper.IsSet("db_driver") {
		logger.Fatal(
			"CCF_DB_DRIVER is not set. Please set it in the environment or .env file. Expected values: ",
			strings.Join(driverOptions, ", "),
		)
	}

	environment := strings.ToLower(viper.GetString("environment"))
	if environment == "" {
		environment = "production"
	}

	dbDriver := stripQuotes(strings.ToLower(viper.GetString("db_driver")))

	if !slices.Contains(driverOptions, dbDriver) {
		logger.Fatal(
			"CCF_DB_DRIVER is set to an unsupported value: ",
			viper.GetString("db_driver"),
			". Supported values are: ",
			strings.Join(driverOptions, ", "),
		)
	}

	if !viper.IsSet("db_connection") {
		logger.Fatal("CCF_DB_CONNECTION is not set. Please set it in the environment or .env file.")
	}

	if !viper.IsSet("jwt_secret") {
		logger.Warn("Using 'change-me' as JWT secret. This is insecure and should be changed in production.")
		viper.Set("jwt_secret", "change-me")
	}

	var (
		jwtPrivateKey *rsa.PrivateKey
		jwtPublicKey  *rsa.PublicKey
		err           error
	)
	if !viper.IsSet("jwt_private_key") || !viper.IsSet("jwt_public_key") {
		logger.Warn("No JWT key files have been provided. Generating new keys. Any previously-created JWTs will no longer be valid.")
		jwtPrivateKey, jwtPublicKey, err = GenerateKeyPair(2048)
		if err != nil {
			logger.Fatalw("Failed to generate RSA key pair", "error", err)
		}
	} else {
		jwtPrivateKeyPath := stripQuotes(viper.GetString("jwt_private_key"))
		jwtPublicKeyPath := stripQuotes(viper.GetString("jwt_public_key"))

		jwtPrivateKey, err = loadRSAPrivateKey(jwtPrivateKeyPath)
		if err != nil {
			logger.Fatalw("Failed to load RSA private key", "error", err, "path", jwtPrivateKeyPath)
		}
		jwtPublicKey, err = loadRSAPublicKey(jwtPublicKeyPath)
		if err != nil {
			logger.Fatalw("Failed to load RSA public key", "error", err, "path", jwtPublicKeyPath)
		}
	}

	appPort := viper.GetString("app_port")
	if !strings.HasPrefix(appPort, ":") {
		appPort = ":" + appPort
	}

	allowedOrigins := []string{"http://localhost:3000"} // Default fallback
	if viper.IsSet("api_allowed_origins") {
		originsStr := viper.GetString("api_allowed_origins")
		if originsStr != "" {
			// Split by comma and trim whitespace
			origins := make([]string, 0)
			for _, origin := range strings.Split(originsStr, ",") {
				trimmed := strings.TrimSpace(origin)
				if trimmed != "" {
					origins = append(origins, trimmed)
				}
			}
			if len(origins) > 0 {
				allowedOrigins = origins
			}
		} else {
			logger.Warnw("api_allowed_origins is set but empty. Setting to the default", "origins", allowedOrigins)
		}
	}

	metricsEnabled := viper.GetBool("metrics_enabled")
	metricsPort := viper.GetString("metrics_port")

	webBaseURL := viper.GetString("web_base_url")
	if webBaseURL == "" {
		webBaseURL = "http://localhost:8000" // Default fallback
	}

	ssoConfigPath := viper.GetString("sso_config")
	if ssoConfigPath == "" {
		ssoConfigPath = "sso.yaml"
	}
	ssoConfig, err := LoadSSOConfig(ssoConfigPath)
	if err != nil {
		logger.Warnw("Failed to load OIDC config, OIDC will be disabled", "error", err, "path", ssoConfigPath)
		ssoConfig = &SSOConfig{Enabled: false}
	}

	emailConfigPath := viper.GetString("email_config")
	if emailConfigPath == "" {
		emailConfigPath = "email.yaml"
	}
	emailConfig, err := LoadEmailConfig(emailConfigPath)
	if err != nil {
		logger.Warnw("Failed to load email config, email will be disabled", "error", err, "path", emailConfigPath)
		emailConfig = &EmailConfig{Enabled: false}
	}

	// Evidence default expiry in months (default: 1 month)
	evidenceDefaultExpiryMonths := viper.GetInt("evidence_default_expiry_months")
	if evidenceDefaultExpiryMonths <= 0 {
		evidenceDefaultExpiryMonths = 1
	}

	// Digest configuration
	digestEnabled := viper.GetBool("digest_enabled")
	digestSchedule := viper.GetString("digest_schedule")
	if digestSchedule == "" {
		digestSchedule = "@weekly"
	}

	workflowConfigPath := viper.GetString("workflow_config")
	if workflowConfigPath == "" {
		workflowConfigPath = "workflow.yaml"
	}
	workflowConfig, err := LoadWorkflowConfig(workflowConfigPath)
	if err != nil {
		logger.Warnw("Failed to load workflow config, scheduler will be disabled", "error", err, "path", workflowConfigPath)
		workflowConfig = &WorkflowConfig{SchedulerEnabled: false}
	}

	riskConfigPath := viper.GetString("risk_config")
	if riskConfigPath == "" {
		riskConfigPath = "risk.yaml"
	}
	riskConfig, err := LoadRiskConfig(riskConfigPath)
	if err != nil {
		logger.Warnw("Failed to load risk config, risk jobs will be disabled", "error", err, "path", riskConfigPath)
		riskConfig = DefaultRiskConfig()
	}

	// Worker configuration
	workerConfig := DefaultWorkerConfig()
	if viper.IsSet("worker_enabled") {
		workerConfig.Enabled = viper.GetBool("worker_enabled")
	}
	if viper.IsSet("worker_count") {
		workerConfig.Workers = viper.GetInt("worker_count")
	}
	if viper.IsSet("worker_queue") {
		workerConfig.Queue = viper.GetString("worker_queue")
	}

	return &Config{
		AppPort:                     appPort,
		Environment:                 environment,
		DBDriver:                    dbDriver,
		DBConnectionString:          stripQuotes(viper.GetString("db_connection")),
		DBDebug:                     viper.GetBool("db_debug"),
		JWTSecret:                   stripQuotes(viper.GetString("jwt_secret")),
		JWTPrivateKey:               jwtPrivateKey,
		JWTPublicKey:                jwtPublicKey,
		APIAllowedOrigins:           allowedOrigins,
		MetricsEnabled:              metricsEnabled,
		MetricsPort:                 metricsPort,
		WebBaseURL:                  webBaseURL,
		SSO:                         ssoConfig,
		Email:                       emailConfig,
		Worker:                      workerConfig,
		EvidenceDefaultExpiryMonths: evidenceDefaultExpiryMonths,
		DigestEnabled:               digestEnabled,
		DigestSchedule:              digestSchedule,
		Workflow:                    workflowConfig,
		Risk:                        riskConfig,
	}

}

func stripQuotes(s string) string {
	if len(s) >= 2 {
		if (s[0] == '"' && s[len(s)-1] == '"') || (s[0] == '\'' && s[len(s)-1] == '\'') {
			return s[1 : len(s)-1]
		}
	}
	return s
}

// LoadRSAPrivateKey reads an RSA private key from a PEM file at the given path.
func loadRSAPrivateKey(path string) (*rsa.PrivateKey, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("unable to read private key file %s: %w", path, err)
	}
	block, _ := pem.Decode(data)
	if block == nil || (block.Type != "RSA PRIVATE KEY" && block.Type != "PRIVATE KEY") {
		return nil, fmt.Errorf("failed to decode PEM block containing private key")
	}
	// Try PKCS1
	if block.Type == "RSA PRIVATE KEY" {
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err == nil {
			return key, nil
		}
	}
	// Try PKCS8
	privKeyIfc, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("unable to parse private key: %w", err)
	}
	privKey, ok := privKeyIfc.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("key type is not RSA private")
	}
	return privKey, nil
}

// LoadRSAPublicKey reads an RSA public key from a PEM file at the given path.
func loadRSAPublicKey(path string) (*rsa.PublicKey, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("unable to read public key file %s: %w", path, err)
	}
	block, _ := pem.Decode(data)
	if block == nil || (block.Type != "PUBLIC KEY" && block.Type != "RSA PUBLIC KEY") {
		return nil, fmt.Errorf("failed to decode PEM block containing public key")
	}
	var pubIfc any
	if block.Type == "PUBLIC KEY" {
		pubIfc, err = x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("unable to parse PKIX public key: %w", err)
		}
	} else {
		pubIfc, err = x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("unable to parse PKCS1 public key: %w", err)
		}
	}
	pubKey, ok := pubIfc.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("key type is not RSA public")
	}
	return pubKey, nil
}

func GenerateKeyPair(bitsize int) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	privKey, err := rsa.GenerateKey(rand.Reader, bitsize)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate RSA private key: %w", err)
	}
	err = privKey.Validate()
	if err != nil {
		return nil, nil, fmt.Errorf("generated RSA private key is invalid: %w", err)
	}

	return privKey, &privKey.PublicKey, nil
}