package main
import (
_ "embed"
"encoding/json"
"fmt"
"github.com/Jinnrry/pmail/dto/parsemail"
"github.com/Jinnrry/pmail/hooks/framework"
"github.com/Jinnrry/pmail/hooks/spam_block/tools"
"github.com/Jinnrry/pmail/models"
"github.com/Jinnrry/pmail/utils/context"
log "github.com/sirupsen/logrus"
"github.com/spf13/cast"
"io"
"net/http"
"os"
"strings"
"time"
)
type SpamBlock struct {
cfg SpamBlockConfig
hc *http.Client
}
func (s *SpamBlock) SendBefore(ctx *context.Context, email *parsemail.Email) {
}
func (s *SpamBlock) SendAfter(ctx *context.Context, email *parsemail.Email, err map[string]error) {
}
func (s *SpamBlock) ReceiveParseBefore(ctx *context.Context, email *[]byte) {
}
// GetName 获取插件名称
func (s *SpamBlock) GetName(ctx *context.Context) string {
return "SpamBlock"
}
//go:embed static/index.html
var index string
//go:embed static/jquery.js
var jquery string
// SettingsHtml 插件页面
func (s *SpamBlock) SettingsHtml(ctx *context.Context, url string, requestData string) string {
if strings.Contains(url, "jquery.js") {
return jquery
}
if strings.Contains(url, "index.html") {
if !ctx.IsAdmin {
return fmt.Sprintf(`
Please contact the administrator for configuration.
`)
}
return fmt.Sprintf(index, s.cfg.ApiURL, s.cfg.ApiTimeout, s.cfg.Threshold)
}
var cfg SpamBlockConfig
var tempCfg map[string]string
err := json.Unmarshal([]byte(requestData), &tempCfg)
if err != nil {
return err.Error()
}
cfg.ApiURL = tempCfg["url"]
cfg.Threshold = cast.ToFloat64(tempCfg["threshold"])
cfg.ApiTimeout = cast.ToInt(tempCfg["timeout"])
err = saveConfig(cfg)
if err != nil {
return err.Error()
}
s.cfg = cfg
return "success"
}
type ModelResponse struct {
Predictions [][]float64 `json:"predictions"`
}
type ApiRequest struct {
Instances []InstanceItem `json:"instances"`
}
type InstanceItem struct {
Token []string `json:"token"`
}
func (s *SpamBlock) ReceiveParseAfter(ctx *context.Context, email *parsemail.Email) {
if s.cfg.ApiURL == "" {
return
}
content := tools.Trim(tools.TrimHtml(string(email.HTML)))
if content == "" {
content = tools.Trim(string(email.Text))
}
reqData := ApiRequest{
Instances: []InstanceItem{
{
Token: []string{
fmt.Sprintf("%s %s", email.Subject, content),
},
},
},
}
str, _ := json.Marshal(reqData)
resp, err := s.hc.Post(s.cfg.ApiURL, "application/json", strings.NewReader(string(str)))
if err != nil {
log.Errorf("API Error: %v", err)
return
}
body, _ := io.ReadAll(resp.Body)
modelResponse := ModelResponse{}
err = json.Unmarshal(body, &modelResponse)
if err != nil {
log.WithContext(ctx).Errorf("API Error: %v", err)
return
}
if len(modelResponse.Predictions) == 0 {
log.WithContext(ctx).Errorf("API Response Error: %v", string(body))
return
}
classes := modelResponse.Predictions[0]
if len(classes) != 3 {
return
}
var maxScore float64
var maxClass int
for i, score := range classes {
if score > maxScore {
maxScore = score
maxClass = i
}
}
switch maxClass {
case 0:
log.WithContext(ctx).Infof("[Spam Check Result: %f Normal] %s", maxScore, email.Subject)
case 1:
log.WithContext(ctx).Infof("[Spam Check Result: %f Spam ] %s", maxScore, email.Subject)
case 2:
log.WithContext(ctx).Infof("[Spam Check Result: %f Blackmail ] %s", maxScore, email.Subject)
}
if maxClass != 0 && maxScore > s.cfg.Threshold/100 {
email.Status = 5
}
}
func (s *SpamBlock) ReceiveSaveAfter(ctx *context.Context, email *parsemail.Email, ue []*models.UserEmail) {
}
type SpamBlockConfig struct {
ApiURL string `json:"apiURL"`
ApiTimeout int `json:"apiTimeout"` // 单位毫秒
Threshold float64 `json:"threshold"`
}
func NewSpamBlockHook() *SpamBlock {
var pluginConfig SpamBlockConfig
if _, err := os.Stat("./plugins/spam_block_config.json"); err == nil {
cfgData, err := os.ReadFile("./plugins/spam_block_config.json")
if err == nil {
json.Unmarshal(cfgData, &pluginConfig)
}
}
log.Infof("Config: %+v", pluginConfig)
if pluginConfig.ApiTimeout == 0 {
pluginConfig.ApiTimeout = 3000
}
if pluginConfig.Threshold == 0 {
pluginConfig.Threshold = 20
}
hc := &http.Client{
Timeout: time.Duration(pluginConfig.ApiTimeout) * time.Millisecond,
}
return &SpamBlock{
cfg: pluginConfig,
hc: hc,
}
}
func saveConfig(cfg SpamBlockConfig) error {
data, _ := json.Marshal(cfg)
err := os.WriteFile("./plugins/spam_block_config.json", data, 0777)
return err
}
func main() {
log.Infof("SpamBlockPlug Star Success")
instance := NewSpamBlockHook()
if instance == nil {
return
}
framework.CreatePlugin("spam_block", instance).Run()
}