// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build plus

package nodes

import (
	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
	"github.com/TeaOSLab/EdgeNode/internal/events"
	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
	"github.com/TeaOSLab/EdgeNode/internal/uam"
	"github.com/TeaOSLab/EdgeNode/internal/utils/agents"
	"github.com/TeaOSLab/EdgeNode/internal/utils/counters"
	"github.com/TeaOSLab/EdgeNode/internal/utils/ttlcache"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/types"
	"net/http"
	"strings"
)

var sharedUAMManager *uam.Manager

func init() {
	if !teaconst.IsMain {
		return
	}

	events.On(events.EventLoaded, func() {
		if sharedUAMManager != nil {
			return
		}
		manager, _ := uam.NewManager(sharedNodeConfig.NodeId, sharedNodeConfig.Secret)
		if manager != nil {
			sharedUAMManager = manager
		}
	})

	events.On(events.EventReload, func() {
		if sharedUAMManager != nil {
			return
		}
		manager, _ := uam.NewManager(sharedNodeConfig.NodeId, sharedNodeConfig.Secret)
		if manager != nil {
			sharedUAMManager = manager
		}
	})
}

func (this *HTTPRequest) isUAMRequest() bool {
	if this.web.UAM == nil || !this.web.UAM.IsOn || this.RawReq.Method != http.MethodPost {
		return false
	}

	var cookiesString = this.RawReq.Header.Get("Cookie")
	if len(cookiesString) == 0 {
		return false
	}

	return strings.HasPrefix(cookiesString, uam.CookiePrevKey+"=") ||
		strings.Contains(cookiesString, " "+uam.CookiePrevKey+"=")
}

// UAM
// TODO 需要检查是否为plus
func (this *HTTPRequest) doUAM() (block bool) {
	var serverId int64
	if this.ReqServer != nil {
		serverId = this.ReqServer.Id
	}

	var uamConfig = this.web.UAM
	if uamConfig == nil ||
		!uamConfig.IsOn ||
		!uamConfig.MatchURL(this.requestScheme()+"://"+this.ReqHost+this.Path()) ||
		!uamConfig.MatchRequest(this.Format) {
		return
	}

	var policy = this.nodeConfig.FindUAMPolicyWithClusterId(this.ReqServer.ClusterId)
	if policy == nil {
		policy = nodeconfigs.NewUAMPolicy()
	}
	if policy == nil || !policy.IsOn {
		return
	}

	// 获取UAM管理器
	var manager = sharedUAMManager
	if manager == nil {
		return false
	}

	// 忽略URL白名单
	if this.RawReq.URL.Path == "/favicon.ico" || this.RawReq.URL.Path == "/favicon.png" {
		return false
	}

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

	// 检查UAM白名单
	if uamConfig.AddToWhiteList && ttlcache.SharedInt64Cache.Read("UAM:WHITE:"+remoteAddr) != nil {
		return false
	}

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

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

	// 如果是搜索引擎直接通过
	var userAgent = this.RawReq.UserAgent()
	if len(userAgent) == 0 {
		// 如果User-Agent为空，则直接阻止
		this.writer.WriteHeader(http.StatusForbidden)

		// 增加失败次数
		if manager.IncreaseFails(policy, remoteAddr, serverId) {
			this.isAttack = true
		}

		return true
	}

	// 不管是否开启允许搜索引擎，这里都放行，避免收录了拦截的代码
	if agents.SharedManager.ContainsIP(remoteAddr) {
		return false
	}
	if policy.AllowSearchEngines {
		if searchEngineRegex.MatchString(userAgent) {
			return false
		}
	}

	// 如果是python之类的直接拦截
	if policy.DenySpiders && spiderRegexp.MatchString(userAgent) {
		this.writer.WriteHeader(http.StatusForbidden)

		// 增加失败次数
		if manager.IncreaseFails(policy, remoteAddr, serverId) {
			this.isAttack = true
		}

		return true
	}

	// 检查预生成Key
	var step = this.Header().Get("X-GE-UA-Step")
	if step == uam.StepPrev {
		if this.Method() != http.MethodPost {
			this.writer.WriteHeader(http.StatusForbidden)
			return true
		}
		if manager.CheckPrevKey(policy, uamConfig, this.RawReq, remoteAddr, this.writer) {
			_, _ = this.writer.Write([]byte(`{"ok": true}`))
		} else {
			_, _ = this.writer.Write([]byte(`{"ok": false}`))
		}

		// 增加失败次数
		if manager.IncreaseFails(policy, remoteAddr, serverId) {
			this.isAttack = true
		}

		return true
	}

	// 检查是否启用QPS
	if uamConfig.MinQPSPerIP > 0 && len(step) == 0 {
		var avgQPS = counters.SharedCounter.IncreaseKey("UAM:"+remoteAddr, 60) / 60
		if avgQPS <= 0 {
			avgQPS = 1
		}
		if avgQPS < types.Uint32(uamConfig.MinQPSPerIP) {
			return false
		}
	}

	// 检查Cookie中的Key
	isVerified, isAttack, _ := manager.CheckKey(policy, this.RawReq, this.writer, remoteAddr, serverId, uamConfig.KeyLife)
	if isVerified {
		return false
	}
	this.isAttack = isAttack

	// 检查是否已生成有效的prev key，如果已经生成，则表示当前请求是附带请求（比如favicon.ico），不再重新生成新的
	// TODO 考虑这里的必要性
	if manager.ExistsActivePreKey(this.RawReq) {
		return true
	}

	// 显示加载页面
	err := manager.LoadPage(policy, this.RawReq, this.Format, remoteAddr, this.writer)
	if err != nil {
		return false
	}

	return true
}