156 lines
3.2 KiB
Go
156 lines
3.2 KiB
Go
package module
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"git.hlsq.asia/alex/service-common/config"
|
|
"git.hlsq.asia/alex/service-common/cos"
|
|
)
|
|
|
|
var SensitiveWordIns *SensitiveWord
|
|
|
|
type SensitiveWord struct {
|
|
DefaultModule
|
|
cfg *config.SensitiveWordConfig
|
|
}
|
|
|
|
func NewSensitiveWord(cfg *config.SensitiveWordConfig) Module {
|
|
SensitiveWordIns = &SensitiveWord{
|
|
cfg: cfg,
|
|
}
|
|
return SensitiveWordIns
|
|
}
|
|
|
|
func (m *SensitiveWord) Init() error {
|
|
if m.cfg == nil {
|
|
//下载文件
|
|
contentBytes, err := cos.DownloadSimple(context.Background(), "server_file/sensitive_words.txt", "sensitive-word.txt")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var sensitiveWords map[string]struct{}
|
|
err = json.Unmarshal(contentBytes, &sensitiveWords)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
sensitiveWordsList := make([]string, 0, len(sensitiveWords))
|
|
for word := range sensitiveWords {
|
|
sensitiveWordsList = append(sensitiveWordsList, word)
|
|
}
|
|
sensitiveTrie = newTrie(sensitiveWordsList)
|
|
}
|
|
// todo 根据配置文件
|
|
return nil
|
|
}
|
|
|
|
// ContainsSensitiveWords 判断是否包含敏感词
|
|
func (m *SensitiveWord) ContainsSensitiveWords(text string) bool {
|
|
return sensitiveTrie.containsSensitiveWords(text)
|
|
}
|
|
|
|
// FilterSensitiveWords 过滤敏感词
|
|
func (m *SensitiveWord) FilterSensitiveWords(text string) string {
|
|
return sensitiveTrie.filterSensitiveWords(text)
|
|
}
|
|
|
|
var sensitiveTrie *Trie
|
|
|
|
type trieNode struct {
|
|
children map[rune]*trieNode
|
|
isEnd bool
|
|
}
|
|
|
|
type Trie struct {
|
|
root *trieNode
|
|
}
|
|
|
|
func newTrie(worlds []string) *Trie {
|
|
t := &Trie{
|
|
root: &trieNode{
|
|
children: make(map[rune]*trieNode),
|
|
isEnd: false,
|
|
},
|
|
}
|
|
for _, str := range worlds {
|
|
t.insert(str)
|
|
}
|
|
return t
|
|
}
|
|
|
|
const replace = '*'
|
|
|
|
func (sensitiveTrie *Trie) insert(word string) {
|
|
node := sensitiveTrie.root
|
|
for _, c := range word {
|
|
if _, ok := node.children[c]; !ok {
|
|
node.children[c] = &trieNode{
|
|
children: make(map[rune]*trieNode),
|
|
isEnd: false,
|
|
}
|
|
}
|
|
node = node.children[c]
|
|
}
|
|
node.isEnd = true
|
|
}
|
|
|
|
func (sensitiveTrie *Trie) containsSensitiveWords(text string) bool {
|
|
if len(text) == 0 {
|
|
return false
|
|
}
|
|
length := len([]rune(text))
|
|
runes := make([]rune, length)
|
|
index := 0
|
|
for _, char := range text {
|
|
runes[index] = char
|
|
index++
|
|
}
|
|
for pointIndex := 0; pointIndex < length; pointIndex++ {
|
|
//判断
|
|
node := sensitiveTrie.root
|
|
for currentIndex := pointIndex; currentIndex < length; currentIndex++ {
|
|
tn := node.children[runes[currentIndex]]
|
|
if tn == nil {
|
|
break
|
|
}
|
|
if tn.isEnd {
|
|
return true
|
|
}
|
|
node = tn
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (sensitiveTrie *Trie) filterSensitiveWords(text string) string {
|
|
if len(text) == 0 {
|
|
return text
|
|
}
|
|
length := len([]rune(text))
|
|
runes := make([]rune, length)
|
|
index := 0
|
|
for _, char := range text {
|
|
runes[index] = char
|
|
index++
|
|
}
|
|
for pointIndex := 0; pointIndex < length; pointIndex++ {
|
|
//判断
|
|
node := sensitiveTrie.root
|
|
for currentIndex := pointIndex; currentIndex < length; currentIndex++ {
|
|
tn := node.children[runes[currentIndex]]
|
|
if tn == nil {
|
|
break
|
|
}
|
|
if tn.isEnd {
|
|
//替换字符
|
|
for i := pointIndex; i < currentIndex+1; i++ {
|
|
runes[i] = replace
|
|
}
|
|
pointIndex = currentIndex
|
|
break
|
|
}
|
|
node = tn
|
|
}
|
|
}
|
|
return string(runes)
|
|
}
|