tamabox/internal/db/ipgeo_test.go

// Copyright 2026 tamakyi. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package db

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

// TestIPGeoStore 覆盖地理定位缓存的读写、批量取、覆盖与过期清理。
//
// 这张表是配额的保护层：同一批攻击者会在多份报告里反复出现，
// 「查一次、之后全靠缓存」能不能成立全看这里的语义对不对。
func TestIPGeoStore(t *testing.T) {
	database, cleanup := newTestDB(t)
	ctx := context.Background()

	t.Cleanup(func() {
		if err := cleanup("ip_geo_cache"); err != nil {
			t.Fatal(err)
		}
	})

	store := NewIPGeoStore(database)
	now := time.Now()

	records := []*IPGeoRecord{
		{
			IPKey: "203.0.113.1", Country: "美国", CountryCode: "US",
			Region: "Virginia", City: "Ashburn", ISP: "Amazon", ASN: "AS14618",
			Lat: 39.04, Lon: -77.49, HasCoords: true, Source: "ipapi", ResolvedAt: now,
		},
		{
			IPKey: "2001:db8::/64", Country: "中国", CountryCode: "CN",
			Region: "北京市", City: "北京", Lat: 39.9, Lon: 116.4,
			HasCoords: true, Source: "tencent", ResolvedAt: now,
		},
		{
			// 查不到：Failed 标记要能落库，否则每次报告都会重试同一批地址
			IPKey: "198.51.100.9", Failed: true, ResolvedAt: now,
		},
	}
	require.NoError(t, store.UpsertMany(ctx, records))

	got, err := store.GetMany(ctx, []string{"203.0.113.1", "2001:db8::/64", "198.51.100.9", "10.0.0.1"})
	require.NoError(t, err)
	require.Len(t, got, 3, "不存在的键不应出现在结果里")
	require.Equal(t, "Ashburn", got["203.0.113.1"].City)
	require.Equal(t, "AS14618", got["203.0.113.1"].ASN)
	require.True(t, got["203.0.113.1"].HasCoords)
	require.Equal(t, "tencent", got["2001:db8::/64"].Source)
	require.True(t, got["198.51.100.9"].Failed)

	// 覆盖写：同一个键再写一次应当更新而不是报主键冲突
	require.NoError(t, store.UpsertMany(ctx, []*IPGeoRecord{{
		IPKey: "203.0.113.1", Country: "美国", CountryCode: "US",
		City: "Reston", Lat: 38.96, Lon: -77.35, HasCoords: true,
		Source: "tencent", ResolvedAt: now.Add(time.Hour),
	}}))
	got2, err := store.GetMany(ctx, []string{"203.0.113.1"})
	require.NoError(t, err)
	require.Equal(t, "Reston", got2["203.0.113.1"].City)
	require.Equal(t, "tencent", got2["203.0.113.1"].Source)

	n, err := store.Count(ctx)
	require.NoError(t, err)
	require.Equal(t, int64(3), n)
}

// TestIPGeoStorePurge 覆盖两套 TTL 分开清理：成功留 90 天、失败留 3 天。
func TestIPGeoStorePurge(t *testing.T) {
	database, cleanup := newTestDB(t)
	ctx := context.Background()

	t.Cleanup(func() {
		if err := cleanup("ip_geo_cache"); err != nil {
			t.Fatal(err)
		}
	})

	store := NewIPGeoStore(database)
	now := time.Now()

	require.NoError(t, store.UpsertMany(ctx, []*IPGeoRecord{
		{IPKey: "1.0.0.1", Country: "澳大利亚", ResolvedAt: now.AddDate(0, 0, -100)}, // 成功且过期
		{IPKey: "1.0.0.2", Country: "日本", ResolvedAt: now.AddDate(0, 0, -10)},    // 成功且新鲜
		{IPKey: "1.0.0.3", Failed: true, ResolvedAt: now.AddDate(0, 0, -4)},      // 失败且过期
		{IPKey: "1.0.0.4", Failed: true, ResolvedAt: now},                        // 失败且新鲜
	}))

	// 清失败的旧记录：只删 1.0.0.3
	n, err := store.PurgeResolvedBefore(ctx, now.AddDate(0, 0, -3), true)
	require.NoError(t, err)
	require.Equal(t, int64(1), n)

	// 清成功的旧记录：只删 1.0.0.1（不能连失败的一起删，两者 TTL 不同）
	n, err = store.PurgeResolvedBefore(ctx, now.AddDate(0, 0, -90), false)
	require.NoError(t, err)
	require.Equal(t, int64(1), n)

	left, err := store.GetMany(ctx, []string{"1.0.0.1", "1.0.0.2", "1.0.0.3", "1.0.0.4"})
	require.NoError(t, err)
	require.Len(t, left, 2)
	require.Contains(t, left, "1.0.0.2")
	require.Contains(t, left, "1.0.0.4")
}

// TestSecurityReportStore 覆盖报告台账的幂等语义。
//
// 这是「重启不会重复发同一期周报」的唯一保证：period_key 唯一索引 + upsert。
func TestSecurityReportStore(t *testing.T) {
	database, cleanup := newTestDB(t)
	ctx := context.Background()

	t.Cleanup(func() {
		if err := cleanup("security_report_logs"); err != nil {
			t.Fatal(err)
		}
	})

	store := NewSecurityReportStore(database)
	start := time.Date(2026, 10, 19, 0, 0, 0, 0, time.Local)
	end := start.AddDate(0, 0, 7)

	// 首次：不存在
	rec, err := store.Get(ctx, "week:2026-W43")
	require.NoError(t, err)
	require.Nil(t, rec)

	require.NoError(t, store.Save(ctx, &SecurityReportLog{
		PeriodKey: "week:2026-W43", Kind: "week",
		WindowStart: start, WindowEnd: end,
		Hits: 42, UniqueSources: 7, Countries: 3,
		Recipient: "[email protected]", Attempts: 1, Error: "smtp timeout",
	}))

	rec, err = store.Get(ctx, "week:2026-W43")
	require.NoError(t, err)
	require.NotNil(t, rec)
	require.False(t, rec.Sent)
	require.Equal(t, 1, rec.Attempts)

	// 第二次：同一 period_key 再写应当覆盖，而不是插入第二行 ——
	// 否则「这一期发过没」会被多行记录搞乱。
	require.NoError(t, store.Save(ctx, &SecurityReportLog{
		PeriodKey: "week:2026-W43", Kind: "week",
		WindowStart: start, WindowEnd: end,
		Hits: 42, UniqueSources: 7, Countries: 3,
		Recipient: "[email protected]", Attempts: 2, Sent: true,
	}))

	rec, err = store.Get(ctx, "week:2026-W43")
	require.NoError(t, err)
	require.True(t, rec.Sent)
	require.Equal(t, 2, rec.Attempts)

	recent, err := store.Recent(ctx, 10)
	require.NoError(t, err)
	require.Len(t, recent, 1, "同一期只能有一行台账")

	// 不同周期应各自一行
	require.NoError(t, store.Save(ctx, &SecurityReportLog{
		PeriodKey: "month:2026-10", Kind: "month",
		WindowStart: start, WindowEnd: end, Sent: true,
	}))
	recent, err = store.Recent(ctx, 10)
	require.NoError(t, err)
	require.Len(t, recent, 2)
}