FranklinWH-Toolbox/service/cmd/mcp/tools/helper/pii_test.go

package helper

import (
	"encoding/json"
	"strings"
	"testing"

	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
)

// ---------- 掩码函数单元测试 ----------

func TestMaskName(t *testing.T) {
	cases := map[string]string{
		"":      "",
		"张三":    "张***",
		"John":  "J***",
		"李":     "李***",
	}
	for in, want := range cases {
		if got := maskName(in); got != want {
			t.Errorf("maskName(%q)=%q, want %q", in, got, want)
		}
	}
}

func TestMaskEmail(t *testing.T) {
	if got := maskEmail("john@corp.com"); got != "j***@corp.com" {
		t.Errorf("maskEmail got %q", got)
	}
	// 非法邮箱退化为完全脱敏
	if got := maskEmail("notanemail"); got != maskFull() {
		t.Errorf("maskEmail(invalid) got %q", got)
	}
	if got := maskEmail(""); got != "" {
		t.Errorf("maskEmail(empty) got %q", got)
	}
}

func TestMaskAccount(t *testing.T) {
	if got := maskAccount("abcdefgh"); got != "ab***gh" {
		t.Errorf("maskAccount got %q", got)
	}
	// 过短完全脱敏
	if got := maskAccount("abc"); got != maskFull() {
		t.Errorf("maskAccount(short) got %q", got)
	}
}

func TestMaskFullFixedLength(t *testing.T) {
	if got := maskFull(); got != strings.Repeat(MaskChar, fullMaskLen) {
		t.Errorf("maskFull got %q", got)
	}
}

// ---------- 优化 B：流式 token 脱敏测试 ----------

func TestMaskJSONBytes_FlatFields(t *testing.T) {
	in := []byte(`{
		"sn": "SN123",
		"userName": "Alice",
		"contactNumber": "13812345678",
		"email": "a@b.com",
		"longitude": 121.5,
		"latitude": 31.2,
		"address": "123 Main St",
		"city": "Shanghai",
		"province": "SH",
		"country": "CN",
		"iccid": "8986012345678901234",
		"imei": "356938035643809",
		"account": "installer01"
	}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("MaskJSONBytes err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("output not valid JSON: %v\n%s", err, out)
	}

	// 保留字段
	if m["sn"] != "SN123" {
		t.Errorf("sn should be preserved, got %v", m["sn"])
	}
	if m["iccid"] != "8986012345678901234" {
		t.Errorf("iccid should be preserved, got %v", m["iccid"])
	}
	if m["city"] != "Shanghai" || m["province"] != "SH" || m["country"] != "CN" {
		t.Errorf("admin region should be preserved: city=%v province=%v country=%v", m["city"], m["province"], m["country"])
	}

	// 完全脱敏字段
	full := maskFull()
	for _, k := range []string{"contactNumber", "longitude", "latitude", "address", "imei"} {
		if m[k] != full {
			t.Errorf("%s should be fully masked, got %v", k, m[k])
		}
	}
	// 姓名
	if m["userName"] != "A***" {
		t.Errorf("userName masked wrong: %v", m["userName"])
	}
	// 邮箱
	if m["email"] != "a***@b.com" {
		t.Errorf("email masked wrong: %v", m["email"])
	}
	// 账号
	if m["account"] != "in***01" {
		t.Errorf("account masked wrong: %v", m["account"])
	}
}

func TestMaskJSONBytes_Nested(t *testing.T) {
	in := []byte(`{
		"data": [
			{"userName": "Bob", "sn": "S1", "nested": {"phone": "111", "keep": "ok"}},
			{"userName": "Carol", "sn": "S2"}
		],
		"total": 2
	}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("invalid JSON: %v\n%s", err, out)
	}
	arr := m["data"].([]any)
	first := arr[0].(map[string]any)
	if first["userName"] != "B***" {
		t.Errorf("nested list userName wrong: %v", first["userName"])
	}
	if first["sn"] != "S1" {
		t.Errorf("sn should be preserved in list: %v", first["sn"])
	}
	nested := first["nested"].(map[string]any)
	if nested["phone"] != maskFull() {
		t.Errorf("deeply nested phone should be masked: %v", nested["phone"])
	}
	if nested["keep"] != "ok" {
		t.Errorf("non-sensitive nested field should be preserved: %v", nested["keep"])
	}
	if m["total"].(float64) != 2 {
		t.Errorf("total wrong: %v", m["total"])
	}
}

func TestMaskJSONBytes_SensitiveSubtreeRedacted(t *testing.T) {
	// 敏感 key 对应对象/数组时，整个子树被脱敏
	in := []byte(`{"address": {"street": "secret", "no": 42}, "list_phone": ["a","b"]}`)
	// 注意 list_phone 不在规则表，仅测 address 子树
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("invalid JSON: %v\n%s", err, out)
	}
	addr := m["address"].(map[string]any)
	if addr["street"] != maskFull() {
		t.Errorf("address.street should be masked: %v", addr["street"])
	}
	// number 子字段被替换为掩码字符串
	if addr["no"] != maskFull() {
		t.Errorf("address.no should be masked: %v", addr["no"])
	}
}

func TestMaskJSONBytes_InvalidJSONFailsClosed(t *testing.T) {
	_, err := MaskJSONBytes([]byte(`{not valid json`))
	if err == nil {
		t.Fatal("expected error for invalid JSON (fail-closed)")
	}
}

func TestMaskJSONBytes_NullPreserved(t *testing.T) {
	in := []byte(`{"userName": null, "phone": null}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !strings.Contains(string(out), `"userName":null`) {
		t.Errorf("null should be preserved: %s", out)
	}
}

// ---------- 优化 A：强类型脱敏测试 ----------

func TestSanitizeSiteInfoDTO_NoUpstreamPollution(t *testing.T) {
	name := "Alice"
	phone := "13812345678"
	addr := "123 Main St"
	iccid := "8986012345678901234"
	orig := &dto.SiteInfoDTO{
		Sn:            "SN1",
		UserName:      &name,
		ContactNumber: &phone,
		Address:       &addr,
		Longitude:     "121.5",
		Latitude:      "31.2",
		City:          "Shanghai",
		ICCID:         iccid,
		IMEI:          "356938035643809",
	}

	sanitized, ok := sanitizeStrong(orig)
	if !ok {
		t.Fatal("expected strong-type hit for *dto.SiteInfoDTO")
	}
	s := sanitized.(*dto.SiteInfoDTO)

	// 脱敏结果正确
	if *s.UserName != "A***" {
		t.Errorf("userName: %v", *s.UserName)
	}
	if *s.ContactNumber != maskFull() {
		t.Errorf("contactNumber: %v", *s.ContactNumber)
	}
	if *s.Address != maskFull() {
		t.Errorf("address: %v", *s.Address)
	}
	if s.Longitude != maskFull() || s.Latitude != maskFull() {
		t.Errorf("geo not masked: %v %v", s.Longitude, s.Latitude)
	}
	if s.IMEI != maskFull() {
		t.Errorf("imei: %v", s.IMEI)
	}
	// 保留字段
	if s.ICCID != iccid {
		t.Errorf("iccid should be preserved: %v", s.ICCID)
	}
	if s.City != "Shanghai" {
		t.Errorf("city should be preserved: %v", s.City)
	}

	// 上游未被污染
	if *orig.UserName != "Alice" || *orig.ContactNumber != "13812345678" || *orig.Address != "123 Main St" {
		t.Errorf("upstream polluted: name=%v phone=%v addr=%v", *orig.UserName, *orig.ContactNumber, *orig.Address)
	}
	if orig.Longitude != "121.5" || orig.IMEI != "356938035643809" {
		t.Errorf("upstream scalar polluted: lng=%v imei=%v", orig.Longitude, orig.IMEI)
	}
}

func TestSanitizeResPageSiteInfo(t *testing.T) {
	name := "Bob"
	orig := &dto.ResPage[dto.SiteInfoDTO]{
		Data:  []dto.SiteInfoDTO{{Sn: "S1", UserName: &name, IMEI: "111"}},
		Total: 1,
	}
	sanitized, ok := sanitizeStrong(orig)
	if !ok {
		t.Fatal("expected hit")
	}
	s := sanitized.(*dto.ResPage[dto.SiteInfoDTO])
	if *s.Data[0].UserName != "B***" {
		t.Errorf("page item userName: %v", *s.Data[0].UserName)
	}
	// 上游 slice 未被污染
	if *orig.Data[0].UserName != "Bob" {
		t.Errorf("upstream page polluted: %v", *orig.Data[0].UserName)
	}
}

func TestSanitizeDomainSimCardInfo(t *testing.T) {
	orig := &domain.SimCardInfo{GICCID: "GIC19", IMEI: "imei1", IMSI: "imsi1", MSISDN: "msisdn1"}
	sanitized, ok := sanitizeStrong(orig)
	if !ok {
		t.Fatal("expected hit")
	}
	s := sanitized.(*domain.SimCardInfo)
	if s.IMEI != maskFull() || s.IMSI != maskFull() || s.MSISDN != maskFull() {
		t.Errorf("sim identifiers not masked: %+v", s)
	}
	if s.GICCID != "GIC19" {
		t.Errorf("giccid should be preserved: %v", s.GICCID)
	}
	// 上游未污染
	if orig.IMEI != "imei1" {
		t.Errorf("upstream sim polluted: %v", orig.IMEI)
	}
}

// ---------- A/B 一致性测试 ----------

func TestAB_Consistency_SiteInfo(t *testing.T) {
	name := "Alice"
	phone := "13812345678"
	addr := "123 Main St"
	site := &dto.SiteInfoDTO{
		Sn:            "SN1",
		UserName:      &name,
		ContactNumber: &phone,
		Address:       &addr,
		Longitude:     "121.5",
		City:          "Shanghai",
		ICCID:         "8986012345678901234",
		IMEI:          "356938035643809",
	}

	// A 路径
	sanitized, _ := sanitizeStrong(site)
	aBytes, err := json.Marshal(sanitized)
	if err != nil {
		t.Fatal(err)
	}

	// B 路径：直接对原始结构体序列化后走流式脱敏
	rawBytes, _ := json.Marshal(site)
	bBytes, err := MaskJSONBytes(rawBytes)
	if err != nil {
		t.Fatalf("B err: %v", err)
	}

	// 归一化比较（字段顺序可能不同，比较解析后的 map）
	var ma, mb map[string]any
	if err := json.Unmarshal(aBytes, &ma); err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(bBytes, &mb); err != nil {
		t.Fatal(err)
	}

	for _, k := range []string{"userName", "contactNumber", "address", "longitude", "imei", "iccId", "city", "sn"} {
		if ma[k] != mb[k] {
			t.Errorf("A/B mismatch on %q: A=%v B=%v", k, ma[k], mb[k])
		}
	}
}