SafeLine/management/webserver/api/policygroup.go

package api

import (
	"encoding/json"
	"net/http"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"

	"chaitin.cn/dev/go/errors"

	"chaitin.cn/patronus/safeline-2/management/webserver/api/response"
	"chaitin.cn/patronus/safeline-2/management/webserver/model"
	"chaitin.cn/patronus/safeline-2/management/webserver/pkg/constants"
	"chaitin.cn/patronus/safeline-2/management/webserver/pkg/database"
	"chaitin.cn/patronus/safeline-2/management/webserver/pkg/fvm"
)

func PutPolicyGroupGlobal(ctx *gin.Context) {
	var params model.PolicyGroup
	if err := ctx.BindJSON(&params); err != nil {
		logger.Error(err)
		response.Error(ctx, response.ErrorParamNotOK, http.StatusInternalServerError)
		return
	}

	db := database.GetDB()
	err := db.Transaction(func(tx *gorm.DB) error {
		var pggOption model.Options
		res := tx.Where(&model.Options{Key: constants.PolicyGroupGlobal}).First(&pggOption)
		if res.Error != nil {
			return res.Error
		}
		if res.RowsAffected == 0 {
			return errors.New("Data queried does not exist")
		}

		pggStr, err := json.Marshal(params)
		if err != nil {
			logger.Error(err)
			response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
			return err
		}

		pggOption.Value = string(pggStr)
		tx.Save(&pggOption)

		if err := fvm.PushFSL(tx); err != nil {
			return errors.New("Rules compile error, please check your params.")
		}

		return nil
	})
	if err != nil {
		logger.Error(err)
		response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
		return
	}

	response.Success(ctx, nil)
}

func GetPolicyGroupGlobal(ctx *gin.Context) {
	var pggOption model.Options
	database.GetDB().Where(&model.Options{Key: constants.PolicyGroupGlobal}).First(&pggOption)

	var pgg model.PolicyGroup
	err := json.Unmarshal([]byte(pggOption.Value), &pgg)
	if err != nil {
		logger.Error(err)
		response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
		return
	}

	response.Success(ctx, gin.H{"data": pgg})
}

func PutSrcIPConfig(ctx *gin.Context) {
	var params model.SrcIPConfig
	if err := ctx.BindJSON(&params); err != nil {
		logger.Error(err)
		response.Error(ctx, response.ErrorParamNotOK, http.StatusInternalServerError)
		return
	}

	db := database.GetDB()
	err := db.Transaction(func(tx *gorm.DB) error {
		var scOption model.Options
		res := tx.Where(&model.Options{Key: constants.SrcIPConfig}).First(&scOption)
		if res.Error != nil {
			return res.Error
		}
		if res.RowsAffected == 0 {
			return errors.New("Data queried does not exist")
		}

		scStr, err := json.Marshal(params)
		if err != nil {
			logger.Error(err)
			response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
			return err
		}

		scOption.Value = string(scStr)
		tx.Save(&scOption)

		if err := fvm.PushFSL(tx); err != nil {
			return errors.New("Rules compile error, please check your params.")
		}

		return nil
	})
	if err != nil {
		logger.Error(err)
		response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
		return
	}

	response.Success(ctx, nil)
}

func GetSrcIPConfig(ctx *gin.Context) {
	srcIPConfig, err := model.GetSrcIPConfig(database.GetDB().DB)
	if err != nil {
		return
	}
	response.Success(ctx, gin.H{"data": srcIPConfig})
}

type PolicyGroupPresetRequest struct {
	Preset string `json:"preset" binding:"required,oneof=strict default disable"`
}

func PostPolicyGroupPreset(ctx *gin.Context) {
	var params PolicyGroupPresetRequest
	if err := ctx.BindJSON(&params); err != nil {
		logger.Error(err)
		response.Error(ctx, response.ErrorParamNotOK, http.StatusInternalServerError)
		return
	}

	// All 17 detection modules
	modules := []string{
		"m_asp_code_injection", "m_cmd_injection", "m_csrf",
		"m_file_include", "m_file_upload", "m_http",
		"m_java", "m_java_unserialize", "m_php_code_injection",
		"m_php_unserialize", "m_response", "m_rule",
		"m_scanner", "m_sqli", "m_ssrf", "m_ssti", "m_xss",
	}

	pgg := make(model.PolicyGroup)
	for _, m := range modules {
		pgg[m] = params.Preset
	}

	db := database.GetDB()
	err := db.Transaction(func(tx *gorm.DB) error {
		var pggOption model.Options
		res := tx.Where(&model.Options{Key: constants.PolicyGroupGlobal}).First(&pggOption)
		if res.Error != nil {
			return res.Error
		}

		pggStr, err := json.Marshal(pgg)
		if err != nil {
			return err
		}

		pggOption.Value = string(pggStr)
		tx.Save(&pggOption)

		if err := fvm.PushFSL(tx); err != nil {
			return errors.New("Rules compile error")
		}

		return nil
	})
	if err != nil {
		logger.Error(err)
		response.Error(ctx, response.JSONBody{Err: response.ErrInternalError, Msg: err.Error()}, http.StatusInternalServerError)
		return
	}

	response.Success(ctx, gin.H{"preset": params.Preset, "data": pgg})
}