diff --git a/.gitignore b/.gitignore index 31c0992..a1ebf22 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ .idea *.mmdb +# Ignore all mmdb files in data-snapshots directory, used for unit-testing +!data-snapshots/*.mmdb **/venv/* **/__pycache__/* # zannotate binary diff --git a/README.md b/README.md index 745f477..86e8835 100644 --- a/README.md +++ b/README.md @@ -33,3 +33,25 @@ Check that it was installed correctly with: ```shell zannotate --help ``` + +# Acquiring Datasets + +> [!NOTE] +> URLs and instructions may change over time. These are up-to-date as of September 2025. + +Below are instructions for getting datasets from the below providers. + +### GeoLite2-ASN +1. [Sign-up form](https://www.maxmind.com/en/geolite2/signup) for MaxMind GeoLite Access +2. Login to your account +3. Go to the "GeoIP / GeoLite" > "Download files" section where you should see a list of available databases +4. Download the `.mmdb` files for GeoLite ASN +5. Unzip the downloaded file and test with: + +```shell +echo "1.1.1.1" | zannotate --geoasn --geoasn-database=/path-to-downloaded-file/GeoLite2-ASN_20250923/GeoLite2-ASN.mmdb +``` + +```shell +{"ip":"1.1.1.1","geoasn":{"asn":13335,"org":"CLOUDFLARENET"}} +``` diff --git a/data-snapshots/geolite2_asn.mmdb b/data-snapshots/geolite2_asn.mmdb new file mode 100644 index 0000000..62987b7 Binary files /dev/null and b/data-snapshots/geolite2_asn.mmdb differ diff --git a/geoipasn_test.go b/geoipasn_test.go new file mode 100644 index 0000000..2665708 --- /dev/null +++ b/geoipasn_test.go @@ -0,0 +1,77 @@ +/* + * ZAnnotate Copyright 2025 Regents of the University of Michigan + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy + * of the License at http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + * implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package zannotate + +import ( + "net" + "reflect" + "testing" +) + +func TestGeoIPASNAnnotator(t *testing.T) { + tests := []struct { + testName string + ipAddr net.IP + expectedResult *GeoIPASNOutput + }{ + { + testName: "Positive Test Case, IPv4", + ipAddr: net.ParseIP("1.1.1.1"), + expectedResult: &GeoIPASNOutput{ + ASN: 13335, + ASNOrg: "CLOUDFLARENET", + }, + }, { + testName: "Positive Test Case, IPv6", + ipAddr: net.ParseIP("2606:4700:4700::1111"), + expectedResult: &GeoIPASNOutput{ + ASN: 13335, + ASNOrg: "CLOUDFLARENET", + }, + }, { + testName: "Negative Test Case, Invalid IP", + ipAddr: net.ParseIP("999.999.999.999"), + expectedResult: &GeoIPASNOutput{}, + }, { + testName: "Negative Test Case, Private IP", + ipAddr: net.ParseIP("127.0.0.1"), + expectedResult: &GeoIPASNOutput{}, + }, + } + factory := &GeoIPASNAnnotatorFactory{ + Path: "./data-snapshots/geolite2_asn.mmdb", + Mode: "mmap", + } + err := factory.Initialize(nil) + if err != nil { + t.Fatalf("Failed to initialize factory: %v", err) + } + for _, tt := range tests { + t.Run(tt.testName, func(t *testing.T) { + annotator := factory.MakeAnnotator(0).(*GeoIPASNAnnotator) + err = annotator.Initialize() + if err != nil { + t.Fatalf("Failed to initialize annotator: %v", err) + } + result := annotator.Annotate(tt.ipAddr) + if tt.expectedResult == nil && result == nil { + return // pass + } + if !reflect.DeepEqual(result, tt.expectedResult) { + t.Errorf("Annotating IP %s gave = %v; expected: %v", tt.ipAddr, result, tt.expectedResult) + } + }) + } +}