add sanitizer for privacy

This commit is contained in:
t 2026-03-09 16:00:30 +08:00
parent 771e53e857
commit 8a219e1873
2 changed files with 533 additions and 0 deletions

307
pkg/utils/sanitizer.go Normal file
View file

@ -0,0 +1,307 @@
// PicoClaw - Ultra-lightweight personal AI assistant
// License: MIT
package utils
import (
"regexp"
"sort"
"strings"
"sync"
)
// SanitizerConfig 脱敏配置
type SanitizerConfig struct {
Enabled bool `json:"enabled"`
Keywords []KeywordRule `json:"keywords"`
CustomPatterns []CustomPatternRule `json:"custom_patterns"`
}
// KeywordRule 关键词规则
type KeywordRule struct {
Word string `json:"word"`
Tag string `json:"tag"`
}
// CustomPatternRule 自定义正则规则
type CustomPatternRule struct {
Name string `json:"name"`
Pattern string `json:"pattern"`
Tag string `json:"tag"`
}
// compiledPattern 编译后的正则规则
type compiledPattern struct {
regex *regexp.Regexp
tag string
name string
}
// Sanitizer 脱敏器
type Sanitizer struct {
config SanitizerConfig
patterns []compiledPattern
mu sync.RWMutex
}
// 全局脱敏器实例
var globalSanitizer *Sanitizer
var globalSanitizerMu sync.RWMutex
// 内置正则规则(按优先级排序)
var builtinPatterns = []compiledPattern{
// 邮箱
{regexp.MustCompile(`[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}`), "[EMAIL]", "email"},
// 中国手机号
{regexp.MustCompile(`\b1[3-9]\d{9}\b`), "[PHONE]", "phone"},
// 中国身份证号18位
{regexp.MustCompile(`\b\d{17}[\dXx]\b`), "[ID_CARD]", "id_card"},
// API KeyOpenAI 格式)
{regexp.MustCompile(`(?i)\bsk-[a-zA-Z0-9]{20,}\b`), "[API_KEY]", "api_key"},
// 信用卡号
{regexp.MustCompile(`\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b`), "[CREDIT_CARD]", "credit_card"},
// IP 地址
{regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`), "[IP]", "ip"},
}
// NewSanitizer 创建脱敏器
func NewSanitizer(config SanitizerConfig) *Sanitizer {
s := &Sanitizer{
config: config,
patterns: make([]compiledPattern, 0),
}
// 添加内置规则
s.patterns = append(s.patterns, builtinPatterns...)
// 添加自定义正则规则
for _, cp := range config.CustomPatterns {
re, err := regexp.Compile(cp.Pattern)
if err != nil {
continue // 忽略无效正则
}
s.patterns = append(s.patterns, compiledPattern{
regex: re,
tag: cp.Tag,
name: cp.Name,
})
}
return s
}
// InitGlobalSanitizer 初始化全局脱敏器
func InitGlobalSanitizer(config SanitizerConfig) {
globalSanitizerMu.Lock()
defer globalSanitizerMu.Unlock()
globalSanitizer = NewSanitizer(config)
}
// GetGlobalSanitizer 获取全局脱敏器
func GetGlobalSanitizer() *Sanitizer {
globalSanitizerMu.RLock()
defer globalSanitizerMu.RUnlock()
return globalSanitizer
}
// SanitizeResult 脱敏结果
type SanitizeResult struct {
Sanitized string // 脱敏后的文本
Mappings map[string]string // 占位符 -> 原始值 的映射
}
// Sanitize 对文本进行脱敏
func (s *Sanitizer) Sanitize(text string) SanitizeResult {
if s == nil || !s.config.Enabled {
return SanitizeResult{
Sanitized: text,
Mappings: make(map[string]string),
}
}
s.mu.RLock()
defer s.mu.RUnlock()
mappings := make(map[string]string)
result := text
// 先处理关键词(精确匹配)
for _, kw := range s.config.Keywords {
if kw.Word == "" {
continue
}
placeholder := s.findUniquePlaceholder(kw.Tag, mappings)
if strings.Contains(result, kw.Word) {
mappings[placeholder] = kw.Word
result = strings.ReplaceAll(result, kw.Word, placeholder)
}
}
// 再处理正则规则
for _, p := range s.patterns {
matches := p.regex.FindAllString(result, -1)
for i, match := range matches {
// 检查是否已经被脱敏(避免重复脱敏)
if s.isAlreadySanitized(match) {
continue
}
placeholder := s.findUniquePlaceholderIndexed(p.tag, i, mappings)
mappings[placeholder] = match
result = strings.Replace(result, match, placeholder, 1)
}
}
return SanitizeResult{
Sanitized: result,
Mappings: mappings,
}
}
// Restore 还原脱敏内容
func (s *Sanitizer) Restore(text string, mappings map[string]string) string {
if s == nil || len(mappings) == 0 {
return text
}
result := text
// 按占位符长度降序排序,避免短占位符先被替换导致长占位符无法匹配
type kv struct {
key string
value string
}
var sorted []kv
for k, v := range mappings {
sorted = append(sorted, kv{k, v})
}
sort.Slice(sorted, func(i, j int) bool {
return len(sorted[i].key) > len(sorted[j].key)
})
for _, kv := range sorted {
result = strings.ReplaceAll(result, kv.key, kv.value)
}
return result
}
// findUniquePlaceholder 找到唯一的占位符
func (s *Sanitizer) findUniquePlaceholder(tag string, existing map[string]string) string {
base := tag
if base == "" {
base = "[REDACTED]"
}
// 如果 tag 已经是完整格式(如 [EMAIL]),直接使用
if strings.HasPrefix(base, "[") && strings.HasSuffix(base, "]") {
placeholder := base
// 检查是否已存在
for i := 0; ; i++ {
if i == 0 {
if _, exists := existing[placeholder]; !exists {
return placeholder
}
} else {
numbered := placeholder[:len(placeholder)-1] + "_" + intToString(i) + "]"
if _, exists := existing[numbered]; !exists {
return numbered
}
}
}
}
// 否则添加方括号
placeholder := "[" + base + "]"
for i := 0; ; i++ {
if i == 0 {
if _, exists := existing[placeholder]; !exists {
return placeholder
}
} else {
numbered := "[" + base + "_" + intToString(i) + "]"
if _, exists := existing[numbered]; !exists {
return numbered
}
}
}
}
// findUniquePlaceholderIndexed 找到带索引的唯一占位符
func (s *Sanitizer) findUniquePlaceholderIndexed(tag string, idx int, existing map[string]string) string {
base := tag
if base == "" {
base = "[REDACTED]"
}
// 确保 tag 有方括号
if !strings.HasPrefix(base, "[") {
base = "[" + base + "]"
}
// 如果是第一个匹配,使用基础格式
if idx == 0 {
if _, exists := existing[base]; !exists {
return base
}
}
// 否则使用带索引的格式
for i := idx; ; i++ {
numbered := base[:len(base)-1] + "_" + intToString(i) + "]"
if _, exists := existing[numbered]; !exists {
return numbered
}
}
}
// isAlreadySanitized 检查文本是否已被脱敏
func (s *Sanitizer) isAlreadySanitized(text string) bool {
// 检查是否是占位符格式
return strings.HasPrefix(text, "[") && strings.HasSuffix(text, "]")
}
// intToString 简单的整数转字符串
func intToString(n int) string {
if n == 0 {
return "0"
}
var negative bool
if n < 0 {
negative = true
n = -n
}
var digits []byte
for n > 0 {
digits = append([]byte{byte('0' + n%10)}, digits...)
n /= 10
}
if negative {
digits = append([]byte{'-'}, digits...)
}
return string(digits)
}
// SanitizeText 全局脱敏函数
func SanitizeText(text string) SanitizeResult {
s := GetGlobalSanitizer()
if s == nil {
return SanitizeResult{
Sanitized: text,
Mappings: make(map[string]string),
}
}
return s.Sanitize(text)
}
// RestoreText 全局还原函数
func RestoreText(text string, mappings map[string]string) string {
s := GetGlobalSanitizer()
if s == nil {
return text
}
return s.Restore(text, mappings)
}

226
pkg/utils/sanitizer_test.go Normal file
View file

@ -0,0 +1,226 @@
// PicoClaw - Ultra-lightweight personal AI assistant
// License: MIT
package utils
import (
"testing"
)
func TestSanitizerBasic(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Test email sanitization
text := "My email is test@example.com"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("Email should be sanitized, got: %s", result.Sanitized)
}
if _, ok := result.Mappings["[EMAIL]"]; !ok {
t.Errorf("Should have [EMAIL] mapping, got: %v", result.Mappings)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerPhone(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Test Chinese phone number
text := "Call me at 13812345678"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("Phone should be sanitized, got: %s", result.Sanitized)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerIDCard(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Test Chinese ID card
text := "My ID is 110101199003077654"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("ID card should be sanitized, got: %s", result.Sanitized)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerAPIKey(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Test API key
text := "API key: sk-abcdefghijklmnopqrstuvwxyz123456"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("API key should be sanitized, got: %s", result.Sanitized)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerMultiple(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Test multiple sensitive data
text := "Email: test@example.com, Phone: 13812345678, ID: 110101199003077654"
result := s.Sanitize(text)
// Should have 3 mappings
if len(result.Mappings) < 3 {
t.Errorf("Should have at least 3 mappings, got: %d", len(result.Mappings))
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerKeywords(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
Keywords: []KeywordRule{
{Word: "secret_password", Tag: "[PASSWORD]"},
},
}
s := NewSanitizer(config)
text := "My password is secret_password"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("Keyword should be sanitized, got: %s", result.Sanitized)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerCustomPattern(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
CustomPatterns: []CustomPatternRule{
{Name: "employee_id", Pattern: `EMP\d{6}`, Tag: "[EMP_ID]"},
},
}
s := NewSanitizer(config)
text := "Employee ID: EMP123456"
result := s.Sanitize(text)
if result.Sanitized == text {
t.Errorf("Custom pattern should be sanitized, got: %s", result.Sanitized)
}
// Test restore
restored := s.Restore(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Restore failed, expected: %s, got: %s", text, restored)
}
}
func TestSanitizerDisabled(t *testing.T) {
config := SanitizerConfig{
Enabled: false,
}
s := NewSanitizer(config)
text := "My email is test@example.com"
result := s.Sanitize(text)
if result.Sanitized != text {
t.Errorf("When disabled, text should not be modified, got: %s", result.Sanitized)
}
}
func TestSanitizerNil(t *testing.T) {
var s *Sanitizer
text := "My email is test@example.com"
result := s.Sanitize(text)
if result.Sanitized != text {
t.Errorf("Nil sanitizer should not modify text, got: %s", result.Sanitized)
}
restored := s.Restore(text, map[string]string{"[EMAIL]": "test@example.com"})
if restored != text {
t.Errorf("Nil sanitizer restore should not modify text, got: %s", restored)
}
}
func TestSanitizerNoDoubleSanitize(t *testing.T) {
config := SanitizerConfig{
Enabled: true,
}
s := NewSanitizer(config)
// Already sanitized text should not be double-sanitized
text := "My email is [EMAIL]"
result := s.Sanitize(text)
if result.Sanitized != text {
t.Errorf("Already sanitized text should not be modified, got: %s", result.Sanitized)
}
}
func TestGlobalSanitizer(t *testing.T) {
// Test global sanitizer functions
InitGlobalSanitizer(SanitizerConfig{
Enabled: true,
})
text := "Email: test@example.com"
result := SanitizeText(text)
if result.Sanitized == text {
t.Errorf("Global sanitize should work, got: %s", result.Sanitized)
}
restored := RestoreText(result.Sanitized, result.Mappings)
if restored != text {
t.Errorf("Global restore should work, got: %s", restored)
}
}