fscan/webscan/lib/Eval.go

package lib

import (
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand" //nolint:gosec // G404: math/rand用于生成测试数据，非加密用途
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/cel-go/cel"
	"github.com/google/cel-go/checker/decls"
	"github.com/google/cel-go/common/types"
	"github.com/google/cel-go/common/types/ref"
	"github.com/google/cel-go/interpreter/functions"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
	exprpb "google.golang.org/genproto/googleapis/api/expr/v1alpha1"
)

// 基础CEL环境缓存（避免重复创建，减少内存分配）
var (
	baseEnvMu      sync.Mutex
	baseEnvInited  bool
	baseEnv        *cel.Env
	baseProgramOpt []cel.ProgramOption
)

const maxPOCResponseBodyBytes = 8 << 20

// 包级POC配置（atomic 保证并发安全）
var pocDNSLog atomic.Bool

// InitPOCConfig 初始化POC配置（在扫描开始前调用）
// 这样CEL回调函数可以使用包级变量而非GetGlobalConfig
func InitPOCConfig(dnsLog bool) {
	pocDNSLog.Store(dnsLog)
}

// NewEnv 创建一个新的 CEL 环境（使用缓存避免重复注册函数）
func NewEnv(c *CustomLib) (*cel.Env, error) {
	cachedCELEnvOnce.Do(func() {
		cachedCELEnv, cachedCELEnvErr = cel.NewEnv(cel.Lib(c))
	})
	return cachedCELEnv, cachedCELEnvErr
}

// initBaseEnv 初始化基础CEL环境（失败后允许重试）
func initBaseEnv() {
	baseEnvMu.Lock()
	defer baseEnvMu.Unlock()
	if baseEnvInited {
		return
	}

	var allDeclarations []*exprpb.Decl
	allDeclarations = append(allDeclarations, registerStringDeclarations()...)
	allDeclarations = append(allDeclarations, registerEncodingDeclarations()...)
	allDeclarations = append(allDeclarations, registerCryptoDeclarations()...)
	allDeclarations = append(allDeclarations, registerRandomDeclarations()...)
	allDeclarations = append(allDeclarations, registerMiscDeclarations()...)

	var allImplementations []*functions.Overload
	allImplementations = append(allImplementations, registerStringImplementations()...)
	allImplementations = append(allImplementations, registerEncodingImplementations()...)
	allImplementations = append(allImplementations, registerCryptoImplementations()...)
	allImplementations = append(allImplementations, registerRandomImplementations()...)
	allImplementations = append(allImplementations, registerMiscImplementations()...)

	//nolint:staticcheck // SA1019: cel.Functions已废弃但CEL库尚未提供替代方案
	baseProgramOpt = []cel.ProgramOption{cel.Functions(allImplementations...)}

	var err error
	baseEnv, err = cel.NewEnv(
		cel.Container("lib"),
		cel.Types(&UrlType{}, &Request{}, &Response{}, &Reverse{}),
		//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
		cel.Declarations(
			decls.NewIdent("request", decls.NewObjectType("lib.Request"), nil),
			decls.NewIdent("response", decls.NewObjectType("lib.Response"), nil),
			decls.NewIdent("reverse", decls.NewObjectType("lib.Reverse"), nil),
		),
		//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
		cel.Declarations(allDeclarations...),
	)
	if err != nil {
		common.LogError(i18n.Tr("webscan_cel_init_failed", err))
		return
	}
	baseEnvInited = true
}

// GetBaseEnv 获取基础CEL环境
func GetBaseEnv() *cel.Env {
	initBaseEnv()
	return baseEnv
}

// GetBaseProgramOptions 获取基础程序选项
func GetBaseProgramOptions() []cel.ProgramOption {
	initBaseEnv()
	return baseProgramOpt
}

// ExtendEnvWithVars 扩展基础环境，添加POC特定的变量声明
func ExtendEnvWithVars(varDecls []*exprpb.Decl) (*cel.Env, error) {
	base := GetBaseEnv()
	if base == nil {
		return nil, fmt.Errorf("%s", i18n.GetText("webscan_cel_env_not_initialized"))
	}
	if len(varDecls) == 0 {
		return base, nil
	}
	//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
	return base.Extend(cel.Declarations(varDecls...))
}

// MakeVarDecl 根据变量名和表达式创建变量声明
func MakeVarDecl(key, value string) *exprpb.Decl {
	switch {
	case strings.HasPrefix(value, "randomInt"):
		//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
		return decls.NewIdent(key, decls.Int, nil)
	case strings.HasPrefix(value, "newReverse"):
		//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
		return decls.NewIdent(key, decls.NewObjectType("lib.Reverse"), nil)
	default:
		//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
		return decls.NewIdent(key, decls.String, nil)
	}
}

// CelProgCache 缓存编译后的 CEL Program，避免同一 POC 内重复编译
// 在 executePoc 中创建，同一个 POC 的所有规则/参数组合共享
type CelProgCache map[string]cel.Program

// Evaluate 评估 CEL 表达式（无缓存，用于 Set/Sets 求值等低频路径）
func Evaluate(env *cel.Env, expression string, params map[string]interface{}) (ref.Val, error) {
	return EvaluateCached(env, expression, params, nil)
}

// EvaluateCached 评估 CEL 表达式（带编译缓存，用于规则执行热路径）
func EvaluateCached(env *cel.Env, expression string, params map[string]interface{}, cache CelProgCache) (ref.Val, error) {
	if expression == "" {
		return types.Bool(true), nil
	}

	var program cel.Program

	if cache != nil {
		if cached, ok := cache[expression]; ok {
			program = cached
		}
	}

	if program == nil {
		ast, issues := env.Compile(expression)
		if issues.Err() != nil {
			return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_expression_compile_failed"), issues.Err())
		}

		var err error
		program, err = env.Program(ast, GetBaseProgramOptions()...)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_program_create_failed"), err)
		}

		if cache != nil {
			cache[expression] = program
		}
	}

	result, _, err := program.Eval(params)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_expression_eval_failed"), err)
	}

	return result, nil
}

// URLTypeToString 将 TargetURL 结构体转换为字符串
func URLTypeToString(u *UrlType) string {
	var builder strings.Builder

	// 处理 scheme 部分
	if u.Scheme != "" {
		builder.WriteString(u.Scheme)
		builder.WriteByte(':')
	}

	// 处理 host 部分
	if u.Scheme != "" || u.Host != "" {
		if u.Host != "" || u.Path != "" {
			builder.WriteString("//")
		}
		if host := u.Host; host != "" {
			builder.WriteString(urlTypeHost(host))
		}
	}

	// 处理 path 部分
	path := u.Path
	if path != "" && path[0] != '/' && u.Host != "" {
		builder.WriteByte('/')
	}

	// 处理相对路径
	if builder.Len() == 0 {
		if i := strings.IndexByte(path, ':'); i > -1 && strings.IndexByte(path[:i], '/') == -1 {
			builder.WriteString("./")
		}
	}
	builder.WriteString(path)

	// 处理查询参数
	if u.Query != "" {
		builder.WriteByte('?')
		builder.WriteString(u.Query)
	}

	// 处理片段标识符
	if u.Fragment != "" {
		builder.WriteByte('#')
		builder.WriteString(u.Fragment)
	}

	return builder.String()
}

// CustomLib 自定义CEL库配置
type CustomLib struct {
	envOptions     []cel.EnvOption
	programOptions []cel.ProgramOption
}

// 缓存CustomLib实例和CEL环境，避免重复注册CEL函数导致冲突
var (
	cachedCustomLib     CustomLib
	cachedCustomLibOnce sync.Once
	cachedCELEnv        *cel.Env
	cachedCELEnvOnce    sync.Once
	cachedCELEnvErr     error
)

// NewEnvOption 创建新的CEL环境配置（使用缓存避免重复注册）
func NewEnvOption() CustomLib {
	cachedCustomLibOnce.Do(func() {
		cachedCustomLib = createCustomLib()
	})
	return cachedCustomLib
}

// createCustomLib 实际创建CustomLib（只执行一次）
func createCustomLib() CustomLib {
	c := CustomLib{}

	// 收集所有函数声明
	var allDeclarations []*exprpb.Decl
	allDeclarations = append(allDeclarations, registerStringDeclarations()...)
	allDeclarations = append(allDeclarations, registerEncodingDeclarations()...)
	allDeclarations = append(allDeclarations, registerCryptoDeclarations()...)
	allDeclarations = append(allDeclarations, registerRandomDeclarations()...)
	allDeclarations = append(allDeclarations, registerMiscDeclarations()...)

	c.envOptions = []cel.EnvOption{
		cel.Container("lib"),
		cel.Types(&UrlType{}, &Request{}, &Response{}, &Reverse{}),
		//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
		cel.Declarations(
			decls.NewIdent("request", decls.NewObjectType("lib.Request"), nil),
			decls.NewIdent("response", decls.NewObjectType("lib.Response"), nil),
			decls.NewIdent("reverse", decls.NewObjectType("lib.Reverse"), nil),
		),
		//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
		cel.Declarations(allDeclarations...),
	}

	// 收集所有函数实现
	var allImplementations []*functions.Overload
	allImplementations = append(allImplementations, registerStringImplementations()...)
	allImplementations = append(allImplementations, registerEncodingImplementations()...)
	allImplementations = append(allImplementations, registerCryptoImplementations()...)
	allImplementations = append(allImplementations, registerRandomImplementations()...)
	allImplementations = append(allImplementations, registerMiscImplementations()...)

	c.programOptions = []cel.ProgramOption{
		//nolint:staticcheck // SA1019: cel.Functions已废弃但CEL库尚未提供替代方案
		cel.Functions(allImplementations...),
	}

	return c
}

// CompileOptions 返回环境编译选项
func (c *CustomLib) CompileOptions() []cel.EnvOption {
	return c.envOptions
}

// ProgramOptions 返回程序运行选项
// 返回空切片，函数实现通过 GetBaseProgramOptions() 在 Evaluate() 时注入
// 这避免了多次创建环境时重复注册函数导致的冲突
func (c *CustomLib) ProgramOptions() []cel.ProgramOption {
	return nil
}

// UpdateCompileOptions 更新编译选项，处理不同类型的变量声明
func (c *CustomLib) UpdateCompileOptions(args StrMap) {
	for _, item := range args {
		key, value := item.Key, item.Value

		// 根据函数前缀确定变量类型
		var declaration *exprpb.Decl
		switch {
		case strings.HasPrefix(value, "randomInt"):
			// randomInt 函数返回整型
			//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
			declaration = decls.NewIdent(key, decls.Int, nil)
		case strings.HasPrefix(value, "newReverse"):
			// newReverse 函数返回 Reverse 对象
			//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
			declaration = decls.NewIdent(key, decls.NewObjectType("lib.Reverse"), nil)
		default:
			// 默认声明为字符串类型
			//nolint:staticcheck // SA1019: decls.NewIdent已废弃但CEL库尚未提供替代方案
			declaration = decls.NewIdent(key, decls.String, nil)
		}

		//nolint:staticcheck // SA1019: cel.Declarations已废弃但CEL库尚未提供替代方案
		c.envOptions = append(c.envOptions, cel.Declarations(declaration))
	}
}

// 随机数生成器（带互斥锁保护，确保并发安全）
//
//nolint:gosec // G404: 用于生成测试数据，非加密用途
var (
	randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
	randMu     sync.Mutex
)

// randomLowercase 生成指定长度的小写字母随机字符串
func randomLowercase(n int) string {
	const lowercase = "abcdefghijklmnopqrstuvwxyz"
	randMu.Lock()
	defer randMu.Unlock()
	return RandomStr(randSource, lowercase, n)
}

// randomUppercase 生成指定长度的大写字母随机字符串
func randomUppercase(n int) string {
	const uppercase = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	randMu.Lock()
	defer randMu.Unlock()
	return RandomStr(randSource, uppercase, n)
}

// randomString 生成指定长度的随机字符串（包含大小写字母和数字）
func randomString(n int) string {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	randMu.Lock()
	defer randMu.Unlock()
	return RandomStr(randSource, charset, n)
}

// reverseCheck 检查 DNS 记录是否存在
// 使用包级pocDNSLog变量，由InitPOCConfig初始化
func reverseCheck(r *Reverse, timeout int64) bool {
	// 检查必要条件（使用包级配置变量）
	if ceyeAPI == "" || r.Domain == "" || !pocDNSLog.Load() {
		return false
	}

	// 等待指定时间
	time.Sleep(time.Second * time.Duration(timeout))

	// 提取子域名
	sub := strings.Split(r.Domain, ".")[0]

	// 构造 API 请求 TargetURL
	apiURL := fmt.Sprintf("http://api.ceye.io/v1/records?token=%s&type=dns&filter=%s",
		ceyeAPI, sub)

	// 创建并发送请求（带超时控制，避免 ceye API 无响应时阻塞）
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, "GET", apiURL, nil)
	if err != nil {
		return false
	}
	// nil session: CEL 回调无法获取 session，回退到全局限速（反连检查请求量极低，可接受）
	resp, err := DoRequest(req, false, nil)
	if err != nil {
		return false
	}

	// 检查响应内容
	hasData := !bytes.Contains(resp.Body, []byte(`"data": []`))
	isOK := bytes.Contains(resp.Body, []byte(`"message": "OK"`))

	if hasData && isOK {
		common.LogDebug(apiURL)
		return true
	}
	return false
}

// RandomStr 生成指定长度的随机字符串
func RandomStr(randSource *rand.Rand, letterBytes string, n int) string {
	if n <= 0 || letterBytes == "" {
		return ""
	}
	const (
		// 用 6 位比特表示一个字母索引
		letterIdxBits = 6
		// 生成掩码：000111111
		letterIdxMask = 1<<letterIdxBits - 1
		// 63 位能存储的字母索引数量
		letterIdxMax = 63 / letterIdxBits
	)

	// 预分配结果数组
	randBytes := make([]byte, n)

	// 使用位操作生成随机字符串
	for i, cache, remain := n-1, randSource.Int63(), letterIdxMax; i >= 0; {
		// 当可用的随机位用完时，重新获取随机数
		if remain == 0 {
			cache, remain = randSource.Int63(), letterIdxMax
		}

		// 获取字符集中的随机索引
		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
			randBytes[i] = letterBytes[idx]
			i--
		}

		// 右移已使用的位，更新计数器
		cache >>= letterIdxBits
		remain--
	}

	return string(randBytes)
}

// DoRequest 执行 HTTP 请求
// session 为 nil 时回退到全局 state（兼容 CEL runtime 等无 session 场景）
func DoRequest(req *http.Request, redirect bool, session *common.ScanSession) (*Response, error) {
	// 处理请求头
	if req.Body != nil && req.Body != http.NoBody {
		body, err := io.ReadAll(req.Body)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_request_body_read_failed"), err)
		}
		_ = req.Body.Close()
		req.Body = io.NopCloser(bytes.NewReader(body))
		req.GetBody = func() (io.ReadCloser, error) {
			return io.NopCloser(bytes.NewReader(body)), nil
		}
		req.ContentLength = int64(len(body))

		// 设置 Content-Length
		req.Header.Set("Content-Length", strconv.FormatInt(req.ContentLength, 10))

		// 如果未指定 Content-Type，设置默认值
		if req.Header.Get("Content-Type") == "" {
			req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		}
	}

	// 执行请求
	// 检查发包限制
	var state *common.State
	if session != nil {
		state = session.State
		if canSend, err := common.CanSendPacketWith(session.Config, state); !canSend {
			reason := ""
			if err != nil {
				reason = err.Error()
			}
			common.LogError(i18n.Tr("webscan_request_restricted", req.URL.String(), reason))
			return nil, fmt.Errorf("%s", i18n.Tr("network_rate_limited", reason))
		}
	} else {
		state = common.GetGlobalState()
		if canSend, reason := common.CanSendPacket(); !canSend {
			common.LogError(i18n.Tr("webscan_request_restricted", req.URL.String(), reason))
			return nil, fmt.Errorf("%s", i18n.Tr("network_rate_limited", reason))
		}
	}

	var (
		oResp *http.Response
		err   error
	)

	if redirect {
		oResp, err = requestClient(true).Do(req)
	} else {
		oResp, err = requestClient(false).Do(req)
	}

	// 标准TLS握手级别失败时，尝试国密TLS客户端
	// 跳过连接超时、拒绝等非 TLS 相关错误，避免无意义的国密握手尝试
	if err != nil && req.URL.Scheme == "https" && maybeGMTLSError(err) {
		if req.GetBody != nil {
			if body, bodyErr := req.GetBody(); bodyErr == nil {
				req.Body = body
			}
		}
		if redirect {
			if clientGM := gmRequestClient(true); clientGM != nil {
				if oResp2, err2 := clientGM.Do(req); err2 == nil {
					oResp, err = oResp2, nil
				}
			}
		} else {
			if clientGM := gmRequestClient(false); clientGM != nil {
				if oResp2, err2 := clientGM.Do(req); err2 == nil {
					oResp, err = oResp2, nil
				}
			}
		}
	}

	if err != nil {
		// HTTP请求失败，计为TCP失败
		state.IncrementTCPFailedPacketCount()
		return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_request_execute_failed"), err)
	}

	// HTTP请求成功，计为TCP成功
	state.IncrementTCPSuccessPacketCount()
	defer func() { _ = oResp.Body.Close() }()

	// 解析响应
	resp, err := ParseResponse(oResp)
	if err != nil {
		common.LogError(i18n.Tr("webscan_response_parse_failed", err))
	}

	return resp, err
}

func requestClient(redirect bool) *http.Client {
	if redirect {
		if Client != nil {
			return Client
		}
		return http.DefaultClient
	}
	if ClientNoRedirect != nil {
		return ClientNoRedirect
	}
	return &http.Client{
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
}

func gmRequestClient(redirect bool) *http.Client {
	if redirect {
		return ClientGM
	}
	return ClientNoRedirectGM
}

// ParseURL 解析 TargetURL 并转换为自定义 TargetURL 类型
func ParseURL(u *url.URL) *UrlType {
	return &UrlType{
		Scheme:   u.Scheme,
		Domain:   u.Hostname(),
		Host:     u.Host,
		Port:     u.Port(),
		Path:     u.EscapedPath(),
		Query:    u.RawQuery,
		Fragment: u.Fragment,
	}
}

func urlTypeHost(host string) string {
	if strings.HasPrefix(host, "[") {
		return host
	}
	if ip := net.ParseIP(host); ip != nil && strings.Contains(host, ":") {
		return "[" + host + "]"
	}
	return host
}

// ParseRequest 将标准 HTTP 请求转换为自定义请求对象
func ParseRequest(oReq *http.Request) (*Request, error) {
	req := &Request{
		Method:      oReq.Method,
		URL:         ParseURL(oReq.URL),
		Headers:     make(map[string]string),
		ContentType: oReq.Header.Get("Content-Type"),
	}

	// 复制请求头
	for k := range oReq.Header {
		req.Headers[k] = oReq.Header.Get(k)
	}

	// 处理请求体
	if oReq.Body != nil && oReq.Body != http.NoBody {
		data, err := io.ReadAll(oReq.Body)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_request_body_read_failed"), err)
		}
		req.Body = data
		// 重新设置请求体，允许后续重复读取
		oReq.Body = io.NopCloser(bytes.NewBuffer(data))
	}

	return req, nil
}

// ParseResponse 将标准 HTTP 响应转换为自定义响应对象
func ParseResponse(oResp *http.Response) (*Response, error) {
	var respURL *UrlType
	if oResp.Request != nil {
		respURL = ParseURL(oResp.Request.URL)
	} else {
		respURL = &UrlType{}
	}
	resp := Response{
		Status:      int32(oResp.StatusCode),
		URL:         respURL,
		Headers:     make(map[string]string),
		ContentType: oResp.Header.Get("Content-Type"),
	}

	// 复制响应头，合并多值头部为分号分隔的字符串
	for k := range oResp.Header {
		resp.Headers[k] = strings.Join(oResp.Header.Values(k), ";")
	}

	// 读取并解析响应体
	body, err := getRespBody(oResp)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_response_body_process_failed"), err)
	}
	resp.Body = body

	return &resp, nil
}

// getRespBody 读取 HTTP 响应体并处理可能的 gzip 压缩
func getRespBody(oResp *http.Response) ([]byte, error) {
	// 读取原始响应体
	body, err := io.ReadAll(io.LimitReader(oResp.Body, maxPOCResponseBodyBytes))
	if err != nil && !errors.Is(err, io.EOF) && len(body) == 0 {
		return nil, err
	}

	// 处理 gzip 压缩
	if strings.Contains(oResp.Header.Get("Content-Encoding"), "gzip") {
		reader, err := gzip.NewReader(bytes.NewReader(body))
		if err != nil {
			return body, nil // 如果解压失败，返回原始数据
		}
		defer func() { _ = reader.Close() }()

		decompressed, err := io.ReadAll(io.LimitReader(reader, maxPOCResponseBodyBytes))
		if err != nil && !errors.Is(err, io.EOF) && len(decompressed) == 0 {
			return nil, err
		}
		if len(decompressed) == 0 && len(body) != 0 {
			return body, nil
		}
		return decompressed, nil
	}

	return body, nil
}

func maybeGMTLSError(err error) bool {
	if err == nil {
		return false
	}
	s := err.Error()
	return strings.Contains(s, "handshake failure") ||
		strings.Contains(s, "protocol version") ||
		strings.Contains(s, "no mutual") ||
		strings.Contains(s, "cipher suite")
}