// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package nodes

import (
	"fmt"
	iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
	"github.com/TeaOSLab/EdgeNode/internal/conns"
	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
	"github.com/TeaOSLab/EdgeNode/internal/utils"
	"github.com/TeaOSLab/EdgeNode/internal/utils/agents"
	"github.com/TeaOSLab/EdgeNode/internal/utils/counters"
	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
	maputils "github.com/TeaOSLab/EdgeNode/internal/utils/maps"
	"github.com/TeaOSLab/EdgeNode/internal/utils/ttlcache"
	"github.com/TeaOSLab/EdgeNode/internal/waf"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/types"
	stringutil "github.com/iwind/TeaGo/utils/string"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
)

type ccBlockedCounter struct {
	count     int32
	updatedAt int64
}

var ccBlockedMap = maputils.NewFixedMap[string, *ccBlockedCounter](65535)

func (this *HTTPRequest) doCC() (block bool) {
	if this.nodeConfig == nil || this.ReqServer == nil {
		return
	}

	var validatePath = "/GE/CC/VALIDATOR"
	var maxConnections = 30

	// 策略
	var httpCCPolicy = this.nodeConfig.FindHTTPCCPolicyWithClusterId(this.ReqServer.ClusterId)
	var scope = firewallconfigs.FirewallScopeGlobal
	if httpCCPolicy != nil {
		if !httpCCPolicy.IsOn {
			return
		}
		if len(httpCCPolicy.RedirectsChecking.ValidatePath) > 0 {
			validatePath = httpCCPolicy.RedirectsChecking.ValidatePath
		}
		if httpCCPolicy.MaxConnectionsPerIP > 0 {
			maxConnections = httpCCPolicy.MaxConnectionsPerIP
		}
		scope = httpCCPolicy.FirewallScope()
	}

	var ccConfig = this.web.CC
	if ccConfig == nil || !ccConfig.IsOn || (this.RawReq.URL.Path != validatePath && !ccConfig.MatchURL(this.requestScheme()+"://"+this.ReqHost+this.Path())) {
		return
	}

	// 忽略常用文件
	if ccConfig.IgnoreCommonFiles {
		if len(this.RawReq.Referer()) > 0 {
			var ext = filepath.Ext(this.RawReq.URL.Path)
			if len(ext) > 0 && utils.IsCommonFileExtension(ext) {
				return
			}
		}
	}

	// 检查白名单
	var remoteAddr = this.requestRemoteAddr(true)

	// 检查是否为白名单直连
	if !Tea.IsTesting() && this.nodeConfig.IPIsAutoAllowed(remoteAddr) {
		return
	}

	// 是否在全局名单中
	canGoNext, isInAllowedList, _ := iplibrary.AllowIP(remoteAddr, this.ReqServer.Id)
	if !canGoNext {
		this.disableLog = true
		this.Close()
		return true
	}
	if isInAllowedList {
		return false
	}

	// WAF黑名单
	if waf.SharedIPBlackList.Contains(waf.IPTypeAll, scope, this.ReqServer.Id, remoteAddr) {
		this.disableLog = true
		this.Close()
		return true
	}

	// 检查是否启用QPS
	if ccConfig.MinQPSPerIP > 0 && this.RawReq.URL.Path != validatePath {
		var avgQPS = counters.SharedCounter.IncreaseKey("QPS:"+remoteAddr, 60) / 60
		if avgQPS <= 0 {
			avgQPS = 1
		}
		if avgQPS < types.Uint32(ccConfig.MinQPSPerIP) {
			return false
		}
	}

	// 检查连接数
	if conns.SharedMap.CountIPConns(remoteAddr) >= maxConnections {
		this.ccForbid(5)

		var forbiddenTimes = this.increaseCCCounter(remoteAddr)
		waf.SharedIPBlackList.RecordIP(waf.IPTypeAll, scope, this.ReqServer.Id, remoteAddr, fasttime.Now().Unix()+int64(forbiddenTimes*1800), 0, scope == firewallconfigs.FirewallScopeGlobal, 0, 0, "CC防护拦截：并发连接数超出"+types.String(maxConnections)+"个")

		return true
	}

	// GET302验证
	if ccConfig.EnableGET302 &&
		this.RawReq.Method == http.MethodGet &&
		!agents.SharedManager.ContainsIP(remoteAddr) /** 搜索引擎亲和性 **/ &&
		!strings.HasPrefix(this.RawReq.URL.Path, "/baidu_verify_") /** 百度验证 **/ &&
		!strings.HasPrefix(this.RawReq.URL.Path, "/google") /** Google验证 **/ {
		// 忽略搜索引擎
		var canSkip302 = false
		var ipResult = iplib.LookupIP(remoteAddr)
		if ipResult != nil && ipResult.IsOk() {
			var providerName = ipResult.ProviderName()
			canSkip302 = strings.Contains(providerName, "百度") || strings.Contains(providerName, "谷歌") || strings.Contains(providerName, "baidu") || strings.Contains(providerName, "google")
		}

		if !canSkip302 {
			// 检查参数
			var ccWhiteListKey = "HTTP-CC-GET302-" + remoteAddr
			var currentTime = fasttime.Now().Unix()
			if this.RawReq.URL.Path == validatePath {
				this.DisableAccessLog()
				this.DisableStat()

				// TODO 根据浏览器信息决定是否校验referer

				if !this.checkCCRedirects(httpCCPolicy, remoteAddr) {
					return true
				}

				var key = this.RawReq.URL.Query().Get("key")
				var pieces = strings.Split(key, ".") // key1.key2.timestamp
				if len(pieces) != 3 {
					this.ccForbid(1)
					return true
				}
				var urlKey = pieces[0]
				var timestampKey = pieces[1]
				var timestamp = pieces[2]

				var targetURL = this.RawReq.URL.Query().Get("url")

				var realURLKey = stringutil.Md5(nodeConfig().Secret + "@" + targetURL + "@" + remoteAddr)
				if urlKey != realURLKey {
					this.ccForbid(2)
					return true
				}

				// 校验时间
				if timestampKey != stringutil.Md5(nodeConfig().Secret+"@"+timestamp) {
					this.ccForbid(3)
					return true
				}

				var elapsedSeconds = currentTime - types.Int64(timestamp)
				if elapsedSeconds > 10 /** 10秒钟 **/ { // 如果校验时间过长，则可能阻止当前访问
					if elapsedSeconds > 300 /** 5分钟 **/ { // 如果超时时间过长就跳回原URL
						httpRedirect(this.writer, this.RawReq, targetURL, http.StatusFound)
					} else {
						this.ccForbid(4)
					}
					return true
				}

				// 加入到临时白名单
				ttlcache.SharedInt64Cache.Write(ccWhiteListKey, 1, currentTime+600 /** 10分钟 **/)

				// 跳转回原来URL
				httpRedirect(this.writer, this.RawReq, targetURL, http.StatusFound)

				return true
			} else {
				// 检查临时白名单
				if ttlcache.SharedInt64Cache.Read(ccWhiteListKey) == nil {
					if !this.checkCCRedirects(httpCCPolicy, remoteAddr) {
						return true
					}

					var urlKey = stringutil.Md5(nodeConfig().Secret + "@" + this.URL() + "@" + remoteAddr)
					var timestampKey = stringutil.Md5(nodeConfig().Secret + "@" + types.String(currentTime))

					// 跳转到验证URL
					this.DisableStat()
					httpRedirect(this.writer, this.RawReq, validatePath+"?key="+urlKey+"."+timestampKey+"."+types.String(currentTime)+"&url="+url.QueryEscape(this.URL()), http.StatusFound)

					return true
				}
			}
		}
	} else if this.RawReq.URL.Path == validatePath {
		// 直接跳回
		var targetURL = this.RawReq.URL.Query().Get("url")
		httpRedirect(this.writer, this.RawReq, targetURL, http.StatusFound)
		return true
	}

	// Key
	var ccKeys = []string{}
	if ccConfig.WithRequestPath {
		ccKeys = append(ccKeys, "HTTP-CC-"+remoteAddr+"-"+this.Path()) // 这里可以忽略域名，因为一个正常用户同时访问多个域名的可能性不大
	} else {
		ccKeys = append(ccKeys, "HTTP-CC-"+remoteAddr)
	}

	// 指纹
	if this.IsHTTPS && ccConfig.EnableFingerprint {
		var requestConn = this.RawReq.Context().Value(HTTPConnContextKey)
		if requestConn != nil {
			clientConn, ok := requestConn.(ClientConnInterface)
			if ok {
				var fingerprint = clientConn.Fingerprint()
				if len(fingerprint) > 0 {
					var fingerprintString = fmt.Sprintf("%x", fingerprint)
					if ccConfig.WithRequestPath {
						ccKeys = append(ccKeys, "HTTP-CC-"+fingerprintString+"-"+this.Path()) // 这里可以忽略域名，因为一个正常用户同时访问多个域名的可能性不大
					} else {
						ccKeys = append(ccKeys, "HTTP-CC-"+fingerprintString)
					}
				}
			}
		}
	}

	// 检查阈值
	var thresholds = ccConfig.Thresholds
	if len(thresholds) == 0 || ccConfig.UseDefaultThresholds {
		if httpCCPolicy != nil && len(httpCCPolicy.Thresholds) > 0 {
			thresholds = httpCCPolicy.Thresholds
		} else {
			thresholds = serverconfigs.DefaultHTTPCCThresholds
		}
	}
	if len(thresholds) == 0 {
		return
	}

	var currentTime = fasttime.Now().Unix()

	for _, threshold := range thresholds {
		if threshold.PeriodSeconds <= 0 || threshold.MaxRequests <= 0 {
			continue
		}

		for _, ccKey := range ccKeys {
			var count = counters.SharedCounter.IncreaseKey(ccKey+"-T"+types.String(threshold.PeriodSeconds), int(threshold.PeriodSeconds))
			if count >= types.Uint32(threshold.MaxRequests) {
				this.writeCode(http.StatusTooManyRequests, "Too many requests, please wait for a few minutes.", "访问过于频繁，请稍等片刻后再访问。")

				// 记录到黑名单
				if threshold.BlockSeconds > 0 {
					// 如果被重复拦截，则加大惩罚力度
					var forbiddenTimes = this.increaseCCCounter(remoteAddr)
					waf.SharedIPBlackList.RecordIP(waf.IPTypeAll, scope, this.ReqServer.Id, remoteAddr, currentTime+int64(threshold.BlockSeconds*forbiddenTimes), 0, scope == firewallconfigs.FirewallScopeGlobal, 0, 0, "CC防护拦截：在"+types.String(threshold.PeriodSeconds)+"秒内请求超过"+types.String(threshold.MaxRequests)+"次")
				}

				this.tags = append(this.tags, "CCProtection")
				this.isAttack = true

				// 关闭连接
				this.writer.Flush()
				this.Close()

				// 关闭同一个IP其他连接
				conns.SharedMap.CloseIPConns(remoteAddr)

				return true
			}
		}
	}

	return
}

// TODO 对forbidden比较多的IP进行惩罚
func (this *HTTPRequest) ccForbid(code int) {
	this.writeCode(http.StatusForbidden, "The request has been blocked by cc policy.", "当前请求已被CC策略拦截。")
}

// 检查跳转次数
func (this *HTTPRequest) checkCCRedirects(httpCCPolicy *nodeconfigs.HTTPCCPolicy, remoteAddr string) bool {
	// 如果无效跳转次数太多，则拦截
	var ccRedirectKey = "HTTP-CC-GET302-" + remoteAddr + "-REDIRECTS"

	var maxRedirectDurationSeconds = 120
	var maxRedirects uint32 = 30
	var blockSeconds int64 = 3600

	if httpCCPolicy != nil && httpCCPolicy.IsOn {
		if httpCCPolicy.RedirectsChecking.DurationSeconds > 0 {
			maxRedirectDurationSeconds = httpCCPolicy.RedirectsChecking.DurationSeconds
		}
		if httpCCPolicy.RedirectsChecking.MaxRedirects > 0 {
			maxRedirects = types.Uint32(httpCCPolicy.RedirectsChecking.MaxRedirects)
		}
		if httpCCPolicy.RedirectsChecking.BlockSeconds > 0 {
			blockSeconds = int64(httpCCPolicy.RedirectsChecking.BlockSeconds)
		}
	}

	var countRedirects = counters.SharedCounter.IncreaseKey(ccRedirectKey, maxRedirectDurationSeconds)
	if countRedirects >= maxRedirects {
		// 加入黑名单
		var scope = firewallconfigs.FirewallScopeGlobal
		if httpCCPolicy != nil {
			scope = httpCCPolicy.FirewallScope()
		}

		waf.SharedIPBlackList.RecordIP(waf.IPTypeAll, scope, this.ReqServer.Id, remoteAddr, fasttime.Now().Unix()+blockSeconds, 0, scope == firewallconfigs.FirewallScopeGlobal, 0, 0, "CC防护拦截：在"+types.String(maxRedirectDurationSeconds)+"秒内无效跳转"+types.String(maxRedirects)+"次")
		this.Close()
		return false
	}

	return true
}

// 对CC禁用次数进行计数
func (this *HTTPRequest) increaseCCCounter(remoteAddr string) int32 {
	counter, ok := ccBlockedMap.Get(remoteAddr)
	if !ok {
		counter = &ccBlockedCounter{
			count:     1,
			updatedAt: fasttime.Now().Unix(),
		}
		ccBlockedMap.Put(remoteAddr, counter)
	} else {
		if counter.updatedAt < fasttime.Now().Unix()-86400 /** 有效期间不要超过1天，防止无限期封禁 **/ {
			counter.count = 0
		}
		counter.updatedAt = fasttime.Now().Unix()
		if counter.count < 32 {
			counter.count++
		}
	}
	return counter.count
}