From 8a219e1873816cf127aa95e6bc4c918f66dc697c Mon Sep 17 00:00:00 2001 From: t <452214834@qq.com> Date: Mon, 9 Mar 2026 16:00:30 +0800 Subject: [PATCH] add sanitizer for privacy --- pkg/utils/sanitizer.go | 307 ++++++++++++++++++++++++++++++++++++ pkg/utils/sanitizer_test.go | 226 ++++++++++++++++++++++++++ 2 files changed, 533 insertions(+) create mode 100644 pkg/utils/sanitizer.go create mode 100644 pkg/utils/sanitizer_test.go diff --git a/pkg/utils/sanitizer.go b/pkg/utils/sanitizer.go new file mode 100644 index 000000000..c6e760754 --- /dev/null +++ b/pkg/utils/sanitizer.go @@ -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 Key（OpenAI 格式） + {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) +} \ No newline at end of file diff --git a/pkg/utils/sanitizer_test.go b/pkg/utils/sanitizer_test.go new file mode 100644 index 000000000..bf2da44da --- /dev/null +++ b/pkg/utils/sanitizer_test.go @@ -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) + } +} \ No newline at end of file