ResolveSpec/pkg/security/hooks.go

package security

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/bitechdev/ResolveSpec/pkg/common"
	"github.com/bitechdev/ResolveSpec/pkg/logger"
	"github.com/bitechdev/ResolveSpec/pkg/modelregistry"
)

// SecurityContext is a generic interface that any spec can implement to integrate with security features
// This interface abstracts the common security context needs across different specs
type SecurityContext interface {
	GetContext() context.Context
	GetUserID() (int, bool)
	// GetUserRef returns an opaque user identifier for row security lookups.
	// Unlike GetUserID, it is not required to be an integer: implementations backed by
	// non-integer identifiers (e.g. UUIDs) can return a string, or the full
	// *security.UserContext so a RowSecurityProvider can read JWT claims directly.
	GetUserRef() (any, bool)
	GetSchema() string
	GetEntity() string
	GetModel() interface{}
	GetQuery() interface{}
	SetQuery(interface{})
	GetResult() interface{}
	SetResult(interface{})
}

// loadSecurityRules loads security configuration for the user and entity (generic version)
func loadSecurityRules(secCtx SecurityContext, securityList *SecurityList) error {
	// Extract user ID from context
	userID, ok := secCtx.GetUserID()
	if !ok {
		logger.Warn("No user ID in context for security check")
		return nil
	}

	schema := secCtx.GetSchema()
	tablename := secCtx.GetEntity()

	logger.Debug("Loading security rules for user=%d, schema=%s, table=%s", userID, schema, tablename)

	// Load column security rules using the provider
	err := securityList.LoadColumnSecurity(secCtx.GetContext(), userID, schema, tablename, false)
	if err != nil {
		logger.Warn("Failed to load column security: %v", err)
		// Don't fail the request if no security rules exist
		// return err
	}

	// Load row security rules using the provider. Row security uses the opaque
	// user ref (not the int-only user ID) so non-integer user identifiers work.
	userRef, refOK := secCtx.GetUserRef()
	if !refOK {
		userRef = userID
	}
	_, err = securityList.LoadRowSecurity(secCtx.GetContext(), userRef, schema, tablename, false)
	if err != nil {
		logger.Warn("Failed to load row security: %v", err)
		// Don't fail the request if no security rules exist
		// return err
	}

	return nil
}

// applyRowSecurity applies row-level security filters to the query (generic version)
func applyRowSecurity(secCtx SecurityContext, securityList *SecurityList) error {
	userRef, ok := secCtx.GetUserRef()
	if !ok {
		userID, idOK := secCtx.GetUserID()
		if !idOK {
			return nil // No user context, skip
		}
		userRef = userID
	}

	schema := secCtx.GetSchema()
	tablename := secCtx.GetEntity()

	// Get row security template
	rowSec, err := securityList.GetRowSecurityTemplate(userRef, schema, tablename)
	if err != nil {
		if errors.Is(err, ErrNoRowSecurity) {
			// No row security defined for this user/table: nothing to apply.
			logger.Debug("No row security for %s.%s", schema, tablename)
			return nil
		}
		// Anything else (including a recovered panic) fails closed.
		return fmt.Errorf("row security failed for %s.%s: %w", schema, tablename, err)
	}

	// Check if user has a blocking rule
	if rowSec.HasBlock {
		logger.Warn("User %v blocked from accessing %s.%s", userRef, schema, tablename)
		return fmt.Errorf("access denied to %s", tablename)
	}

	// If there's a security template, apply it as a WHERE clause
	if rowSec.Template != "" {
		model := secCtx.GetModel()
		if model == nil {
			logger.Debug("No model available for row security on %s.%s", schema, tablename)
			return nil
		}

		// Get primary key name from model
		modelType := reflect.TypeOf(model)
		if modelType.Kind() == reflect.Pointer {
			modelType = modelType.Elem()
		}

		// Find primary key field
		pkName := "id" // default
		for i := 0; i < modelType.NumField(); i++ {
			field := modelType.Field(i)
			if tag := field.Tag.Get("bun"); tag != "" {
				// Check for primary key tag
				if contains(tag, "pk") || contains(tag, "primary_key") {
					if sqlName := extractSQLName(tag); sqlName != "" {
						pkName = sqlName
					}
					break
				}
			}
		}

		// Generate the WHERE clause and bind arguments from the template
		whereClause, whereArgs, err := rowSec.GetTemplate(pkName, modelType)
		if err != nil {
			return fmt.Errorf("row security failed for %s.%s: %w", schema, tablename, err)
		}

		logger.Debug("Applying row security filter on %s.%s: %s", schema, tablename, whereClause)

		// A filter that cannot be attached must fail the request; silently
		// skipping it would expose every row.
		selectQuery, ok := secCtx.GetQuery().(common.SelectQuery)
		if !ok {
			return fmt.Errorf("row security: query type %T on %s.%s does not support Where", secCtx.GetQuery(), schema, tablename)
		}
		secCtx.SetQuery(selectQuery.Where(whereClause, whereArgs...))
	}

	return nil
}

// applyColumnSecurity applies column-level security (masking/hiding) to results (generic version)
func applyColumnSecurity(secCtx SecurityContext, securityList *SecurityList) error {
	userID, ok := secCtx.GetUserID()
	if !ok {
		return nil // No user context, skip
	}

	schema := secCtx.GetSchema()
	tablename := secCtx.GetEntity()

	// Get result data
	result := secCtx.GetResult()
	if result == nil {
		return nil
	}

	logger.Debug("Applying column security for user=%d, schema=%s, table=%s", userID, schema, tablename)

	model := secCtx.GetModel()
	if model == nil {
		logger.Debug("No model available for column security on %s.%s", schema, tablename)
		return nil
	}

	// Get model type
	modelType := reflect.TypeOf(model)
	if modelType.Kind() == reflect.Pointer {
		modelType = modelType.Elem()
	}

	// Apply column security masking
	resultValue := reflect.ValueOf(result)
	if resultValue.Kind() == reflect.Pointer {
		resultValue = resultValue.Elem()
	}

	maskedResult, err := securityList.ApplyColumnSecurity(resultValue, modelType, userID, schema, tablename)
	if err != nil {
		if errors.Is(err, ErrNoColumnSecurity) {
			// No rules for this user/table: nothing to mask.
			logger.Debug("No column security for %s.%s", schema, tablename)
			return nil
		}
		// Anything else (including a recovered panic) fails closed rather
		// than returning unmasked data.
		return fmt.Errorf("column security failed for %s.%s: %w", schema, tablename, err)
	}

	// Update the result with masked data
	if maskedResult.IsValid() && maskedResult.CanInterface() {
		secCtx.SetResult(maskedResult.Interface())
	}

	return nil
}

// logDataAccess logs all data access for audit purposes (generic version)
func logDataAccess(secCtx SecurityContext) error {
	userID, _ := secCtx.GetUserID()

	logger.Info("AUDIT: User %d accessed %s.%s",
		userID,
		secCtx.GetSchema(),
		secCtx.GetEntity(),
	)

	// TODO: Write to audit log table or external audit service
	// auditLog := AuditLog{
	//     UserID:    userID,
	//     Schema:    secCtx.GetSchema(),
	//     Entity:    secCtx.GetEntity(),
	//     Action:    "READ",
	//     Timestamp: time.Now(),
	// }
	// db.Create(&auditLog)

	return nil
}

// LogDataAccess is a public wrapper for logDataAccess that accepts a SecurityContext
// This allows other packages to use the audit logging functionality
func LogDataAccess(secCtx SecurityContext) error {
	return logDataAccess(secCtx)
}

// LoadSecurityRules is a public wrapper for loadSecurityRules that accepts a SecurityContext
// This allows other packages to load security rules using the generic interface
func LoadSecurityRules(secCtx SecurityContext, securityList *SecurityList) error {
	return loadSecurityRules(secCtx, securityList)
}

// ApplyRowSecurity is a public wrapper for applyRowSecurity that accepts a SecurityContext
// This allows other packages to apply row-level security using the generic interface
func ApplyRowSecurity(secCtx SecurityContext, securityList *SecurityList) error {
	// Spec adapters that expose the dispatched operation can enforce the same
	// model-rule bypass even when ApplyRowSecurity is called directly.
	if operationCtx, ok := secCtx.(interface{ GetOperation() string }); ok &&
		ShouldSkipRowSecurity(secCtx, operationCtx.GetOperation()) {
		return nil
	}
	return applyRowSecurity(secCtx, securityList)
}

// ShouldSkipRowSecurity reports whether row-security enforcement should be
// skipped for the operation. It uses the same model-rule resolution as
// CheckModelAuthAllowed so the model registry remains the single source of
// truth for security behavior.
func ShouldSkipRowSecurity(secCtx SecurityContext, operation string) bool {
	rules, ok := resolveModelRules(secCtx)
	if !ok {
		return false
	}

	return rules.SecurityDisabled || (operation == "read" && rules.CanPublicRead)
}

// IsModelSecurityDisabled reports whether all model-level security processing
// is disabled for the model. This is distinct from ShouldSkipRowSecurity:
// CanPublicRead skips row filtering for reads but must still allow other read
// security, such as column masking, to be loaded.
func IsModelSecurityDisabled(secCtx SecurityContext) bool {
	rules, ok := resolveModelRules(secCtx)
	return ok && rules.SecurityDisabled
}

// ApplyColumnSecurity is a public wrapper for applyColumnSecurity that accepts a SecurityContext
// This allows other packages to apply column-level security using the generic interface
func ApplyColumnSecurity(secCtx SecurityContext, securityList *SecurityList) error {
	return applyColumnSecurity(secCtx, securityList)
}

// checkModelUpdateAllowed returns an error if CanUpdate is false for the model.
// Rules are read from context (set by NewModelAuthMiddleware) with a fallback to the model registry.
func checkModelUpdateAllowed(secCtx SecurityContext) error {
	rules, ok := GetModelRulesFromContext(secCtx.GetContext())
	if !ok {
		schema := secCtx.GetSchema()
		entity := secCtx.GetEntity()
		var err error
		if schema != "" {
			rules, err = modelregistry.GetModelRulesByName(fmt.Sprintf("%s.%s", schema, entity))
		}
		if err != nil || schema == "" {
			rules, err = modelregistry.GetModelRulesByName(entity)
		}
		if err != nil {
			if errors.Is(err, modelregistry.ErrModelNotFound) {
				return nil // model not registered, allow by default
			}
			return err
		}
	}
	if !rules.CanUpdate {
		return fmt.Errorf("update not allowed for %s", secCtx.GetEntity())
	}
	return nil
}

// checkModelDeleteAllowed returns an error if CanDelete is false for the model.
// Rules are read from context (set by NewModelAuthMiddleware) with a fallback to the model registry.
func checkModelDeleteAllowed(secCtx SecurityContext) error {
	rules, ok := GetModelRulesFromContext(secCtx.GetContext())
	if !ok {
		schema := secCtx.GetSchema()
		entity := secCtx.GetEntity()
		var err error
		if schema != "" {
			rules, err = modelregistry.GetModelRulesByName(fmt.Sprintf("%s.%s", schema, entity))
		}
		if err != nil || schema == "" {
			rules, err = modelregistry.GetModelRulesByName(entity)
		}
		if err != nil {
			if errors.Is(err, modelregistry.ErrModelNotFound) {
				return nil // model not registered, allow by default
			}
			return err
		}
	}
	if !rules.CanDelete {
		return fmt.Errorf("delete not allowed for %s", secCtx.GetEntity())
	}
	return nil
}

// CheckModelAuthAllowed checks whether the requested operation is permitted based on
// model rules and the current user's authentication state. It is intended for use in
// a BeforeHandle hook, fired after model resolution.
//
// Logic:
//  1. Load model rules from context (set by NewModelAuthMiddleware) or fall back to registry.
//  2. SecurityDisabled → allow.
//  3. operation == "read" && CanPublicRead → allow.
//  4. operation == "create" && CanPublicCreate → allow.
//  5. operation == "update" && CanPublicUpdate → allow.
//  6. operation == "delete" && CanPublicDelete → allow.
//  7. Guest (UserID == 0) → return "authentication required".
//  8. Authenticated user → allow (operation-specific checks remain in BeforeUpdate/BeforeDelete).
func CheckModelAuthAllowed(secCtx SecurityContext, operation string) error {
	rules, ok := resolveModelRules(secCtx)
	if !ok {
		// Model not registered - fall through to auth check
		userID, _ := secCtx.GetUserID()
		if userID == 0 {
			return fmt.Errorf("authentication required")
		}
		return nil
	}

	if rules.SecurityDisabled {
		return nil
	}
	if operation == "read" && rules.CanPublicRead {
		return nil
	}
	if operation == "create" && rules.CanPublicCreate {
		return nil
	}
	if operation == "update" && rules.CanPublicUpdate {
		return nil
	}
	if operation == "delete" && rules.CanPublicDelete {
		return nil
	}

	userID, _ := secCtx.GetUserID()
	if userID == 0 {
		return fmt.Errorf("authentication required")
	}
	return nil
}

// resolveModelRules returns model rules from the request context first, then
// falls back to the schema-qualified and unqualified registry names.
func resolveModelRules(secCtx SecurityContext) (modelregistry.ModelRules, bool) {
	if rules, ok := GetModelRulesFromContext(secCtx.GetContext()); ok {
		return rules, true
	}

	schema := secCtx.GetSchema()
	entity := secCtx.GetEntity()
	var err error
	if schema != "" {
		var rules modelregistry.ModelRules
		rules, err = modelregistry.GetModelRulesByName(fmt.Sprintf("%s.%s", schema, entity))
		if err == nil {
			return rules, true
		}
	}

	rules, err := modelregistry.GetModelRulesByName(entity)
	if err != nil {
		return modelregistry.ModelRules{}, false
	}
	return rules, true
}

// CheckModelUpdateAllowed is the public wrapper for checkModelUpdateAllowed.
func CheckModelUpdateAllowed(secCtx SecurityContext) error {
	return checkModelUpdateAllowed(secCtx)
}

// CheckModelDeleteAllowed is the public wrapper for checkModelDeleteAllowed.
func CheckModelDeleteAllowed(secCtx SecurityContext) error {
	return checkModelDeleteAllowed(secCtx)
}

// Helper functions

func contains(s, substr string) bool {
	return len(s) >= len(substr) && s[:len(substr)] == substr ||
		len(s) > len(substr) && s[len(s)-len(substr):] == substr
}

func extractSQLName(tag string) string {
	// Simple parser for "column:name" or just "name"
	// This is a simplified version
	parts := splitTag(tag, ',')
	for _, part := range parts {
		if part != "" && !contains(part, ":") {
			return part
		}
		if contains(part, "column:") {
			return part[7:] // Skip "column:"
		}
	}
	return ""
}

func splitTag(tag string, sep rune) []string {
	return strings.FieldsFunc(tag, func(r rune) bool { return r == sep })
}