| 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 | } |