waf-platform/EdgeAPI/internal/db/models/user_script_dao_plus.go

//go:build plus

package models

import (
	"encoding/json"
	"errors"
	"github.com/TeaOSLab/EdgeAPI/internal/utils"
	_ "github.com/go-sql-driver/mysql"
	"github.com/iwind/TeaGo/dbs"
	"github.com/iwind/TeaGo/lists"
	"github.com/iwind/TeaGo/types"
	stringutil "github.com/iwind/TeaGo/utils/string"
	"time"
)

// EnableUserScript 启用条目
func (this *UserScriptDAO) EnableUserScript(tx *dbs.Tx, id int64) error {
	_, err := this.Query(tx).
		Pk(id).
		Set("state", UserScriptStateEnabled).
		Update()
	return err
}

// DisableUserScript 禁用条目
func (this *UserScriptDAO) DisableUserScript(tx *dbs.Tx, id int64) error {
	_, err := this.Query(tx).
		Pk(id).
		Set("state", UserScriptStateDisabled).
		Update()
	return err
}

// FindEnabledUserScript 查找启用中的条目
func (this *UserScriptDAO) FindEnabledUserScript(tx *dbs.Tx, id int64) (*UserScript, error) {
	result, err := this.Query(tx).
		Pk(id).
		State(UserScriptStateEnabled).
		Find()
	if result == nil {
		return nil, err
	}
	return result.(*UserScript), err
}

// CreateUserScript 创建脚本
func (this *UserScriptDAO) CreateUserScript(tx *dbs.Tx, userId int64, code string) (int64, error) {
	var op = NewUserScriptOperator()
	op.UserId = userId
	op.Code = code
	op.CodeMD5 = stringutil.Md5(code)
	op.IsRejected = false
	op.IsPassed = false
	op.State = UserScriptStateEnabled
	return this.SaveInt64(tx, op)
}

// ExistsPassedCodeMD5 根据代码MD5检查脚本是否已经存在并通过审核
func (this *UserScriptDAO) ExistsPassedCodeMD5(tx *dbs.Tx, codeMD5 string) (bool, error) {
	return this.Query(tx).
		State(UserScriptStateEnabled).
		Attr("isPassed", true).
		Attr("codeMD5", codeMD5).
		Exist()
}

// FindUserScriptIdWithCodeMD5 根据代码MD5查找用户脚本ID
func (this *UserScriptDAO) FindUserScriptIdWithCodeMD5(tx *dbs.Tx, userId int64, codeMD5 string) (int64, error) {
	if userId <= 0 {
		return 0, errors.New("invalid userId '" + types.String(userId) + "'")
	}
	return this.Query(tx).
		State(UserScriptStateEnabled).
		Attr("userId", userId).
		Attr("codeMD5", codeMD5).
		ResultPk().
		FindInt64Col(0)
}

// FindUserScriptWithCodeMD5 根据代码MD5查找代码
func (this *UserScriptDAO) FindUserScriptWithCodeMD5(tx *dbs.Tx, codeMD5 string) (*UserScript, error) {
	if len(codeMD5) != 32 {
		return nil, nil
	}
	one, err := this.Query(tx).
		Attr("codeMD5", codeMD5).
		State(UserScriptStateEnabled).
		Find()
	if err != nil || one == nil {
		return nil, err
	}
	return one.(*UserScript), nil
}

// CountUserScripts 查询脚本数量
func (this *UserScriptDAO) CountUserScripts(tx *dbs.Tx, userId int64, isAuditing bool) (count int64, err error) {
	var query = this.Query(tx)
	if userId > 0 {
		query.Attr("userId", userId)
	}
	if isAuditing {
		query.Where("(isRejected=0 AND isPassed=0)")
	} else {
		query.Where("(isRejected=1 OR isPassed=1)")
	}
	query.State(UserScriptStateEnabled)
	return query.Count()
}

// ListUserScripts 读取单页脚本
func (this *UserScriptDAO) ListUserScripts(tx *dbs.Tx, userId int64, isAuditing bool, offset int64, size int64) (result []*UserScript, err error) {
	var query = this.Query(tx).
		DescPk().
		Offset(offset).
		Limit(size)
	if userId > 0 {
		query.Attr("userId", userId)
	}
	if isAuditing {
		query.Where("(isRejected=0 AND isPassed=0)")
	} else {
		query.Where("(isRejected=1 OR isPassed=1)")
	}
	query.State(UserScriptStateEnabled)
	query.Slice(&result)
	_, err = query.FindAll()
	return
}

// RejectUserScript 驳回
func (this *UserScriptDAO) RejectUserScript(tx *dbs.Tx, scriptId int64, adminId int64, reason string) error {
	if scriptId <= 0 {
		return errors.New("invalid 'scriptId'")
	}
	var op = NewUserScriptOperator()
	op.Id = scriptId
	op.AdminId = adminId

	op.IsPassed = 0

	op.IsRejected = 1
	op.RejectedAt = time.Now().Unix()
	op.RejectedReason = utils.LimitString(reason, 100) // prevent too long reason for table column
	return this.Save(tx, op)
}

// PassUserScript 通过脚本
func (this *UserScriptDAO) PassUserScript(tx *dbs.Tx, scriptId int64, adminId int64) error {
	if scriptId <= 0 {
		return errors.New("invalid 'scriptId'")
	}
	var op = NewUserScriptOperator()
	op.Id = scriptId
	op.AdminId = adminId

	op.IsRejected = 0

	op.IsPassed = 1
	op.PassedAt = time.Now().Unix()

	err := this.Save(tx, op)
	if err != nil {
		return err
	}

	// 通知审核结果
	return this.NotifyUserScriptPassed(tx, scriptId)
}

// AddWebIdToUserScript 为脚本添加关联的WebId
func (this *UserScriptDAO) AddWebIdToUserScript(tx *dbs.Tx, scriptId int64, webId int64) error {
	if scriptId <= 0 {
		return errors.New("invalid 'scriptId'")
	}
	if webId <= 0 {
		return errors.New("invalid 'webId'")
	}

	webIdsJSON, err := this.Query(tx).
		Pk(scriptId).
		Result("webIds").
		FindJSONCol()
	if err != nil {
		return err
	}
	var webIds = []int64{}
	if len(webIdsJSON) > 0 {
		err = json.Unmarshal(webIdsJSON, &webIds)
		if err != nil {
			return err
		}
		if !lists.ContainsInt64(webIds, webId) {
			webIds = append(webIds, webId)
		}
	} else {
		webIds = append(webIds, webId)
	}
	webIdsJSON, err = json.Marshal(webIds)
	if err != nil {
		return err
	}
	return this.Query(tx).
		Pk(scriptId).
		Set("webIds", webIdsJSON).
		UpdateQuickly()
}

// NotifyUserScriptPassed 通知审核通过结果
func (this *UserScriptDAO) NotifyUserScriptPassed(tx *dbs.Tx, scriptId int64) error {
	if scriptId <= 0 {
		return nil
	}
	one, err := this.Query(tx).
		Pk(scriptId).
		Result("codeMD5", "webIds").
		Find()
	if err != nil || one == nil {
		return err
	}
	var userScript = one.(*UserScript)
	var webIds []int64
	if len(userScript.WebIds) > 0 {
		err = json.Unmarshal(userScript.WebIds, &webIds)
		if err != nil {
			return err
		}
	}

	for _, webId := range webIds {
		err = SharedHTTPWebDAO.UpdateWebRequestScriptsAsPassed(tx, webId, userScript.CodeMD5)
		if err != nil {
			return err
		}
	}

	return nil
}