master
go 221 lines 6.68 KB
Raw
1 // SPDX-License-Identifier: GPL-3.0-or-later
2
3 package main
4
5 import (
6 "encoding/json"
7 "net"
8 "net/netip"
9 "os"
10 "path/filepath"
11 "testing"
12
13 "github.com/maxmind/mmdbwriter"
14 "github.com/maxmind/mmdbwriter/mmdbtype"
15 "github.com/stretchr/testify/require"
16 )
17
18 func TestWriteOutputsAndClassifications(t *testing.T) {
19 cfg := defaultConfig()
20 cfg.output.directory = t.TempDir()
21 cfg.output.asnFile = "asn.mmdb"
22 cfg.output.geoFile = "geo.mmdb"
23 cfg.output.metadataFile = "meta.json"
24 cfg.policy.localhostCIDRs = []string{"127.0.0.0/8"}
25 cfg.policy.privateCIDRs = []string{"10.0.0.0/8"}
26 cfg.policy.interestingCIDRs = []string{"203.0.113.0/24"}
27
28 asnRanges := []asnRange{{
29 start: mustParseAddr(t, "1.0.0.0"),
30 end: mustParseAddr(t, "1.0.0.255"),
31 asn: 13335,
32 org: "Cloudflare",
33 }}
34 geoRanges := []geoRange{{
35 start: mustParseAddr(t, "1.0.0.0"),
36 end: mustParseAddr(t, "1.0.0.255"),
37 country: "US",
38 state: "California",
39 city: "Los Angeles",
40 latitude: 34.0522,
41 longitude: -118.2437,
42 hasLocation: true,
43 }}
44 sources := []generationDatasetRef{
45 {
46 Name: "dbip-asn",
47 Family: sourceFamilyASN,
48 Provider: providerDBIP,
49 Artifact: artifactDBIPASNLite,
50 Source: "builtin",
51 Format: formatMMDB,
52 ResolvedURL: "https://download.db-ip.com/free/dbip-asn-lite-2026-03.mmdb.gz",
53 },
54 {
55 Name: "dbip-geo",
56 Family: sourceFamilyGeo,
57 Provider: providerDBIP,
58 Artifact: artifactDBIPCityLite,
59 Source: "builtin",
60 Format: formatMMDB,
61 ResolvedURL: "https://download.db-ip.com/free/dbip-city-lite-2026-03.mmdb.gz",
62 },
63 }
64
65 require.NoError(t, writeOutputs(cfg, asnRanges, geoRanges, sources))
66
67 loadOpts := mmdbwriter.Options{IncludeReservedNetworks: true, DisableIPv4Aliasing: true}
68 asnDB, err := mmdbwriter.Load(filepath.Join(cfg.output.directory, cfg.output.asnFile), loadOpts)
69 require.NoError(t, err)
70
71 geoDB, err := mmdbwriter.Load(filepath.Join(cfg.output.directory, cfg.output.geoFile), loadOpts)
72 require.NoError(t, err)
73
74 publicRec := mustLookupMap(t, asnDB, "1.0.0.1")
75 require.Equal(t, uint32(13335), uint32(publicRec["autonomous_system_number"].(mmdbtype.Uint32)))
76 require.Equal(t, "Cloudflare", string(publicRec["autonomous_system_organization"].(mmdbtype.String)))
77
78 localhostRec := mustLookupMap(t, asnDB, "127.0.0.1")
79 require.Equal(t, "localhost", netdataClass(localhostRec))
80 require.True(t, netdataTrackIndividual(localhostRec))
81
82 privateRec := mustLookupMap(t, asnDB, "10.1.2.3")
83 require.Equal(t, "private", netdataClass(privateRec))
84 require.True(t, netdataTrackIndividual(privateRec))
85
86 interestingRec := mustLookupMap(t, asnDB, "203.0.113.55")
87 require.Equal(t, "interesting", netdataClass(interestingRec))
88 require.True(t, netdataTrackIndividual(interestingRec))
89
90 geoRec := mustLookupMap(t, geoDB, "1.0.0.8")
91 countryMap, ok := geoRec["country"].(mmdbtype.Map)
92 require.True(t, ok)
93 require.Equal(t, "US", string(countryMap["iso_code"].(mmdbtype.String)))
94 require.Equal(t, "California", string(geoRec["region"].(mmdbtype.String)))
95
96 cityMap, ok := geoRec["city"].(mmdbtype.Map)
97 require.True(t, ok)
98 cityNames := cityMap["names"].(mmdbtype.Map)
99 require.Equal(t, "Los Angeles", string(cityNames["en"].(mmdbtype.String)))
100
101 locationMap, ok := geoRec["location"].(mmdbtype.Map)
102 require.True(t, ok)
103 require.Equal(t, float64(34.0522), float64(locationMap["latitude"].(mmdbtype.Float64)))
104 require.Equal(t, float64(-118.2437), float64(locationMap["longitude"].(mmdbtype.Float64)))
105
106 metaPath := filepath.Join(cfg.output.directory, cfg.output.metadataFile)
107 metaBlob, err := os.ReadFile(metaPath)
108 require.NoError(t, err)
109
110 var meta generationMetadata
111 require.NoError(t, json.Unmarshal(metaBlob, &meta))
112 require.Equal(t, "topology-ip-intel-downloader", meta.GeneratedBy)
113 require.Len(t, meta.Sources, 2)
114 require.Equal(t, cfg.output.asnFile, meta.Output.AsnFile)
115 require.Equal(t, cfg.output.geoFile, meta.Output.GeoFile)
116 require.Equal(t, cfg.output.metadataFile, meta.Output.MetadataFile)
117 require.Equal(t, 1, meta.Counts.AsnRanges)
118 require.Equal(t, 1, meta.Counts.GeoRanges)
119 }
120
121 func TestWriteOutputsRemovesDisabledGeoFile(t *testing.T) {
122 cfg := defaultConfig()
123 cfg.output.directory = t.TempDir()
124 cfg.output.asnFile = "asn.mmdb"
125 cfg.output.geoFile = "geo.mmdb"
126 cfg.output.metadataFile = "meta.json"
127 cfg.sources = []sourceEntry{
128 {
129 family: sourceFamilyASN,
130 provider: providerDBIP,
131 artifact: artifactDBIPASNLite,
132 format: formatMMDB,
133 },
134 }
135
136 geoPath := filepath.Join(cfg.output.directory, cfg.output.geoFile)
137 require.NoError(t, os.WriteFile(geoPath, []byte("stale"), 0o644))
138
139 require.NoError(t, writeOutputs(cfg, nil, nil, nil))
140
141 _, err := os.Stat(geoPath)
142 require.ErrorIs(t, err, os.ErrNotExist)
143
144 metaBlob, err := os.ReadFile(filepath.Join(cfg.output.directory, cfg.output.metadataFile))
145 require.NoError(t, err)
146 var meta generationMetadata
147 require.NoError(t, json.Unmarshal(metaBlob, &meta))
148 require.Equal(t, cfg.output.asnFile, meta.Output.AsnFile)
149 require.Empty(t, meta.Output.GeoFile)
150 }
151
152 func TestWriteOutputsDoesNotPublishHalfBuiltState(t *testing.T) {
153 cfg := defaultConfig()
154 cfg.output.directory = t.TempDir()
155 cfg.output.asnFile = "asn.mmdb"
156 cfg.output.geoFile = "geo.mmdb"
157 cfg.output.metadataFile = "meta.json"
158
159 asnPath := filepath.Join(cfg.output.directory, cfg.output.asnFile)
160 geoPath := filepath.Join(cfg.output.directory, cfg.output.geoFile)
161 require.NoError(t, writeOutputs(cfg, nil, nil, nil))
162
163 asnBefore, err := os.ReadFile(asnPath)
164 require.NoError(t, err)
165 geoBefore, err := os.ReadFile(geoPath)
166 require.NoError(t, err)
167
168 invalidGeo := []geoRange{{
169 start: mustParseAddr(t, "1.0.0.0"),
170 end: mustParseAddr(t, "2001:db8::1"),
171 }}
172
173 err = writeOutputs(cfg, nil, invalidGeo, nil)
174 require.Error(t, err)
175
176 asnAfter, err := os.ReadFile(asnPath)
177 require.NoError(t, err)
178 geoAfter, err := os.ReadFile(geoPath)
179 require.NoError(t, err)
180 require.Equal(t, asnBefore, asnAfter)
181 require.Equal(t, geoBefore, geoAfter)
182 }
183
184 func mustParseAddr(t *testing.T, value string) netip.Addr {
185 t.Helper()
186 addr, err := netip.ParseAddr(value)
187 require.NoError(t, err)
188 return addr
189 }
190
191 func mustLookupMap(t *testing.T, tree *mmdbwriter.Tree, ip string) mmdbtype.Map {
192 t.Helper()
193 _, v := tree.Get(net.ParseIP(ip))
194 m, ok := v.(mmdbtype.Map)
195 require.True(t, ok)
196 return m
197 }
198
199 func netdataClass(record mmdbtype.Map) string {
200 netdata, ok := record["netdata"].(mmdbtype.Map)
201 if !ok {
202 return ""
203 }
204 classValue, ok := netdata["ip_class"].(mmdbtype.String)
205 if !ok {
206 return ""
207 }
208 return string(classValue)
209 }
210
211 func netdataTrackIndividual(record mmdbtype.Map) bool {
212 netdata, ok := record["netdata"].(mmdbtype.Map)
213 if !ok {
214 return false
215 }
216 flag, ok := netdata["track_individual"].(mmdbtype.Bool)
217 if !ok {
218 return false
219 }
220 return bool(flag)
221 }