expr.go

package probe

import (
	"encoding/base64"
	"fmt"
	"math/rand/v2"
	"regexp"
	"strconv"
	"strings"
	"time"

	ex "github.com/expr-lang/expr"
	"github.com/expr-lang/expr/vm"
)

var (
	// Regular expression to find `{{ ... }}` patterns
	templateRegexp = regexp.MustCompile(`\{\{([^{}]+)\}\}`)
	templateStart  = "{{"
	templateEnd    = "}}"

	// Security: Maximum expression length and evaluation timeout
	maxExpressionLength = 1000000
	evaluationTimeout   = 5 * time.Second

	// Security: Maximum string length to prevent memory exhaustion
	maxStringLength = 1000000
)

type Expr struct{}

func (e *Expr) Options(env any) []ex.Option {
	// Security: Create a safe environment for expression evaluation
	safeEnv := e.createSafeEnvironment(env)

	return []ex.Option{
		ex.Env(safeEnv),
		// Security: Allow undefined variables but with safe environment only
		ex.AllowUndefinedVariables(),

		// Security: Disable dangerous built-in functions
		ex.DisableBuiltin("all"),
		ex.DisableBuiltin("any"),
		ex.DisableBuiltin("one"),
		ex.DisableBuiltin("filter"),
		ex.DisableBuiltin("map"),
		ex.DisableBuiltin("count"),

		// Security: Add only safe, whitelisted functions
		ex.Function(
			"match_json",
			func(params ...any) (any, error) {
				if len(params) != 2 {
					return false, fmt.Errorf("match_json requires exactly 2 parameters")
				}
				src, ok1 := params[0].(map[string]any)
				target, ok2 := params[1].(map[string]any)
				if !ok1 || !ok2 {
					return false, fmt.Errorf("match_json parameters must be objects")
				}
				return MatchJSON(src, target), nil
			},
		),
		ex.Function(
			"diff_json",
			func(params ...any) (any, error) {
				if len(params) != 2 {
					return nil, fmt.Errorf("diff_json requires exactly 2 parameters")
				}
				src, ok1 := params[0].(map[string]any)
				target, ok2 := params[1].(map[string]any)
				if !ok1 || !ok2 {
					return nil, fmt.Errorf("diff_json parameters must be objects")
				}
				return DiffJSON(src, target), nil
			},
		),
		ex.Function(
			"random_int",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("random_int requires exactly 1 parameter")
				}
				n, ok := params[0].(int)
				if !ok {
					// Try to convert float64 to int (common in JSON/expr)
					if f, ok := params[0].(float64); ok {
						n = int(f)
					} else {
						return nil, fmt.Errorf("random_int parameter must be an integer")
					}
				}
				if n <= 0 {
					return nil, fmt.Errorf("random_int parameter must be positive")
				}
				return rand.IntN(n), nil
			},
		),
		ex.Function(
			"random_str",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("random_str requires exactly 1 parameter")
				}
				length, ok := params[0].(int)
				if !ok {
					// Try to convert float64 to int (common in JSON/expr)
					if f, ok := params[0].(float64); ok {
						length = int(f)
					} else {
						return nil, fmt.Errorf("random_str parameter must be an integer")
					}
				}
				if length <= 0 {
					return nil, fmt.Errorf("random_str parameter must be positive")
				}
				if length > 1000000 {
					return nil, fmt.Errorf("random_str parameter must be <= 1000000")
				}

				const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
				b := make([]byte, length)
				for i := range b {
					b[i] = charset[rand.IntN(len(charset))]
				}

				return string(b), nil
			},
		),
		ex.Function(
			"unixtime",
			func(params ...any) (any, error) {
				if len(params) != 0 {
					return nil, fmt.Errorf("unixtime takes no parameters")
				}
				return time.Now().Unix(), nil
			},
		),
		ex.Function(
			"parse_float",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("parse_float requires exactly 1 parameter")
				}
				s, ok := params[0].(string)
				if !ok {
					return nil, fmt.Errorf("parse_float parameter must be a string")
				}
				f, err := strconv.ParseFloat(s, 64)
				if err != nil {
					return nil, fmt.Errorf("parse_float error: %w", err)
				}
				return f, nil
			},
		),
		ex.Function(
			"parse_int",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("parse_int requires exactly 1 parameter")
				}

				switch v := params[0].(type) {
				case string:
					i, err := strconv.ParseInt(v, 10, 64)
					if err != nil {
						return nil, fmt.Errorf("parse_int error: %w", err)
					}
					return i, nil
				case int:
					return int64(v), nil
				case int64:
					return v, nil
				case uint64:
					return int64(v), nil
				case float64:
					return int64(v), nil
				default:
					return nil, fmt.Errorf("parse_int parameter must be a string or number, got %T", v)
				}
			},
		),
		ex.Function(
			"encode_base64",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("encode_base64 requires exactly 1 parameter")
				}
				s, ok := params[0].(string)
				if !ok {
					return nil, fmt.Errorf("encode_base64 parameter must be a string")
				}
				if len(s) > maxStringLength {
					return nil, fmt.Errorf("encode_base64 parameter exceeds maximum length (%d chars)", maxStringLength)
				}
				return base64.StdEncoding.EncodeToString([]byte(s)), nil
			},
		),
		ex.Function(
			"parse_json",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("parse_json requires exactly 1 parameter")
				}
				s, ok := params[0].(string)
				if !ok {
					return nil, fmt.Errorf("parse_json parameter must be a string")
				}
				if len(s) > maxStringLength {
					return nil, fmt.Errorf("parse_json parameter exceeds maximum length (%d chars)", maxStringLength)
				}
				return ParseJSON(s)
			},
		),
		ex.Function(
			"decode_base64",
			func(params ...any) (any, error) {
				if len(params) != 1 {
					return nil, fmt.Errorf("decode_base64 requires exactly 1 parameter")
				}
				s, ok := params[0].(string)
				if !ok {
					return nil, fmt.Errorf("decode_base64 parameter must be a string")
				}
				if len(s) > maxStringLength {
					return nil, fmt.Errorf("decode_base64 parameter exceeds maximum length (%d chars)", maxStringLength)
				}
				decoded, err := base64.StdEncoding.DecodeString(s)
				if err != nil {
					return nil, fmt.Errorf("decode_base64 error: %w", err)
				}
				return string(decoded), nil
			},
		),
	}
}

// Security: Create a safe environment by filtering out dangerous variables
func (e *Expr) createSafeEnvironment(env any) any {
	envMap, ok := env.(map[string]any)
	if !ok {
		// For non-map types (like structs with expr tags), return as-is
		// This allows expr library to handle StepContext and similar safe structs
		return env
	}

	safeEnv := make(map[string]any)

	// Security: Whitelist safe environment variables and data
	for key, value := range envMap {
		if e.isSafeEnvKey(key) {
			safeEnv[key] = e.sanitizeValue(value)
		}
	}

	return safeEnv
}

// isSafeEnvKey reports whether the given environment variable name is
// safe to expose to expression evaluation. probe uses a blocklist
// model: any key whose upper-cased form *contains* one of the patterns
// below is rejected, and every other key is allowed. The substring
// match is intentional so compound names like "DB_PASSWORD" or
// "MY_API_KEY" still get rejected, at the cost of also rejecting
// unrelated keys that happen to embed one of these substrings (e.g.
// anything containing "KEY"). That conservative bias is deliberate.
//
// Historical note: this used to layer prefix / safeKeys / testEnvVars /
// fallback whitelists on top of the blocklist, but every one of those
// branches collapsed back to "anything that survived the blocklist is
// allowed" — none of them actually narrowed what was reachable. They
// were removed so the contract here matches the code.
func (e *Expr) isSafeEnvKey(key string) bool {
	upperKey := strings.ToUpper(key)
	blocklist := []string{
		// Shell and host metadata
		"PATH", "HOME", "USER", "USERNAME", "SHELL", "PWD",
		// Credential-bearing names
		"SECRET", "KEY", "TOKEN", "PASSWORD", "CREDENTIAL",
		"API_KEY", "PRIVATE", "CERT", "SSH",
	}
	for _, pattern := range blocklist {
		if strings.Contains(upperKey, pattern) {
			return false
		}
	}
	return true
}

// Security: Sanitize values to prevent injection
func (e *Expr) sanitizeValue(value any) any {
	switch v := value.(type) {
	case string:
		// Security: Limit string length to prevent memory exhaustion
		if len(v) > maxStringLength {
			return v[:maxStringLength] + GetTruncationMessage()
		}
		return v
	case map[string]any:
		safeMap := make(map[string]any)
		for k, val := range v {
			if e.isSafeEnvKey(k) {
				safeMap[k] = e.sanitizeValue(val)
			}
		}
		return safeMap
	case []any:
		// Security: Limit array size to prevent memory exhaustion
		if len(v) > 1000 {
			return v[:1000]
		}
		safeSlice := make([]any, len(v))
		for i, val := range v {
			safeSlice[i] = e.sanitizeValue(val)
		}
		return safeSlice
	default:
		return v
	}
}

// Security: Validate expression for dangerous patterns
func (e *Expr) validateExpression(expression string) error {
	// Security: Check expression length
	if len(expression) > maxExpressionLength {
		return fmt.Errorf("SECURITY: expression exceeds maximum length (%d chars)", maxExpressionLength)
	}

	lowerExpr := strings.ToLower(expression)

	// Security: Special validation for env. patterns - only allow safe environment variables
	if strings.Contains(lowerExpr, "env.") {
		return e.validateEnvAccess(expression)
	}

	return nil
}

// Security: Validate environment variable access patterns
func (e *Expr) validateEnvAccess(expression string) error {
	// Block access to dangerous environment variables
	dangerousEnvPatterns := []string{
		"env.secret", "env.password", "env.credential",
		"env.api_key", "env.private_key", "env.cert", "env.ssh_key", "env.path", "env.home",
	}

	lowerExpr := strings.ToLower(expression)
	for _, pattern := range dangerousEnvPatterns {
		if strings.Contains(lowerExpr, pattern) {
			return fmt.Errorf("SECURITY: attempt to access dangerous environment variable '%s'", pattern)
		}
	}

	return nil
}

func (e *Expr) EvalOrEvalTemplate(input string, env any) (string, error) {
	// Security: Validate input expression
	if err := e.validateExpression(input); err != nil {
		return "", fmt.Errorf("expression validation failed: %w", err)
	}

	if strings.Contains(input, templateStart) && strings.Contains(input, templateEnd) {
		return e.EvalTemplate(input, env)
	}
	output, err := e.Eval(input, env)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%v", output), nil
}

func (e *Expr) Eval(input string, env any) (any, error) {
	// Security: Validate expression before compilation
	if err := e.validateExpression(input); err != nil {
		return nil, fmt.Errorf("expression validation failed: %w", err)
	}

	program, err := ex.Compile(input, e.Options(env)...)
	if err != nil {
		return false, err
	}

	// Security: Execute with timeout to prevent infinite loops
	return e.executeWithTimeout(program, env)
}

// Security: Execute expression with timeout protection
func (e *Expr) executeWithTimeout(program *vm.Program, env any) (any, error) {
	type result struct {
		output any
		err    error
	}

	resultCh := make(chan result, 1)
	done := make(chan bool, 1)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				select {
				case resultCh <- result{nil, fmt.Errorf("expression execution panicked: %v", r)}:
				default:
				}
				done <- true
			}
		}()

		output, err := ex.Run(program, env)
		select {
		case resultCh <- result{output, err}:
		default:
		}
		done <- true
	}()

	select {
	case res := <-resultCh:
		return res.output, res.err
	case <-time.After(evaluationTimeout):
		return nil, fmt.Errorf("SECURITY: expression evaluation timed out after %v", evaluationTimeout)
	}
}

// isWholeStringTemplate checks if the input string contains only a single template expression
func isWholeStringTemplate(input string) bool {
	trimmed := strings.TrimSpace(input)
	if !strings.HasPrefix(trimmed, templateStart) || !strings.HasSuffix(trimmed, templateEnd) {
		return false
	}

	// Count template markers to ensure there's exactly one pair
	startCount := strings.Count(trimmed, templateStart)
	endCount := strings.Count(trimmed, templateEnd)

	return startCount == 1 && endCount == 1
}

// extractTemplateExpression extracts the expression from a whole string template
func extractTemplateExpression(input string) string {
	trimmed := strings.TrimSpace(input)
	if !strings.HasPrefix(trimmed, templateStart) || !strings.HasSuffix(trimmed, templateEnd) {
		return ""
	}

	// Remove template markers and trim whitespace
	expression := trimmed[len(templateStart) : len(trimmed)-len(templateEnd)]
	return strings.TrimSpace(expression)
}

func (e *Expr) EvalTemplate(input string, env any) (string, error) {
	// Security: Validate template input
	if err := e.validateExpression(input); err != nil {
		return "", fmt.Errorf("template validation failed: %w", err)
	}

	re := templateRegexp
	var evalError error

	// Replace matches with evaluated results
	result := re.ReplaceAllFunc([]byte(input), func(match []byte) []byte {
		// Security: Check if we've already encountered an error
		if evalError != nil {
			return match
		}

		// Extract the expression inside `{{ ... }}` using submatch
		submatch := re.FindStringSubmatch(string(match))
		if len(submatch) < 2 {
			evalError = fmt.Errorf("invalid template expression: %s", string(match))
			return []byte("[TemplateError: invalid expression]")
		}
		expression := strings.TrimSpace(submatch[1])

		// Security: Validate individual expression
		if err := e.validateExpression(expression); err != nil {
			evalError = fmt.Errorf("template expression validation failed: %w", err)
			return fmt.Appendf(nil, "[SecurityError: %s]", err.Error())
		}

		// Evaluate the expression using expr
		program, err := ex.Compile(expression, e.Options(env)...)
		if err != nil {
			return fmt.Appendf(nil, "[CompileError: %s]", err.Error())
		}

		// Security: Execute with timeout protection
		output, err := e.executeWithTimeout(program, env)
		if err != nil {
			return fmt.Appendf(nil, "[RuntimeError: %s]", err.Error())
		}

		// Convert the output to string with size limit
		outputStr := fmt.Sprintf("%v", output)
		if len(outputStr) > maxStringLength {
			outputStr = outputStr[:maxStringLength] + GetTruncationMessage()
		}

		return []byte(outputStr)
	})

	if evalError != nil {
		return "", evalError
	}

	return string(result), nil
}

func (e *Expr) EvalTemplateWithTypePreservation(input string, env any) (any, error) {
	// Security: Validate template input
	if err := e.validateExpression(input); err != nil {
		return "", fmt.Errorf("template validation failed: %w", err)
	}

	// Check if this is a whole string template for type preservation
	if isWholeStringTemplate(input) {
		expression := extractTemplateExpression(input)
		if expression == "" {
			return "", fmt.Errorf("failed to extract expression from template: %s", input)
		}

		// Use Eval directly to preserve type
		return e.Eval(expression, env)
	}

	// For partial templates, fall back to string processing
	return e.EvalTemplate(input, env)
}

func (e *Expr) EvalTemplateMap(input map[string]any, env any) map[string]any {
	results := make(map[string]any)

	for key, val := range input {
		// Security: Limit the number of processed keys to prevent DoS
		if len(results) > 1000 {
			results["_truncated"] = "Map processing truncated due to size limits"
			break
		}

		switch v := val.(type) {
		case string:
			output, err := e.EvalTemplateWithTypePreservation(v, env)
			if err != nil {
				// Security: Don't expose internal errors, use sanitized error
				results[key] = "[EvaluationError]"
				continue
			}
			results[key] = output

		case map[string]any:
			results[key] = e.EvalTemplateMap(v, env)

		case []any:
			results[key] = e.evalTemplateArray(v, env)

		default:
			results[key] = v
		}
	}

	return results
}

// evalTemplateArray evaluates templates in array elements
func (e *Expr) evalTemplateArray(input []any, env any) []any {
	results := make([]any, len(input))

	for i, val := range input {
		// Security: Limit the number of processed elements to prevent DoS
		if i > 1000 {
			results = append(results, "_truncated: Array processing truncated due to size limits")
			break
		}

		switch v := val.(type) {
		case string:
			output, err := e.EvalTemplateWithTypePreservation(v, env)
			if err != nil {
				// Security: Don't expose internal errors, use sanitized error
				results[i] = "[EvaluationError]"
				continue
			}
			results[i] = output

		case map[string]any:
			results[i] = e.EvalTemplateMap(v, env)

		case []any:
			results[i] = e.evalTemplateArray(v, env)

		default:
			results[i] = v
		}
	}

	return results
}