diff --git a/data-prepper-plugins/geoip-processor/build.gradle b/data-prepper-plugins/geoip-processor/build.gradle index 916f78e2f1..b9a85deefc 100644 --- a/data-prepper-plugins/geoip-processor/build.gradle +++ b/data-prepper-plugins/geoip-processor/build.gradle @@ -1,15 +1,15 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + plugins{ - id 'de.undercouch.download' version '4.1.2' + id 'de.undercouch.download' version '5.5.0' } apply plugin: 'de.undercouch.download' import de.undercouch.gradle.tasks.download.Download -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - dependencies { implementation project(':data-prepper-api') implementation project(path: ':data-prepper-plugins:common') @@ -19,90 +19,64 @@ dependencies { implementation libs.commons.compress implementation 'org.mapdb:mapdb:3.0.8' implementation libs.commons.io - implementation 'software.amazon.awssdk:aws-sdk-java:2.20.67' + implementation 'software.amazon.awssdk:sts' implementation 'software.amazon.awssdk:s3-transfer-manager' - implementation 'software.amazon.awssdk.crt:aws-crt:0.21.17' + implementation 'software.amazon.awssdk.crt:aws-crt:0.29.9' implementation 'com.maxmind.geoip2:geoip2:4.0.1' implementation 'com.maxmind.db:maxmind-db:3.0.0' implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' - implementation libs.commons.lang3 + testImplementation project(':data-prepper-core') testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310' testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' testImplementation project(':data-prepper-test-common') } -def geoIP2='GeoIP2' -def geoLite2= 'GeoLite2' -task downloadFile(type: Download) { - - def urls = [ - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-City-Test.mmdb', - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-Country-Test.mmdb', - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoLite2-ASN-Test.mmdb' - ] - def mmdbFileExtension = '.mmdb' - def baseDirPath = 'src/test/resources/mmdb-file/geo-lite2/' - - urls.each { url -> - src(url) - dest(baseDirPath) - doLast { - def testFileName = url.substring(url.lastIndexOf('/') + 1) - def testMmdbSubString = testFileName.substring(testFileName.lastIndexOf('-')) - def fileName = testFileName.substring(0, testFileName.length() - testMmdbSubString.length()) +def downloadFiles = tasks.register('downloadFiles') - if(fileName.contains(geoIP2)) { - fileName = fileName.replace(geoIP2, geoLite2) - } - File sourceFile = file(baseDirPath+testFileName) - File destinationFile = file( baseDirPath+fileName+mmdbFileExtension) - sourceFile.renameTo(destinationFile) +def databaseNames = [ + 'GeoLite2-City-Test', + 'GeoLite2-Country-Test', + 'GeoLite2-ASN-Test' +] - } +databaseNames.forEach { databaseName -> { + def url = "https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/${databaseName}.mmdb" + def gradleName = databaseName.replaceAll('-', '') + def downloadTask = tasks.register("download${gradleName}", Download) { + src(url) + dest "build/resources/test/mmdb-files/geo-lite2/${databaseName}.mmdb" + overwrite true } + downloadFiles.get().dependsOn downloadTask +}} +def enterpriseDatabaseNames = [ + 'GeoIP2-Enterprise-Test' +] -} -task downloadEnterpriseFile(type: Download) { - dependsOn downloadFile - def urls = [ - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-Enterprise-Test.mmdb' - ] - def mmdbFileExtension = '.mmdb' - def baseDirPath = 'src/test/resources/mmdb-file/geo-enterprise/' - - urls.each { url -> - src(url) - def testFileName = url.substring(url.lastIndexOf('/') + 1) - def testMmdbSubString = testFileName.substring(testFileName.lastIndexOf('-')) - def fileName = testFileName.substring(0, testFileName.length() - testMmdbSubString.length()) - - dest(baseDirPath+testFileName) - doLast { - if(fileName.contains(geoIP2)) { - fileName = fileName.replace(geoIP2, geoLite2) - } - File sourceFile = file(baseDirPath+testFileName) - File destinationFile = file( baseDirPath+fileName+mmdbFileExtension) - sourceFile.renameTo(destinationFile) - } +enterpriseDatabaseNames.forEach { enterpriseDatabaseName -> { + def url = "https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/${enterpriseDatabaseName}.mmdb" + def gradleName = enterpriseDatabaseName.replaceAll('-', '') + def downloadEnterpriseTask = tasks.register("download${gradleName}", Download) { + src(url) + dest "build/resources/test/mmdb-files/geo-ip2/${enterpriseDatabaseName}.mmdb" + overwrite true } + downloadFiles.get().dependsOn downloadEnterpriseTask +}} -} -/*task processTestResources(type: Copy) { - dependsOn downloadEnterpriseFile - from 'src/test/resources' // Source directory containing test resources - into 'build/resources/test' // Destination directory for processed test resources -}*/ -tasks.test.dependsOn 'processTestResources' -tasks.processTestResources.dependsOn 'downloadEnterpriseFile' test { useJUnitPlatform() + dependsOn(downloadFiles) +} + +checkstyleTest { + dependsOn(downloadFiles) } jacocoTestCoverageVerification { @@ -110,40 +84,10 @@ jacocoTestCoverageVerification { violationRules { rule { limit { - minimum = 0.1 // temporarily reduce coverage for the builds to pass + minimum = 0.85 } } } } -check.dependsOn jacocoTestCoverageVerification - -sourceSets { - integrationTest { - java { - compileClasspath += main.output + test.output - runtimeClasspath += main.output + test.output - srcDir file('src/integrationTest/java') - } - resources.srcDir file('src/integrationTest/resources') - } -} - -configurations { - integrationTestImplementation.extendsFrom testImplementation - integrationTestRuntime.extendsFrom testRuntime -} - -task integrationTest(type: Test) { - group = 'verification' - testClassesDirs = sourceSets.integrationTest.output.classesDirs - - useJUnitPlatform() - - classpath = sourceSets.integrationTest.runtimeClasspath - systemProperty 'tests.geoipProcessor.maxmindLicenseKey', System.getProperty('tests.geoipProcessor.maxmindLicenseKey') - - filter { - includeTestsMatching '*IT' - } -} \ No newline at end of file +check.dependsOn jacocoTestCoverageVerification \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java b/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java index 70a80c0684..426686e099 100644 --- a/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java +++ b/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java @@ -81,7 +81,7 @@ public void setUp() throws JsonProcessingException { public GeoIPProcessorService createObjectUnderTest() { // TODO: pass in geoIpServiceConfig object - return new GeoIPProcessorService(null); + return new GeoIPProcessorService(null, null, null); } @Test @@ -93,7 +93,7 @@ void verify_enrichment_of_data_from_maxmind_url() throws UnknownHostException { if (IPValidationCheck.isPublicIpAddress(ipAddress)) { InetAddress inetAddress = InetAddress.getByName(ipAddress); //All attributes are considered by default with the null value - geoData = geoIPProcessorService.getGeoData(inetAddress, null); +// geoData = geoIPProcessorService.getGeoData(inetAddress, null); assertThat(geoData.get("country_iso_code"), equalTo("US")); assertThat(geoData.get("ip"), equalTo("8.8.8.8")); diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java new file mode 100644 index 0000000000..ff22e138b0 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java @@ -0,0 +1,13 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +public enum GeoIPDatabase { + CITY, + COUNTRY, + ASN, + ENTERPRISE; +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java new file mode 100644 index 0000000000..5476b7ea88 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java @@ -0,0 +1,72 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; + +public enum GeoIPField { + CONTINENT_CODE("continent_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + CONTINENT_NAME("continent_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_NAME("country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + IS_COUNTRY_IN_EUROPEAN_UNION("is_country_in_european_union", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_ISO_CODE("country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_CONFIDENCE("country_confidence", GeoIPDatabase.ENTERPRISE), + REGISTERED_COUNTRY_NAME("registered_country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REGISTERED_COUNTRY_ISO_CODE("registered_country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_NAME("represented_country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_ISO_CODE("represented_country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_TYPE("represented_country_type", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + CITY_NAME("city_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + CITY_CONFIDENCE("city_confidence", GeoIPDatabase.ENTERPRISE), + LOCATION("location", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LATITUDE("latitude", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LONGITUDE("longitude", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LOCATION_ACCURACY_RADIUS("location_accuracy_radius", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + METRO_CODE("metro_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + TIME_ZONE("time_zone", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + POSTAL_CODE("postal_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + POSTAL_CODE_CONFIDENCE("postal_code_confidence", GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_NAME("most_specified_subdivision_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_ISO_CODE("most_specified_subdivision_iso_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_CONFIDENCE("most_specified_subdivision_confidence", GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_NAME("least_specified_subdivision_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE("least_specified_subdivision_iso_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE("least_specified_subdivision_confidence", GeoIPDatabase.ENTERPRISE), + + ASN("asn", GeoIPDatabase.ASN), + ASN_ORGANIZATION("asn_organization", GeoIPDatabase.ASN), + NETWORK("network", GeoIPDatabase.ASN), + IP("ip", GeoIPDatabase.ASN); + + private final HashSet geoIPDatabases; + private final String fieldName; + + GeoIPField(final String fieldName, final GeoIPDatabase... geoIPDatabases) { + this.fieldName = fieldName; + this.geoIPDatabases = new HashSet<>(Arrays.asList(geoIPDatabases)); + } + + public static GeoIPField findByName(final String name) { + GeoIPField result = null; + for (GeoIPField geoIPField : values()) { + if (geoIPField.getFieldName().equalsIgnoreCase(name)) { + result = geoIPField; + break; + } + } + return result; + } + + public String getFieldName() { + return fieldName; + } + + public Set getGeoIPDatabases() { + return geoIPDatabases; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java index 1d04ede2e8..c81e2caa0f 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java @@ -15,7 +15,9 @@ import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.InvalidIPAddressException; import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; import org.opensearch.dataprepper.plugins.processor.extension.GeoIpConfigSupplier; import org.opensearch.dataprepper.plugins.processor.utils.IPValidationCheck; @@ -23,11 +25,15 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.IOException; import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.ArrayList; import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; /** * Implementation class of geoIP-processor plugin. It is responsible for enrichment of @@ -47,6 +53,8 @@ public class GeoIPProcessor extends AbstractProcessor, Record tagsOnFailure; private final GeoIPProcessorService geoIPProcessorService; private final ExpressionEvaluator expressionEvaluator; + private final Map> entryFieldsMap; + final Map> entryDatabaseMap; /** * GeoIPProcessor constructor for initialization of required attributes @@ -60,14 +68,20 @@ public GeoIPProcessor(final PluginMetrics pluginMetrics, final GeoIpConfigSupplier geoIpConfigSupplier, final ExpressionEvaluator expressionEvaluator) { super(pluginMetrics); + if (geoIpConfigSupplier.getGeoIPProcessorService().isEmpty()) { + throw new RuntimeException("geoip_service configuration is required when using geoip processor."); + } + this.geoIPProcessorService = geoIpConfigSupplier.getGeoIPProcessorService().get(); this.geoIPProcessorConfig = geoIPProcessorConfig; - this.geoIPProcessorService = geoIpConfigSupplier.getGeoIPProcessorService(); this.tagsOnFailure = geoIPProcessorConfig.getTagsOnFailure(); this.expressionEvaluator = expressionEvaluator; this.geoIpEventsProcessed = pluginMetrics.counter(GEO_IP_EVENTS_PROCESSED); this.geoIpEventsFailedLookup = pluginMetrics.counter(GEO_IP_EVENTS_FAILED_LOOKUP); //TODO: Use the exception metric for exceptions from service this.geoIpEventsFailedEngineException = pluginMetrics.counter(GEO_IP_EVENTS_FAILED_ENGINE_EXCEPTION); + + this.entryFieldsMap = populateGeoIPFields(); + this.entryDatabaseMap = populateGeoIPDatabases(); } /** @@ -79,34 +93,62 @@ public GeoIPProcessor(final PluginMetrics pluginMetrics, public Collection> doExecute(final Collection> records) { Map geoData; + final GeoIPDatabaseReader geoIPDatabaseReader = geoIPProcessorService.getGeoIPDatabaseReader(); + final boolean databasesExpired = geoIPDatabaseReader.isExpired(); + for (final Record eventRecord : records) { final Event event = eventRecord.getData(); + final String whenCondition = geoIPProcessorConfig.getWhenCondition(); if (whenCondition != null && !expressionEvaluator.evaluateConditional(whenCondition, event)) { continue; } + geoIpEventsProcessed.increment(); + + // TODO: Need to decide the behaviour, right now if all databases are expired we don't enrich the data. + if (databasesExpired) { + // TODO: Finalize the tags + continue; + } boolean isEventFailedLookup = false; - geoIpEventsProcessed.increment(); for (final EntryConfig entry : geoIPProcessorConfig.getEntries()) { final String source = entry.getSource(); - final List attributes = entry.getFields(); - final String ipAddress = event.get(source, String.class); + final List fields = entryFieldsMap.get(entry); + final Set databases = entryDatabaseMap.get(entry); + String ipAddress = null; + try { + ipAddress = event.get(source, String.class); + } catch (final Exception e) { + // add tags + LOG.error(DataPrepperMarkers.EVENT, "Failed to get IP address from [{}] in event: [{}]. Caused by:[{}]", + source, event, e.getMessage()); + } + //Lookup from DB if (ipAddress != null && !ipAddress.isEmpty()) { try { if (IPValidationCheck.isPublicIpAddress(ipAddress)) { - geoData = geoIPProcessorService.getGeoData(InetAddress.getByName(ipAddress), attributes); - eventRecord.getData().put(entry.getTarget(), geoData); + geoData = geoIPDatabaseReader.getGeoData(InetAddress.getByName(ipAddress), fields, databases); + if (geoData.isEmpty()) { + isEventFailedLookup = true; + } else { + eventRecord.getData().put(entry.getTarget(), geoData); + } } else { isEventFailedLookup = true; } - } catch (final IOException | EnrichFailedException ex) { + } catch (final InvalidIPAddressException | UnknownHostException e) { + isEventFailedLookup = true; + LOG.error(DataPrepperMarkers.EVENT, "Failed to validate IP address: [{}] in event: [{}]. Caused by:[{}]", + ipAddress, event, e.getMessage()); + } catch (final EnrichFailedException e) { isEventFailedLookup = true; - LOG.error(DataPrepperMarkers.EVENT, "Failed to get Geo data for event: [{}] for the IP address [{}]", event, ipAddress, ex); + LOG.error(DataPrepperMarkers.EVENT, "Failed to get Geo data for event: [{}] for the IP address [{}]. Caused by:{}", + event, ipAddress, e.getMessage()); } } else { //No Enrichment. @@ -119,21 +161,65 @@ public Collection> doExecute(final Collection> recor event.getMetadata().addTags(tagsOnFailure); } } + geoIPDatabaseReader.close(); return records; } + private Map> populateGeoIPFields() { + final Map> entryConfigFieldsMap = new HashMap<>(); + for (final EntryConfig entry: geoIPProcessorConfig.getEntries()) { + final List includeFields = entry.getIncludeFields(); + final List excludeFields = entry.getExcludeFields(); + List geoIPFields = new ArrayList<>(); + if (includeFields != null && !includeFields.isEmpty()) { + for (final String field : includeFields) { + final GeoIPField geoIPField = GeoIPField.findByName(field); + if (geoIPField != null) { + geoIPFields.add(geoIPField); + } + } + } else if (excludeFields != null) { + final List excludeGeoIPFields = new ArrayList<>(); + for (final String field : excludeFields) { + final GeoIPField geoIPField = GeoIPField.findByName(field); + if (geoIPField != null) { + excludeGeoIPFields.add(geoIPField); + } + } + geoIPFields = new ArrayList<>(List.of(GeoIPField.values())); + geoIPFields.removeAll(excludeGeoIPFields); + } + entryConfigFieldsMap.put(entry, geoIPFields); + } + return entryConfigFieldsMap; + } + + private Map> populateGeoIPDatabases() { + final Map> entryConfigGeoIPDatabaseMap = new HashMap<>(); + for (final EntryConfig entry : geoIPProcessorConfig.getEntries()) { + final List geoIPFields = entryFieldsMap.get(entry); + final Set geoIPDatabasesToUse = new HashSet<>(); + for (final GeoIPField geoIPField : geoIPFields) { + final Set geoIPDatabases = geoIPField.getGeoIPDatabases(); + geoIPDatabasesToUse.addAll(geoIPDatabases); + } + entryConfigGeoIPDatabaseMap.put(entry, geoIPDatabasesToUse); + } + return entryConfigGeoIPDatabaseMap; + } + @Override public void prepareForShutdown() { } @Override public boolean isReadyForShutdown() { + geoIPProcessorService.shutdown(); return true; } @Override public void shutdown() { - //TODO: delete mmdb files - LOG.info("GeoIP plugin Shutdown"); + } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java index 222a8698e4..35612401ca 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java @@ -18,10 +18,10 @@ */ public class GeoIPProcessorConfig { - @JsonProperty("entries") + @Valid @NotNull @Size(min = 1) - @Valid + @JsonProperty("entries") private List entries; @JsonProperty("tags_on_failure") diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java index c1e95373a3..d9c22d6143 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java @@ -6,12 +6,13 @@ package org.opensearch.dataprepper.plugins.processor.configuration; import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.NotEmpty; import java.util.List; public class EntryConfig { - static final String DEFAULT_TARGET = "geolocation"; + static final String DEFAULT_TARGET = "geo"; @JsonProperty("source") @NotEmpty private String source; @@ -19,8 +20,11 @@ public class EntryConfig { @JsonProperty("target") private String target = DEFAULT_TARGET; - @JsonProperty("fields") - private List fields; + @JsonProperty("include_fields") + private List includeFields; + + @JsonProperty("exclude_fields") + private List excludeFields; public String getSource() { return source; @@ -30,7 +34,19 @@ public String getTarget() { return target; } - public List getFields() { - return fields; + public List getIncludeFields() { + return includeFields; + } + + public List getExcludeFields() { + return excludeFields; + } + + @AssertTrue(message = "include_fields and exclude_fields are mutually exclusive. include_fields or exclude_fields is required.") + boolean areFieldsValid() { + if (includeFields == null && excludeFields == null) { + return false; + } + return includeFields == null || excludeFields == null; } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java new file mode 100644 index 0000000000..0160918676 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java @@ -0,0 +1,177 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.EnterpriseResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.exception.DatabaseReaderInitializationException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +public class GeoIP2DatabaseReader implements GeoIPDatabaseReader, AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(GeoIP2DatabaseReader.class); + private static final String MAXMIND_GEOIP2_DATABASE_TYPE = "geoip2"; + private static final String ENTERPRISE_DATABASE = "enterprise"; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final String databasePath; + private final int cacheSize; + private final AtomicInteger closeCount; + private final GeoIPFileManager geoIPFileManager; + private final AtomicBoolean isEnterpriseDatabaseExpired; + private DatabaseReader enterpriseDatabaseReader; + private Instant enterpriseDatabaseBuildDate; + + public GeoIP2DatabaseReader(final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final String databasePath, final int cacheSize) { + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.databasePath = databasePath; + this.cacheSize = cacheSize; + closeCount = new AtomicInteger(1); + this.isEnterpriseDatabaseExpired = new AtomicBoolean(false); + buildDatabaseReaders(); + } + + private void buildDatabaseReaders() { + try { + final Optional enterpriseDatabaseName = getDatabaseName(ENTERPRISE_DATABASE, databasePath, MAXMIND_GEOIP2_DATABASE_TYPE); + + if (enterpriseDatabaseName.isPresent()) { + enterpriseDatabaseReader = databaseReaderBuilder.buildReader(Path.of( + databasePath + File.separator + enterpriseDatabaseName.get()), cacheSize); + enterpriseDatabaseBuildDate = enterpriseDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + } catch (final IOException ex) { + throw new DatabaseReaderInitializationException("Exception while creating GeoIP2 DatabaseReaders due to: " + ex.getMessage()); + } + + if (enterpriseDatabaseReader == null) { + throw new NoValidDatabaseFoundException("Unable to initialize GeoIP2 database, make sure it is valid."); + } + } + @Override + public Map getGeoData(final InetAddress inetAddress, final List fields, final Set geoIPDatabases) { + Map geoData = new HashMap<>(); + + try { + if (geoIPDatabases.contains(GeoIPDatabase.ENTERPRISE)) { + final Optional optionalEnterpriseResponse = enterpriseDatabaseReader.tryEnterprise(inetAddress); + optionalEnterpriseResponse.ifPresent(response -> processEnterpriseResponse(response, geoData, fields)); + } + + if (geoIPDatabases.contains(GeoIPDatabase.ASN)) { + final Optional asnResponse = enterpriseDatabaseReader.tryAsn(inetAddress); + asnResponse.ifPresent(response -> processAsnResponse(response, geoData, fields)); + } + + } catch (final GeoIp2Exception e) { + throw new EnrichFailedException("Address not found in database."); + } catch (final IOException e) { + throw new EnrichFailedException("Failed to close database readers gracefully. It can be due to expired databases"); + } + return geoData; + } + + private void processEnterpriseResponse(final EnterpriseResponse enterpriseResponse, final Map geoData, final List fields) { + final Continent continent = enterpriseResponse.getContinent(); + final Country country = enterpriseResponse.getCountry(); + final Country registeredCountry = enterpriseResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = enterpriseResponse.getRepresentedCountry(); + + final City city = enterpriseResponse.getCity(); + final Location location = enterpriseResponse.getLocation(); + final Postal postal = enterpriseResponse.getPostal(); + final Subdivision mostSpecificSubdivision = enterpriseResponse.getMostSpecificSubdivision(); + final Subdivision leastSpecificSubdivision = enterpriseResponse.getLeastSpecificSubdivision(); + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, true); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + extractCityFields(city, geoData, fields, true); + extractLocationFields(location, geoData, fields); + extractPostalFields(postal, geoData, fields, true); + extractMostSpecifiedSubdivisionFields(mostSpecificSubdivision, geoData, fields, true); + extractLeastSpecifiedSubdivisionFields(leastSpecificSubdivision, geoData, fields, true); + } + + private void processAsnResponse(final AsnResponse asnResponse, final Map geoData, final List fields) { + extractAsnFields(asnResponse, geoData, fields); + } + + @Override + public boolean isExpired() { + final Instant instant = Instant.now(); + if (enterpriseDatabaseReader == null) { + return true; + } + if (isEnterpriseDatabaseExpired.get()) { + return true; + } + if (enterpriseDatabaseBuildDate.plus(MAX_EXPIRY_DURATION).isBefore(instant)) { + isEnterpriseDatabaseExpired.set(true); + closeReader(); + } + return isEnterpriseDatabaseExpired.get(); + } + + @Override + public void retain() { + closeCount.incrementAndGet(); + } + + @Override + public void close() { + final int count = closeCount.decrementAndGet(); + if (count == 0) { + LOG.debug("Closing old geoip database readers"); + closeReader(); + } + } + + private void closeReader() { + try { + if (enterpriseDatabaseReader != null) { + enterpriseDatabaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database directory + final File file = new File(databasePath); + geoIPFileManager.deleteDirectory(file); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java new file mode 100644 index 0000000000..b05d67c116 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java @@ -0,0 +1,361 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; + +import java.io.File; +import java.net.InetAddress; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * Interface for storing and maintaining MaxMind database readers + */ + +public interface GeoIPDatabaseReader { + String MAXMIND_DATABASE_EXTENSION = ".mmdb"; + Duration MAX_EXPIRY_DURATION = Duration.ofDays(30); + String LAT = "lat"; + String LON = "lon"; + + /** + * Gets the geo data from the {@link com.maxmind.geoip2.DatabaseReader} + * + * @param inetAddress InetAddress + * @return Map of geo field and value pairs from IP address + * + * @since 2.7 + */ + Map getGeoData(InetAddress inetAddress, List fields, Set geoIPDatabases); + + /** + * Gets if the database is expired from metadata or last updated timestamp + * + * @return boolean indicating if database is expired + * + * @since 2.7 + */ + boolean isExpired(); + + /** + * Retains the reader which prevents from closing if it's being used + * + * @since 2.7 + */ + void retain(); + + /** + * Closes the reader after processing a batch + * + * @since 2.7 + */ + void close(); + + /** + * Enrich attributes + * @param geoData geoData + * @param fieldName fieldName + * @param fieldValue fieldValue + */ + default void enrichData(final Map geoData, final String fieldName, final Object fieldValue) { + if (!geoData.containsKey(fieldName) && fieldValue != null) { + geoData.put(fieldName, fieldValue); + } + } + + default Optional getDatabaseName(final String database, final String databasePath, final String databaseType) { + final File file = new File(databasePath); + if (file.exists() && file.isDirectory()) { + final String[] list = file.list(); + for (final String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (lowerCaseFileName.contains(database) + && fileName.endsWith(MAXMIND_DATABASE_EXTENSION) + && lowerCaseFileName.contains(databaseType)) { + return Optional.of(fileName); + } + } + } + return Optional.empty(); + } + + default void extractContinentFields(final Continent continent, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case CONTINENT_CODE: + enrichData(geoData, GeoIPField.CONTINENT_CODE.getFieldName(), continent.getCode()); + break; + case CONTINENT_NAME: + enrichData(geoData, GeoIPField.CONTINENT_NAME.getFieldName(), continent.getName()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.CONTINENT_CODE.getFieldName(), continent.getCode()); + enrichData(geoData, GeoIPField.CONTINENT_NAME.getFieldName(), continent.getName()); + } + } + + default void extractCountryFields(final Country country, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case COUNTRY_NAME: + enrichData(geoData, GeoIPField.COUNTRY_NAME.getFieldName(), country.getName()); + break; + case IS_COUNTRY_IN_EUROPEAN_UNION: + enrichData(geoData, GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName(), country.isInEuropeanUnion()); + break; + case COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.COUNTRY_ISO_CODE.getFieldName(), country.getIsoCode()); + break; + case COUNTRY_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.COUNTRY_CONFIDENCE.getFieldName(), country.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.COUNTRY_NAME.getFieldName(), country.getName()); + enrichData(geoData, GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName(), country.isInEuropeanUnion()); + enrichData(geoData, GeoIPField.COUNTRY_ISO_CODE.getFieldName(), country.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData,GeoIPField. COUNTRY_CONFIDENCE.getFieldName(), country.getConfidence()); + } + } + + default void extractRegisteredCountryFields(final Country registeredCountry, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case REGISTERED_COUNTRY_NAME: + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_NAME.getFieldName(), registeredCountry.getName()); + break; + case REGISTERED_COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_ISO_CODE.getFieldName(), registeredCountry.getIsoCode()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_NAME.getFieldName(), registeredCountry.getName()); + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_ISO_CODE.getFieldName(), registeredCountry.getIsoCode()); + } + } + + default void extractRepresentedCountryFields(final RepresentedCountry representedCountry, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case REPRESENTED_COUNTRY_NAME: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_NAME.getFieldName(), representedCountry.getName()); + break; + case REPRESENTED_COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_ISO_CODE.getFieldName(), representedCountry.getIsoCode()); + break; + case REPRESENTED_COUNTRY_TYPE: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_TYPE.getFieldName(), representedCountry.getType()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_NAME.getFieldName(), representedCountry.getName()); + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_ISO_CODE.getFieldName(), representedCountry.getIsoCode()); + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_TYPE.getFieldName(), representedCountry.getType()); + } + } + + default void extractCityFields(final City city, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + if (field.equals(GeoIPField.CITY_NAME)) { + enrichData(geoData, GeoIPField.CITY_NAME.getFieldName(), city.getName()); + } else if (isEnterpriseDatabase && field.equals(GeoIPField.CITY_CONFIDENCE)) { + enrichData(geoData, GeoIPField.CITY_CONFIDENCE.getFieldName(), city.getConfidence()); + } + } + } else{ + enrichData(geoData, GeoIPField.CITY_NAME.getFieldName(), city.getName()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.CITY_CONFIDENCE.getFieldName(), city.getConfidence()); + } + } + + default void extractLocationFields(final Location location, + final Map geoData, + final List fields) { + final Map locationObject = new HashMap<>(); + locationObject.put(LAT, location.getLatitude()); + locationObject.put(LON, location.getLongitude()); + + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case LOCATION: + enrichData(geoData, GeoIPField.LOCATION.getFieldName(), locationObject); + break; + case LATITUDE: + enrichData(geoData, GeoIPField.LATITUDE.getFieldName(), location.getLatitude()); + break; + case LONGITUDE: + enrichData(geoData, GeoIPField.LONGITUDE.getFieldName(), location.getLongitude()); + break; + case METRO_CODE: + enrichData(geoData, GeoIPField.METRO_CODE.getFieldName(), location.getMetroCode()); + break; + case TIME_ZONE: + enrichData(geoData, GeoIPField.TIME_ZONE.getFieldName(), location.getTimeZone()); + break; + case LOCATION_ACCURACY_RADIUS: + enrichData(geoData, GeoIPField.LOCATION_ACCURACY_RADIUS.getFieldName(), location.getAccuracyRadius()); + break; + } + } + } else{ + // add all fields - latitude & longitude will be part of location key + enrichData(geoData, GeoIPField.LOCATION.getFieldName(), locationObject); + enrichData(geoData, GeoIPField.METRO_CODE.getFieldName(), location.getMetroCode()); + enrichData(geoData, GeoIPField.TIME_ZONE.getFieldName(), location.getTimeZone()); + enrichData(geoData, GeoIPField.LOCATION_ACCURACY_RADIUS.getFieldName(), location.getAccuracyRadius()); + } + } + + default void extractPostalFields(final Postal postal, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + if (field.equals(GeoIPField.POSTAL_CODE)) { + enrichData(geoData, GeoIPField.POSTAL_CODE.getFieldName(), postal.getCode()); + } else if (isEnterpriseDatabase && field.equals(GeoIPField.POSTAL_CODE_CONFIDENCE)) { + enrichData(geoData, GeoIPField.POSTAL_CODE_CONFIDENCE.getFieldName(), postal.getConfidence()); + } + } + } else{ + enrichData(geoData, GeoIPField.POSTAL_CODE.getFieldName(), postal.getCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.POSTAL_CODE_CONFIDENCE.getFieldName(), postal.getConfidence()); + } + } + + default void extractMostSpecifiedSubdivisionFields(final Subdivision subdivision, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case MOST_SPECIFIED_SUBDIVISION_NAME: + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + break; + case MOST_SPECIFIED_SUBDIVISION_ISO_CODE: + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + break; + case MOST_SPECIFIED_SUBDIVISION_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + } + } + + default void extractLeastSpecifiedSubdivisionFields(final Subdivision subdivision, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case LEAST_SPECIFIED_SUBDIVISION_NAME: + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + break; + case LEAST_SPECIFIED_SUBDIVISION_ISO_CODE: + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + break; + case LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + } + } + + default void extractAsnFields(final AsnResponse asnResponse, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case ASN: + enrichData(geoData, GeoIPField.ASN.getFieldName(), asnResponse.getAutonomousSystemNumber()); + break; + case ASN_ORGANIZATION: + enrichData(geoData, GeoIPField.ASN_ORGANIZATION.getFieldName(), asnResponse.getAutonomousSystemOrganization()); + break; + case NETWORK: + enrichData(geoData, GeoIPField.NETWORK.getFieldName(), asnResponse.getNetwork().toString()); + break; + case IP: + enrichData(geoData, GeoIPField.IP.getFieldName(), asnResponse.getIpAddress()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.ASN.getFieldName(), asnResponse.getAutonomousSystemNumber()); + enrichData(geoData, GeoIPField.ASN_ORGANIZATION.getFieldName(), asnResponse.getAutonomousSystemOrganization()); + enrichData(geoData, GeoIPField.NETWORK.getFieldName(), asnResponse.getNetwork().toString()); + enrichData(geoData, GeoIPField.IP.getFieldName(), asnResponse.getIpAddress()); + } + } + + +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java new file mode 100644 index 0000000000..cf4778978b --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java @@ -0,0 +1,260 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.CityResponse; +import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.exception.DatabaseReaderInitializationException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +public class GeoLite2DatabaseReader implements GeoIPDatabaseReader, AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(GeoLite2DatabaseReader.class); + static final String MAXMIND_GEOLITE2_DATABASE_TYPE = "geolite2"; + static final String CITY_DATABASE = "city"; + static final String COUNTRY_DATABASE = "country"; + static final String ASN_DATABASE = "asn"; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final String databasePath; + private final int cacheSize; + private final AtomicInteger closeCount; + private final AtomicBoolean isCountryDatabaseExpired; + private final AtomicBoolean isCityDatabaseExpired; + private final AtomicBoolean isAsnDatabaseExpired; + private final GeoIPFileManager geoIPFileManager; + private DatabaseReader cityDatabaseReader; + private DatabaseReader countryDatabaseReader; + private DatabaseReader asnDatabaseReader; + private Instant cityDatabaseBuildDate; + private Instant countryDatabaseBuildDate; + private Instant asnDatabaseBuildDate; + + public GeoLite2DatabaseReader(final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final String databasePath, final int cacheSize) { + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.databasePath = databasePath; + this.cacheSize = cacheSize; + this.closeCount = new AtomicInteger(1); + this.isCountryDatabaseExpired = new AtomicBoolean(false); + this.isCityDatabaseExpired = new AtomicBoolean(false); + this.isAsnDatabaseExpired = new AtomicBoolean(false); + buildDatabaseReaders(); + } + + private void buildDatabaseReaders() { + try { + final Optional cityDatabaseName = getDatabaseName(CITY_DATABASE, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + final Optional countryDatabaseName = getDatabaseName(COUNTRY_DATABASE, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + final Optional asnDatabaseName = getDatabaseName(ASN_DATABASE, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + + if (cityDatabaseName.isPresent()) { + cityDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + cityDatabaseName.get()), cacheSize); + cityDatabaseBuildDate = cityDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + if (countryDatabaseName.isPresent()) { + countryDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + countryDatabaseName.get()), cacheSize); + countryDatabaseBuildDate = countryDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + if (asnDatabaseName.isPresent()) { + asnDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + asnDatabaseName.get()), cacheSize); + asnDatabaseBuildDate = asnDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + + } catch (final IOException ex) { + throw new DatabaseReaderInitializationException("Exception while creating GeoLite2 DatabaseReaders due to: " + ex.getMessage()); + } + + if (cityDatabaseReader == null && countryDatabaseReader == null && asnDatabaseReader == null) { + throw new NoValidDatabaseFoundException("Unable to initialize any GeoLite2 database, make sure they are valid."); + } + } + + @Override + public Map getGeoData(final InetAddress inetAddress, final List fields, final Set geoIPDatabases) { + final Map geoData = new HashMap<>(); + + try { + if (geoIPDatabases.contains(GeoIPDatabase.COUNTRY)) { + final Optional countryResponse = countryDatabaseReader.tryCountry(inetAddress); + countryResponse.ifPresent(response -> processCountryResponse(response, geoData, fields)); + } + + if (geoIPDatabases.contains(GeoIPDatabase.CITY)) { + final Optional cityResponse = cityDatabaseReader.tryCity(inetAddress); + cityResponse.ifPresent(response -> processCityResponse(response, geoData, fields, geoIPDatabases)); + } + + if (geoIPDatabases.contains(GeoIPDatabase.ASN)) { + final Optional asnResponse = asnDatabaseReader.tryAsn(inetAddress); + asnResponse.ifPresent(response -> processAsnResponse(response, geoData, fields)); + } + + } catch (final GeoIp2Exception e) { + throw new EnrichFailedException("Address not found in database."); + } catch (final IOException e) { + throw new EnrichFailedException("Failed to close database readers gracefully. It can be due to expired databases."); + } + return geoData; + } + + private void processCountryResponse(final CountryResponse countryResponse, final Map geoData, final List fields) { + final Continent continent = countryResponse.getContinent(); + final Country country = countryResponse.getCountry(); + final Country registeredCountry = countryResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = countryResponse.getRepresentedCountry(); + + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, false); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + } + + private void processCityResponse(final CityResponse cityResponse, + final Map geoData, + final List fields, + final Set geoIPDatabases) { + // Continent and Country fields are added from City database only if they are not extracted from Country database + if (!geoIPDatabases.contains(GeoIPDatabase.COUNTRY)) { + final Continent continent = cityResponse.getContinent(); + final Country country = cityResponse.getCountry(); + final Country registeredCountry = cityResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = cityResponse.getRepresentedCountry(); + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, false); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + } + + final City city = cityResponse.getCity(); + final Location location = cityResponse.getLocation(); + final Postal postal = cityResponse.getPostal(); + final Subdivision mostSpecificSubdivision = cityResponse.getMostSpecificSubdivision(); + final Subdivision leastSpecificSubdivision = cityResponse.getLeastSpecificSubdivision(); + + extractCityFields(city, geoData, fields, false); + extractLocationFields(location, geoData, fields); + extractPostalFields(postal, geoData, fields, false); + extractMostSpecifiedSubdivisionFields(mostSpecificSubdivision, geoData, fields, false); + extractLeastSpecifiedSubdivisionFields(leastSpecificSubdivision, geoData, fields, false); + } + + private void processAsnResponse(final AsnResponse asnResponse, final Map geoData, final List fields) { + extractAsnFields(asnResponse, geoData, fields); + } + + @Override + public void retain() { + closeCount.incrementAndGet(); + } + + @Override + public void close() { + final int count = closeCount.decrementAndGet(); + if (count == 0) { + LOG.debug("Closing old geoip database readers"); + closeReaders(); + } + } + + @Override + public boolean isExpired() { + // TODO: Decide the expiry behaviour + final Instant instant = Instant.now(); + return isDatabaseExpired(instant, countryDatabaseReader, isCountryDatabaseExpired, countryDatabaseBuildDate, COUNTRY_DATABASE) && + isDatabaseExpired(instant, cityDatabaseReader, isCityDatabaseExpired, cityDatabaseBuildDate, CITY_DATABASE) && + isDatabaseExpired(instant, asnDatabaseReader, isAsnDatabaseExpired, asnDatabaseBuildDate, ASN_DATABASE); + } + + private boolean isDatabaseExpired(final Instant instant, + final DatabaseReader databaseReader, + final AtomicBoolean isDatabaseExpired, + final Instant databaseBuildDate, + final String databaseName) { + if (databaseReader == null) { + // no need to delete - no action needed + return true; + } + if (isDatabaseExpired.get()) { + // Another thread already updated status to expired - no action needed + return true; + } + if (databaseBuildDate.plus(MAX_EXPIRY_DURATION).isBefore(instant)) { + isDatabaseExpired.set(true); + closeReader(databaseReader, databaseName); + } + return isDatabaseExpired.get(); + } + + private void closeReaders() { + try { + if (cityDatabaseReader != null) { + cityDatabaseReader.close(); + } + if (countryDatabaseReader != null) { + countryDatabaseReader.close(); + } + if (asnDatabaseReader != null) { + asnDatabaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database directory + final File file = new File(databasePath); + geoIPFileManager.deleteDirectory(file); + } + + private void closeReader(final DatabaseReader databaseReader, final String databaseName) { + try { + if (databaseReader != null) { + databaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database file + final Optional fileName = getDatabaseName(databaseName, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + fileName.ifPresent(response -> { + File file = new File(databasePath + File.separator + response); + geoIPFileManager.deleteFile(file); + }); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java deleted file mode 100644 index 3730be7990..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import java.net.InetAddress; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Interface for Downloading through S3 or URl or from local path - */ -public interface GetGeoData { - - public final String GeoLite2CityDB = "GeoLite2-City.mmdb"; - public final String GeoLite2CountryDB = "GeoLite2-Country.mmdb"; - public final String GeoLite2AsnDB = "GeoLite2-ASN.mmdb"; - public final String GeoIP2EnterpriseDB = "GeoIP2-Enterprise.mmdb"; - - void switchDatabaseReader(); - void closeReader(); - Map getGeoData(InetAddress inetAddress, List attributes, String tempDestDir); - - /** - * Enrich attributes - * @param geoData geoData - * @param attributeName attributeName - * @param attributeValue attributeValue - */ - default void enrichData(final Map geoData, final String attributeName, final String attributeValue) { - if (attributeValue != null) { - geoData.put(attributeName, attributeValue); - } - } - - /** - * Enrich region iso code - * @param geoData geoData - * @param countryIso countryIso - * @param subdivisionIso subdivisionIso - */ - default void enrichRegionIsoCode(final Map geoData, final String countryIso, final String subdivisionIso) { - if (countryIso != null && subdivisionIso != null) { - enrichData(geoData, "region_iso_code", countryIso + "-" + subdivisionIso); - } - } - - /** - * Enrich Location Data - * @param geoData geoData - * @param latitude latitude - * @param longitude longitude - */ - default void enrichLocationData(final Map geoData, final Double latitude, final Double longitude) { - if (latitude != null && longitude != null) { - Map locationObject = new HashMap<>(); - locationObject.put("lat", latitude); - locationObject.put("lon", longitude); - geoData.put("location", locationObject); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java deleted file mode 100644 index 70d888f587..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import com.maxmind.geoip2.DatabaseReader; -import com.maxmind.geoip2.exception.GeoIp2Exception; -import com.maxmind.geoip2.model.EnterpriseResponse; -import com.maxmind.geoip2.record.City; -import com.maxmind.geoip2.record.Continent; -import com.maxmind.geoip2.record.Country; -import com.maxmind.geoip2.record.Subdivision; -import com.maxmind.geoip2.record.Location; -import com.maxmind.geoip2.record.Postal; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderCreate; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.net.InetAddress; -import java.nio.file.Path; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Implementation class for enrichment of enterprise data - */ -public class GetGeoIP2Data implements GetGeoData { - - private static final Logger LOG = LoggerFactory.getLogger(GetGeoIP2Data.class); - public static final String COUNTRY_NAME = "country_name"; - public static final String CONTINENT_NAME = "continent_name"; - public static final String REGION_NAME = "region_name"; - public static final String CITY_NAME = "city_name"; - public static final String COUNTRY_ISO_CODE = "country_iso_code"; - public static final String IP = "ip"; - public static final String REGION_ISO_CODE = "region_iso_code"; - public static final String TIMEZONE = "timezone"; - public static final String LOCATION = "location"; - public static final String POSTAL = "postal"; - private DatabaseReader readerEnterprise; - private Country country; - private Continent continent; - private City city; - private Location location; - private Subdivision subdivision; - private String dbPath; - private int cacheSize; - private Postal postal; - private String tempDestDir; - - /** - * GetGeoLite2Data constructor for initialisation of attributes - * @param dbPath dbPath - * @param cacheSize cacheSize - */ - public GetGeoIP2Data(final String dbPath, final int cacheSize) { - this.dbPath = dbPath; - this.cacheSize = cacheSize; - initDatabaseReader(); - } - - /** - * Initialise all the DatabaseReader - */ - public void initDatabaseReader() { - try { - readerEnterprise = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + tempDestDir + File.separator + GeoIP2EnterpriseDB), cacheSize); - } catch (final IOException ex) { - LOG.error("Exception while creating GeoIP2 DatabaseReader: {0}", ex); - } - } - - /** - * Switch all the DatabaseReader - */ - @Override - public void switchDatabaseReader() { - LOG.info("Switching GeoIP2 DatabaseReader"); - closeReader(); - System.gc(); - File file = new File(dbPath); - DBSource.deleteDirectory(file); - initDatabaseReader(); - } - - /** - * Enrich the GeoData - * @param inetAddress inetAddress - * @param attributes attributes - * @return enriched data Map - */ - public Map getGeoData(InetAddress inetAddress, List attributes, String tempDestDir) { - Map geoData = new HashMap<>(); - if (GeoIPProcessorService.downloadReady) { - this.tempDestDir = tempDestDir; - GeoIPProcessorService.downloadReady = false; - switchDatabaseReader(); - } - try { - EnterpriseResponse enterpriseResponse = readerEnterprise.enterprise(inetAddress); - country = enterpriseResponse.getCountry(); - subdivision = enterpriseResponse.getMostSpecificSubdivision(); - city = enterpriseResponse.getCity(); - location = enterpriseResponse.getLocation(); - continent = enterpriseResponse.getContinent(); - postal = enterpriseResponse.getPostal(); - } catch (IOException | GeoIp2Exception ex) { - LOG.info("Look up Exception : {0}", ex); - } - - try { - if ((attributes != null) && (!attributes.isEmpty())) { - for (String attribute : attributes) { - switch (attribute) { - case IP: - enrichData(geoData, IP, inetAddress.getHostAddress()); - break; - case COUNTRY_ISO_CODE: - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - break; - case COUNTRY_NAME: - enrichData(geoData, COUNTRY_NAME, country.getName()); - break; - case CONTINENT_NAME: - enrichData(geoData, CONTINENT_NAME, continent.getName()); - break; - case REGION_ISO_CODE: - // ISO 3166-2 code for country subdivisions. - // See iso.org/iso-3166-country-codes.html - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - break; - case REGION_NAME: - enrichData(geoData, REGION_NAME, subdivision.getName()); - break; - case CITY_NAME: - enrichData(geoData, CITY_NAME, city.getName()); - break; - case TIMEZONE: - enrichData(geoData, TIMEZONE, location.getTimeZone()); - break; - case LOCATION: - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - break; - case POSTAL: - enrichData(geoData, "postalCode", postal.getCode()); - break; - } - } - } else { - - enrichData(geoData, IP, inetAddress.getHostAddress()); - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - enrichData(geoData, COUNTRY_NAME, country.getName()); - enrichData(geoData, CONTINENT_NAME, continent.getName()); - - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - - enrichData(geoData, REGION_NAME, subdivision.getName()); - enrichData(geoData, CITY_NAME, city.getName()); - enrichData(geoData, "postalCode", postal.getCode()); - - enrichData(geoData, TIMEZONE, location.getTimeZone()); - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - } - } catch (Exception ex) { - throw new EnrichFailedException("Enrichment failed exception" + ex); - } - return geoData; - } - - - /** - * Close the DatabaseReader - */ - @Override - public void closeReader() { - try { - if (readerEnterprise != null) - readerEnterprise.close(); - } catch (IOException ex) { - LOG.info("Close Enterprise DatabaseReader Exception : {0}", ex); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java deleted file mode 100644 index 0d719541e0..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java +++ /dev/null @@ -1,262 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import com.maxmind.db.Network; -import com.maxmind.geoip2.DatabaseReader; -import com.maxmind.geoip2.exception.GeoIp2Exception; -import com.maxmind.geoip2.model.AsnResponse; -import com.maxmind.geoip2.model.CityResponse; -import com.maxmind.geoip2.model.CountryResponse; -import com.maxmind.geoip2.record.City; -import com.maxmind.geoip2.record.Continent; -import com.maxmind.geoip2.record.Country; -import com.maxmind.geoip2.record.Subdivision; -import com.maxmind.geoip2.record.Location; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderCreate; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.net.InetAddress; -import java.nio.file.Path; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Implementation class for enrichment of geoip lite2 data - */ -public class GetGeoLite2Data implements GetGeoData { - - private static final Logger LOG = LoggerFactory.getLogger(GetGeoLite2Data.class); - public static final String COUNTRY_NAME = "country_name"; - public static final String CONTINENT_NAME = "continent_name"; - public static final String REGION_NAME = "region_name"; - public static final String CITY_NAME = "city_name"; - public static final String ORGANIZATION_NAME = "organization_name"; - public static final String NETWORK = "network"; - public static final String COUNTRY_ISO_CODE = "country_iso_code"; - public static final String IP = "ip"; - public static final String REGION_ISO_CODE = "region_iso_code"; - public static final String TIMEZONE = "timezone"; - public static final String LOCATION = "location"; - public static final String ASN = "asn"; - private DatabaseReader readerCity; - private DatabaseReader readerCountry; - private DatabaseReader readerAsn; - private Country country; - private Continent continent; - private City city; - private Location location; - private Subdivision subdivision; - private Long asn; - private String organizationName; - private Network network; - private String dbPath; - private int cacheSize; - private CityResponse responseCity; - private CountryResponse responseCountry; - private AsnResponse responseAsn; - private String tempDestDir; - - - /** - * GetGeoLite2Data constructor for initialisation of attributes - * @param dbPath dbPath - * @param cacheSize cacheSize - */ - public GetGeoLite2Data(final String dbPath, final int cacheSize) { - this.dbPath = dbPath; - this.cacheSize = cacheSize; - initDatabaseReader(); - } - - /** - * Initialise all the DatabaseReader - */ - private void initDatabaseReader() { - try { - readerCity = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2CityDB), cacheSize); - readerCountry = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2CountryDB), cacheSize); - readerAsn = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2AsnDB), cacheSize); - } catch (final IOException ex) { - LOG.error("Exception while creating GeoLite2 DatabaseReader: {0}", ex); - } - } - - /** - * Switch all the DatabaseReader - */ - @Override - public void switchDatabaseReader() { - LOG.info("Switching GeoLite2 DatabaseReaders"); - closeReaderCity(); - closeReaderCountry(); - closeReaderAsn(); - System.gc(); - final File file = new File(dbPath); - DBSource.deleteDirectory(file); - dbPath = tempDestDir; - initDatabaseReader(); - } - - /** - * Enrich the GeoData - * @param inetAddress inetAddress - * @param attributes attributes - * @return enriched data Map - */ - @Override - public Map getGeoData(final InetAddress inetAddress, final List attributes, final String tempDestDir) { - Map geoData = new HashMap<>(); - if (GeoIPProcessorService.downloadReady) { - this.tempDestDir = tempDestDir; - GeoIPProcessorService.downloadReady = false; - switchDatabaseReader(); - } - try { - responseCountry = readerCountry.country(inetAddress); - country = responseCountry.getCountry(); - continent = responseCountry.getContinent(); - - responseCity = readerCity.city(inetAddress); - city = responseCity.getCity(); - location = responseCity.getLocation(); - subdivision = responseCity.getMostSpecificSubdivision(); - - responseAsn = readerAsn.asn(inetAddress); - asn = responseAsn.getAutonomousSystemNumber(); - organizationName = responseAsn.getAutonomousSystemOrganization(); - network = responseAsn.getNetwork(); - } catch (IOException | GeoIp2Exception ex) { - LOG.info("Look up Exception : {0}", ex); - } - - try { - if ((attributes != null) && (!attributes.isEmpty())) { - for (String attribute : attributes) { - switch (attribute) { - case IP: - enrichData(geoData, IP, inetAddress.getHostAddress()); - break; - case COUNTRY_ISO_CODE: - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - break; - case COUNTRY_NAME: - enrichData(geoData, COUNTRY_NAME, country.getName()); - break; - case CONTINENT_NAME: - enrichData(geoData, CONTINENT_NAME, continent.getName()); - break; - case REGION_ISO_CODE: - // ISO 3166-2 code for country subdivisions. - // See iso.org/iso-3166-country-codes.html - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - break; - case REGION_NAME: - enrichData(geoData, REGION_NAME, subdivision.getName()); - break; - case CITY_NAME: - enrichData(geoData, CITY_NAME, city.getName()); - break; - case TIMEZONE: - enrichData(geoData, TIMEZONE, location.getTimeZone()); - break; - case LOCATION: - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - break; - case ASN: - if (asn != null) { - geoData.put(ASN, asn); - } - break; - case ORGANIZATION_NAME: - enrichData(geoData, ORGANIZATION_NAME, organizationName); - break; - case NETWORK: - enrichData(geoData, NETWORK,network!=null? network.toString():null); - break; - } - } - } else { - - enrichData(geoData, IP, inetAddress.getHostAddress()); - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - enrichData(geoData, COUNTRY_NAME, country.getName()); - enrichData(geoData, CONTINENT_NAME, continent.getName()); - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - - enrichData(geoData, REGION_NAME, subdivision.getName()); - enrichData(geoData, CITY_NAME, city.getName()); - enrichData(geoData, TIMEZONE, location.getTimeZone()); - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - - if (asn != null) { - geoData.put(ASN, asn); - } - - enrichData(geoData, ORGANIZATION_NAME, organizationName); - enrichData(geoData, NETWORK,network!=null? network.toString():null); - } - } catch (Exception ex) { - throw new EnrichFailedException("Enrichment failed exception" + ex); - } - return geoData; - } - - /** - * Close the all DatabaseReader - */ - @Override - public void closeReader() { - closeReaderCity(); - closeReaderCountry(); - closeReaderAsn(); - } - - /** - * Close the City DatabaseReader - */ - private void closeReaderCity() { - try { - if (readerCity != null) - readerCity.close(); - readerCity = null; - } catch (IOException ex) { - LOG.info("Close City DatabaseReader Exception : {0}", ex); - } - } - - /** - * Close the Country DatabaseReader - */ - private void closeReaderCountry() { - try { - if (readerCountry != null) - readerCountry.close(); - readerCountry = null; - } catch (IOException ex) { - LOG.info("Close Country DatabaseReader Exception : {0}", ex); - } - } - - /** - * Close the ASN DatabaseReader - */ - private void closeReaderAsn() { - try { - if (readerAsn != null) - readerAsn.close(); - readerAsn = null; - } catch (IOException ex) { - LOG.info("Close Asn DatabaseReader Exception : {0}", ex); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java new file mode 100644 index 0000000000..f51f249220 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class DatabaseReaderInitializationException extends EngineFailureException { + public DatabaseReaderInitializationException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java similarity index 50% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java index b667c3273e..05cf8d9613 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.processor.databaseenrich; +package org.opensearch.dataprepper.plugins.processor.exception; /** * Implementation class for DownloadFailedException Custom exception */ -public class DownloadFailedException extends RuntimeException { - public DownloadFailedException(String exceptionMsg) { +public class DownloadFailedException extends EngineFailureException { + public DownloadFailedException(final String exceptionMsg) { super(exceptionMsg); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java new file mode 100644 index 0000000000..f5958b521a --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class EngineFailureException extends RuntimeException { + public EngineFailureException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java similarity index 67% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java index 9182d83979..0f3dd2bfba 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.processor.databaseenrich; +package org.opensearch.dataprepper.plugins.processor.exception; /** * Implementation class for EnrichFailedException Custom exception */ public class EnrichFailedException extends RuntimeException { - public EnrichFailedException(String exceptionMsg) { + public EnrichFailedException(final String exceptionMsg) { super(exceptionMsg); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java new file mode 100644 index 0000000000..134941b5ad --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class InvalidIPAddressException extends EnrichFailedException { + public InvalidIPAddressException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java new file mode 100644 index 0000000000..218d8c0901 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class NoValidDatabaseFoundException extends EngineFailureException { + public NoValidDatabaseFoundException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java index 08e50381bf..6c07304171 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java @@ -5,15 +5,30 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; + +import java.util.Optional; +import java.util.concurrent.locks.ReentrantReadWriteLock; + public class DefaultGeoIpConfigSupplier implements GeoIpConfigSupplier { private final GeoIpServiceConfig geoIpServiceConfig; + private final ReentrantReadWriteLock.ReadLock readLock; + private final GeoIPDatabaseManager geoIPDatabaseManager; - public DefaultGeoIpConfigSupplier(final GeoIpServiceConfig geoIpServiceConfig) { + public DefaultGeoIpConfigSupplier(final GeoIpServiceConfig geoIpServiceConfig, + final GeoIPDatabaseManager geoIPDatabaseManager, + final ReentrantReadWriteLock.ReadLock readLock + ) { this.geoIpServiceConfig = geoIpServiceConfig; + this.geoIPDatabaseManager = geoIPDatabaseManager; + this.readLock = readLock; } @Override - public GeoIPProcessorService getGeoIPProcessorService() { - return new GeoIPProcessorService(geoIpServiceConfig); + public Optional getGeoIPProcessorService() { + if (geoIpServiceConfig != null) + return Optional.of(new GeoIPProcessorService(geoIpServiceConfig, geoIPDatabaseManager, readLock)); + else + return Optional.empty(); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java index 7df7af684a..1ae19f0c0c 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java @@ -5,138 +5,78 @@ package org.opensearch.dataprepper.plugins.processor.extension; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoDataFactory; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.S3DBService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.HttpDBDownloadService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.LocalDBDownloadService; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.utils.DbSourceIdentification; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.File; -import java.net.InetAddress; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.Objects; +import java.time.Instant; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * Implementation class of geoIP-processor plugin service class. * It is responsible for calling of mmdb files download */ public class GeoIPProcessorService { - private static final Logger LOG = LoggerFactory.getLogger(GeoIPProcessorService.class); - public static final String DATABASE_1 = "first_database_path"; - public static final String DATABASE_2 = "second_database_path"; - private static final String TEMP_PATH_FOLDER = "geoip"; - private final GeoDataFactory geoDataFactory; - private GetGeoData geoData; - private List databasePaths; - private final String tempPath; - private final ScheduledExecutorService scheduledExecutorService; - private final DBSourceOptions dbSourceOptions; private final MaxMindConfig maxMindConfig; - private final LicenseTypeCheck licenseTypeCheck; - public static volatile boolean downloadReady; - private boolean toggle; - private String flipDatabase; - private boolean isDuringInitialization; + private final GeoIPDatabaseManager geoIPDatabaseManager; + private final ReentrantReadWriteLock.ReadLock readLock; + private Instant nextUpdateAt; + private ExecutorService executorService = null; /** * GeoIPProcessorService constructor for initialization of required attributes * * @param geoIpServiceConfig geoIpServiceConfig */ - public GeoIPProcessorService(final GeoIpServiceConfig geoIpServiceConfig) { - this.toggle = false; + public GeoIPProcessorService(final GeoIpServiceConfig geoIpServiceConfig, + final GeoIPDatabaseManager geoIPDatabaseManager, + final ReentrantReadWriteLock.ReadLock readLock + ) { this.maxMindConfig = geoIpServiceConfig.getMaxMindConfig(); - this.databasePaths = maxMindConfig.getDatabasePaths(); - this.isDuringInitialization = true; - flipDatabase = DATABASE_1; - - licenseTypeCheck = new LicenseTypeCheck(); - geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - this.tempPath = System.getProperty("java.io.tmpdir") + File.separator + TEMP_PATH_FOLDER; + this.geoIPDatabaseManager = geoIPDatabaseManager; + this.readLock = readLock; - dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePaths); - final Duration checkInterval = Objects.requireNonNull(maxMindConfig.getDatabaseRefreshInterval()); - scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); - scheduledExecutorService - .scheduleAtFixedRate(this::downloadThroughURLandS3, 0L, checkInterval.toSeconds(), TimeUnit.SECONDS); + geoIPDatabaseManager.initiateDatabaseDownload(); + this.nextUpdateAt = Instant.now().plus(maxMindConfig.getDatabaseRefreshInterval()); + } - synchronized (this) { - try { - while (!downloadReady) { - wait(); - } - } catch (final InterruptedException ex) { - LOG.info("Thread interrupted while waiting for download to complete: {0}", ex); - Thread.currentThread().interrupt(); - } - if (downloadReady) { - geoData = geoDataFactory.create(flipDatabase); - } + public GeoIPDatabaseReader getGeoIPDatabaseReader() { + readLock.lock(); + try { + final GeoIPDatabaseReader geoIPDatabaseReader = geoIPDatabaseManager.getGeoIPDatabaseReader(); + geoIPDatabaseReader.retain(); + checkAndUpdateDatabases(); + return geoIPDatabaseReader; + } finally { + readLock.unlock(); } - downloadReady = false; } - /** - * Calling download method based on the database path type - */ - public synchronized void downloadThroughURLandS3() { - DBSource dbSource; - toggle = !toggle; - if (!toggle) { - flipDatabase = DATABASE_1; - } else { - flipDatabase = DATABASE_2; + private synchronized void checkAndUpdateDatabases() { + if (nextUpdateAt.isBefore(Instant.now())) { + nextUpdateAt = Instant.now().plus(maxMindConfig.getDatabaseRefreshInterval()); + LOG.info("Trying to update geoip Database readers"); + executorService = Executors.newSingleThreadExecutor(); + executorService.execute(geoIPDatabaseManager::updateDatabaseReader); + executorService.shutdown(); } + } - try { - switch (dbSourceOptions) { - case URL: - dbSource = new HttpDBDownloadService(flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - case S3: - dbSource = new S3DBService(maxMindConfig.getAwsAuthenticationOptionsConfig(), flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - case PATH: - dbSource = new LocalDBDownloadService(flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - } - } catch (final Exception ex) { - if (isDuringInitialization) { - throw new DownloadFailedException("Download failed due to: " + ex); - } else { - LOG.error("Download failed due to: {0}. Using previously loaded database files.", ex); + public void shutdown() { + if (executorService != null) { + try { + if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) { + executorService.shutdownNow(); + } + } catch (final InterruptedException e) { + executorService.shutdownNow(); } } - isDuringInitialization = false; - notifyAll(); - } - - /** - * Method to call enrichment of data based on license type - * @param inetAddress inetAddress - * @param attributes attributes - * @return Enriched Map - */ - public Map getGeoData(final InetAddress inetAddress, final List attributes) { - return geoData.getGeoData(inetAddress, attributes, tempPath + File.separator + flipDatabase); + geoIPDatabaseManager.deleteDatabasesOnShutdown(); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java index 200541339d..6f3b9d9b0d 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java @@ -9,18 +9,35 @@ import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.plugin.ExtensionPlugin; import org.opensearch.dataprepper.model.plugin.ExtensionPoints; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; -@DataPrepperExtensionPlugin(modelType = GeoIpServiceConfig.class, rootKeyJsonPath = "/geoip_service") +import java.util.concurrent.locks.ReentrantReadWriteLock; + +@DataPrepperExtensionPlugin(modelType = GeoIpServiceConfig.class, rootKeyJsonPath = "/geoip_service", allowInPipelineConfigurations = true) public class GeoIpConfigExtension implements ExtensionPlugin { private final DefaultGeoIpConfigSupplier defaultGeoIpConfigSupplier; @DataPrepperPluginConstructor public GeoIpConfigExtension(final GeoIpServiceConfig geoIpServiceConfig) { - this.defaultGeoIpConfigSupplier = new DefaultGeoIpConfigSupplier(geoIpServiceConfig != null ? geoIpServiceConfig : new GeoIpServiceConfig()); + final ReentrantReadWriteLock reentrantReadWriteLock = new ReentrantReadWriteLock(true); + GeoIPDatabaseManager geoIPDatabaseManager = null; + if (geoIpServiceConfig != null) { + geoIPDatabaseManager = new GeoIPDatabaseManager( + geoIpServiceConfig.getMaxMindConfig(), + new LicenseTypeCheck(), + new DatabaseReaderBuilder(), + new GeoIPFileManager(), + reentrantReadWriteLock.writeLock() + ); + } + this.defaultGeoIpConfigSupplier = new DefaultGeoIpConfigSupplier(geoIpServiceConfig, geoIPDatabaseManager, reentrantReadWriteLock.readLock()); } @Override public void apply(final ExtensionPoints extensionPoints) { - extensionPoints.addExtensionProvider(new GeoIpConfigProvider(this.defaultGeoIpConfigSupplier)); + extensionPoints.addExtensionProvider(new GeoIpConfigProvider(this.defaultGeoIpConfigSupplier)); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java index a754c9a745..8e00ef6f3e 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import java.util.Optional; + /** * Interface for supplying {@link GeoIPProcessorService} to {@link GeoIpConfigExtension} * @@ -16,5 +18,5 @@ public interface GeoIpConfigSupplier { * * @since 2.7 */ - GeoIPProcessorService getGeoIPProcessorService(); + Optional getGeoIPProcessorService(); } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java index c311f773ca..3dc974e698 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java @@ -11,12 +11,8 @@ public class GeoIpServiceConfig { private static final MaxMindConfig DEFAULT_MAXMIND_CONFIG = new MaxMindConfig(); - public GeoIpServiceConfig() { - // This default constructor is used if geoip_service is not configured - } - - @JsonProperty("maxmind") @Valid + @JsonProperty("maxmind") private MaxMindConfig maxMindConfig = DEFAULT_MAXMIND_CONFIG; /** diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java index 3fea586e67..b03ec864d4 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java @@ -11,12 +11,17 @@ import jakarta.validation.constraints.Min; import org.hibernate.validator.constraints.time.DurationMax; import org.hibernate.validator.constraints.time.DurationMin; +import org.opensearch.dataprepper.plugins.processor.utils.DatabaseSourceIdentification; +import java.io.File; +import java.net.URI; +import java.net.URISyntaxException; import java.time.Duration; import java.util.ArrayList; import java.util.List; public class MaxMindConfig { + private static final boolean DEFAULT_INSECURE = false; private static final String S3_PREFIX = "s3://"; //TODO: Add validations to database paths @@ -24,6 +29,7 @@ public class MaxMindConfig { private static final List DEFAULT_DATABASE_PATHS = new ArrayList<>(); private static final Duration DEFAULT_DATABASE_REFRESH_INTERVAL = Duration.ofDays(7); private static final int DEFAULT_CACHE_SIZE = 4096; + private static final String DEFAULT_DATABASE_DESTINATION = System.getProperty("data-prepper.dir") + File.separator + "data"; @JsonProperty("database_paths") private List databasePaths = DEFAULT_DATABASE_PATHS; @@ -38,17 +44,22 @@ public class MaxMindConfig { //TODO: Add a Max limit on cache size private int cacheSize = DEFAULT_CACHE_SIZE; - //TODO: Add a destination path to store database files - @JsonProperty("aws") @Valid + @JsonProperty("aws") private AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig; + @JsonProperty("insecure") + private boolean insecure = DEFAULT_INSECURE; + + @JsonProperty("database_destination_path") + private String databaseDestination = DEFAULT_DATABASE_DESTINATION; + public MaxMindConfig() { // This default constructor is used if maxmind is not configured } @AssertTrue(message = "aws should be configured if any path in database_paths is S3 bucket path.") - boolean isAwsAuthenticationOptionsRequired() { + public boolean isAwsAuthenticationOptionsRequired() { for (final String databasePath : databasePaths) { if (databasePath.startsWith(S3_PREFIX)) { return awsAuthenticationOptionsConfig != null; @@ -57,6 +68,19 @@ boolean isAwsAuthenticationOptionsRequired() { return true; } + @AssertTrue(message = "database_paths should be https endpoint if using URL and if insecure is set to false") + public boolean isHttpsEndpointOrInsecure() throws URISyntaxException { + if (insecure) { + return true; + } + for (final String databasePath : databasePaths) { + if (DatabaseSourceIdentification.isURL(databasePath)) { + return new URI(databasePath).getScheme().equals("https"); + } + } + return true; + } + /** * Gets the MaxMind database paths * @@ -96,4 +120,14 @@ public int getCacheSize() { public AwsAuthenticationOptionsConfig getAwsAuthenticationOptionsConfig() { return awsAuthenticationOptionsConfig; } + + /** + * Gets the destination folder to store database files + * + * @return The destination folder + * @since 2.7 + */ + public String getDatabaseDestination() { + return databaseDestination + File.separator + "geoip"; + } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java index 3ead66744c..8c30dd80a7 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java @@ -5,17 +5,12 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import javax.net.ssl.TrustManager; import javax.net.ssl.SSLContext; import javax.net.ssl.X509TrustManager; import javax.net.ssl.HttpsURLConnection; import javax.net.ssl.SSLSession; import javax.net.ssl.HostnameVerifier; -import java.io.File; -import java.io.UncheckedIOException; import java.security.KeyManagementException; import java.security.NoSuchAlgorithmException; import java.security.SecureRandom; @@ -24,47 +19,7 @@ import java.util.List; public interface DBSource { - - public static final Logger LOG = LoggerFactory.getLogger(DBSource.class); - public String tempFolderPath = System.getProperty("java.io.tmpdir")+ File.separator +"GeoIP"; - public String tarFolderPath = tempFolderPath + "/tar"; - public String downloadTarFilepath = tarFolderPath + "/out.tar.gz"; - public void initiateDownload(List config) throws Exception; - - /** - * create Folder If Not Exist - * @param outputFilePath Output File Path - * @return File - */ - static File createFolderIfNotExist(String outputFilePath) { - final File destFile = new File(outputFilePath); - try { - if (!destFile.exists()) { - destFile.mkdirs(); - } - } - catch (UncheckedIOException ex) { - LOG.info("Create Folder If NotExist Exception {0}", ex); - } - return destFile; - } - - /** - * Delete Directory - * @param file file - */ - static void deleteDirectory(File file) { - - if (file.exists()) { - for (final File subFile : file.listFiles()) { - if (subFile.isDirectory()) { - deleteDirectory(subFile); - } - subFile.delete(); - } - file.delete(); - } - } + void initiateDownload(List config) throws Exception; /** * initiateSSL diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java index 485274dd9d..60eacff72a 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java @@ -17,7 +17,8 @@ public enum DBSourceOptions { PATH("path"), URL("url"), - S3("s3"); + S3("s3"), + HTTP_MANIFEST("http_manifest"); private final String option; diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java similarity index 80% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java index 47546c6c9f..841b4c7fbc 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java @@ -15,7 +15,9 @@ /** * Implementation class for DatabaseReader Creation */ -public class DatabaseReaderCreate { +public class DatabaseReaderBuilder { + public DatabaseReaderBuilder() { + } /** * Creates DatabaseReader instance based on in memory or cache type @@ -23,7 +25,7 @@ public class DatabaseReaderCreate { * @param cacheSize cacheSize * @return DatabaseReader */ - public static DatabaseReader createLoader(final Path databasePath, final int cacheSize) throws IOException { + public DatabaseReader buildReader(final Path databasePath, final int cacheSize) throws IOException { return new DatabaseReader.Builder(databasePath.toFile()) .fileMode(Reader.FileMode.MEMORY_MAPPED) .withCache(new CHMCache(cacheSize)) diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java deleted file mode 100644 index e05453d37f..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoIP2Data; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoLite2Data; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; - -import java.io.File; - -import static org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource.tempFolderPath; - -public class GeoDataFactory { - private final MaxMindConfig maxMindConfig; - private final LicenseTypeCheck licenseTypeCheck; - - public GeoDataFactory(final MaxMindConfig maxMindConfig, final LicenseTypeCheck licenseTypeCheck) { - this.maxMindConfig = maxMindConfig; - this.licenseTypeCheck = licenseTypeCheck; - } - - /** - * Creates GetGeoData class based on LicenseTypeOptions - */ - public GetGeoData create(final String databasePath) { - final String finalPath = tempFolderPath + File.separator + databasePath; - final LicenseTypeOptions licenseType = licenseTypeCheck.isGeoLite2OrEnterpriseLicense(finalPath); - if (licenseType.equals(LicenseTypeOptions.FREE)) { - return new GetGeoLite2Data(finalPath, maxMindConfig.getCacheSize()); - } else { - return new GetGeoIP2Data(finalPath, maxMindConfig.getCacheSize()); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java new file mode 100644 index 0000000000..4efd3cbdd2 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java @@ -0,0 +1,194 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIP2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoLite2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; +import org.opensearch.dataprepper.plugins.processor.utils.DatabaseSourceIdentification; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.util.List; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock.WriteLock; + +public class GeoIPDatabaseManager { + private static final Logger LOG = LoggerFactory.getLogger(GeoIPDatabaseManager.class); + public static final String FIRST_DATABASE_DIR = "first_database"; + public static final String SECOND_DATABASE_DIR = "second_database"; + private final MaxMindConfig maxMindConfig; + private final LicenseTypeCheck licenseTypeCheck; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final List databasePaths; + private final WriteLock writeLock; + private final int cacheSize; + private final GeoIPFileManager geoIPFileManager; + private DBSourceOptions dbSourceOptions; + private String currentDatabaseDir; + private GeoIPDatabaseReader geoIPDatabaseReader; + private boolean databaseDirToggle; + private boolean downloadReady; + + public GeoIPDatabaseManager(final MaxMindConfig maxMindConfig, + final LicenseTypeCheck licenseTypeCheck, + final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final ReentrantReadWriteLock.WriteLock writeLock + ) { + this.maxMindConfig = maxMindConfig; + this.licenseTypeCheck = licenseTypeCheck; + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.databasePaths = maxMindConfig.getDatabasePaths(); + this.writeLock = writeLock; + this.cacheSize = maxMindConfig.getCacheSize(); + } + + public void initiateDatabaseDownload() { + dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePaths); + try { + downloadDatabases(); + } catch (final Exception e) { + throw new DownloadFailedException(e.getMessage()); + } + synchronized (this) { + try { + while (!downloadReady) { + wait(); + } + } catch (final InterruptedException ex) { + LOG.info("Thread interrupted while waiting for download to complete: {0}", ex); + Thread.currentThread().interrupt(); + } + if (downloadReady) { + try { + geoIPDatabaseReader = createReader(); + } catch (final NoValidDatabaseFoundException e) { + throw new NoValidDatabaseFoundException(e.getMessage()); + } + } + } + downloadReady = false; + } + + public void updateDatabaseReader() { + try { + downloadDatabases(); + } catch (final Exception e) { + LOG.error("Database download failed, using previously loaded database. {}", e.getMessage()); + final File file = new File(currentDatabaseDir); + geoIPFileManager.deleteDirectory(file); + switchDirectory(); + } + + if (downloadReady) { + downloadReady = false; + try { + switchDatabase(); + } catch (final Exception e) { + LOG.error("Failed to update databases. Please make sure the database files exist at configured path amd are valid. " + + "Using previously loaded database. {}", e.getMessage()); + final File file = new File(currentDatabaseDir); + geoIPFileManager.deleteDirectory(file); + switchDirectory(); + } + } + } + + private void switchDatabase() { + writeLock.lock(); + try { + final GeoIPDatabaseReader newGeoipDatabaseReader = createReader(); + final GeoIPDatabaseReader oldGeoipDatabaseReader = geoIPDatabaseReader; + geoIPDatabaseReader = newGeoipDatabaseReader; + oldGeoipDatabaseReader.close(); + } finally { + writeLock.unlock(); + } + } + + private void downloadDatabases() throws Exception { + DBSource dbSource; + switchDirectory(); + + final String destinationPath = maxMindConfig.getDatabaseDestination() + File.separator + currentDatabaseDir; + geoIPFileManager.createDirectoryIfNotExist(destinationPath); + switch (dbSourceOptions) { + case HTTP_MANIFEST: + dbSource = new ManifestDownloadService(destinationPath); + dbSource.initiateDownload(databasePaths); + downloadReady =true; + break; + case URL: + dbSource = new HttpDBDownloadService(destinationPath, geoIPFileManager); + dbSource.initiateDownload(databasePaths); + downloadReady = true; + break; + case S3: + dbSource = new S3DBService(maxMindConfig.getAwsAuthenticationOptionsConfig(), destinationPath); + dbSource.initiateDownload(databasePaths); + downloadReady = true; + break; + case PATH: + dbSource = new LocalDBDownloadService(destinationPath); + dbSource.initiateDownload(databasePaths); + downloadReady = true; + break; + } + } + + private GeoIPDatabaseReader createReader() { + final String finalPath = maxMindConfig.getDatabaseDestination() + File.separator + currentDatabaseDir; + final LicenseTypeOptions licenseType = licenseTypeCheck.isGeoLite2OrEnterpriseLicense(finalPath); + if (licenseType == null) { + throw new NoValidDatabaseFoundException("At least one valid database is required."); + } + GeoIPDatabaseReader newGeoIPDatabaseReader; + if (licenseType.equals(LicenseTypeOptions.FREE)) { + newGeoIPDatabaseReader = new GeoLite2DatabaseReader(databaseReaderBuilder, geoIPFileManager, finalPath, cacheSize); + } else { + newGeoIPDatabaseReader = new GeoIP2DatabaseReader(databaseReaderBuilder, geoIPFileManager, finalPath, cacheSize); + } + return newGeoIPDatabaseReader; + } + + private void switchDirectory() { + databaseDirToggle = !databaseDirToggle; + if (databaseDirToggle) { + currentDatabaseDir = FIRST_DATABASE_DIR; + } else { + currentDatabaseDir = SECOND_DATABASE_DIR; + } + } + + public GeoIPDatabaseReader getGeoIPDatabaseReader() { + return geoIPDatabaseReader; + } + + public void deleteDatabasesOnShutdown() { + geoIPFileManager.deleteDirectory(new File(maxMindConfig.getDatabaseDestination() + File.separator + FIRST_DATABASE_DIR)); + geoIPFileManager.deleteDirectory(new File(maxMindConfig.getDatabaseDestination() + File.separator + SECOND_DATABASE_DIR)); + } + + public void deleteDirectory(final File file) { + + if (file.exists()) { + for (final File subFile : file.listFiles()) { + if (subFile.isDirectory()) { + deleteDirectory(subFile); + } + subFile.delete(); + } + file.delete(); + } + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java new file mode 100644 index 0000000000..66dda2db70 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import java.io.File; + +public class GeoIPFileManager { + public void deleteDirectory(final File file) { + + if (file.exists()) { + for (final File subFile : file.listFiles()) { + if (subFile.isDirectory()) { + deleteDirectory(subFile); + } + subFile.delete(); + } + file.delete(); + } + } + + public void deleteFile(final File file) { + file.delete(); + } + + public void createDirectoryIfNotExist(final String outputFilePath) { + final File destFile = new File(outputFilePath); + if (!destFile.exists()) { + destFile.mkdirs(); + } + } + +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java index d407f9bdb7..d28f6ab67e 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java @@ -26,15 +26,17 @@ public class HttpDBDownloadService implements DBSource { private static final Logger LOG = LoggerFactory.getLogger(HttpDBDownloadService.class); - private final String prefixDir; + private final String destinationDirectory; private static final int DEFAULT_BYTE_SIZE = 1024; + private final GeoIPFileManager geoIPFileManager; /** * HttpDBDownloadService constructor for initialisation of attributes - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ - public HttpDBDownloadService(String prefixDir) { - this.prefixDir = prefixDir; + public HttpDBDownloadService(final String destinationDirectory, final GeoIPFileManager geoIPFileManager) { + this.destinationDirectory = destinationDirectory; + this.geoIPFileManager = geoIPFileManager; } /** @@ -42,14 +44,15 @@ public HttpDBDownloadService(String prefixDir) { * @param urlList urlList */ public void initiateDownload(List urlList) { - final File tmpDir = DBSource.createFolderIfNotExist(tempFolderPath + File.separator + prefixDir); - for(String url : urlList) { - DBSource.createFolderIfNotExist(tarFolderPath); + final String tarDir = destinationDirectory + File.separator + "tar"; + final String downloadTarFilepath = tarDir + File.separator + "out.tar.gz"; + for(final String url : urlList) { + geoIPFileManager.createDirectoryIfNotExist(tarDir); try { initiateSSL(); - buildRequestAndDownloadFile(url); - decompressAndUntarFile(tarFolderPath, downloadTarFilepath, tmpDir); - deleteTarFolder(tarFolderPath); + buildRequestAndDownloadFile(url, downloadTarFilepath); + decompressAndUntarFile(tarDir, downloadTarFilepath, new File(destinationDirectory)); + deleteTarFolder(tarDir); } catch (Exception ex) { LOG.info("InitiateDownload Exception {0} " , ex); } @@ -62,7 +65,7 @@ public void initiateDownload(List urlList) { * @param downloadTarFilepath downloadTarFilepath * @param tmpDir tmpDir */ - private void decompressAndUntarFile(String tarFolderPath, String downloadTarFilepath, File tmpDir) { + private void decompressAndUntarFile(final String tarFolderPath, final String downloadTarFilepath, final File tmpDir) { try { final File inputFile = new File(downloadTarFilepath); final String outputFile = getFileName(inputFile, tarFolderPath); @@ -80,17 +83,17 @@ private void decompressAndUntarFile(String tarFolderPath, String downloadTarFile * Build Request And DownloadFile * @param url url */ - public void buildRequestAndDownloadFile(String... url) { - downloadDBFileFromMaxmind(url[0], downloadTarFilepath); + public void buildRequestAndDownloadFile(final String url, final String downloadTarFilepath) { + downloadDBFileFromMaxmind(url, downloadTarFilepath); } /** * Delete Tar Folder * @param tarFolder Tar Folder */ - private static void deleteTarFolder(String tarFolder) { + private void deleteTarFolder(String tarFolder) { final File file = new File(tarFolder); - DBSource.deleteDirectory(file); + geoIPFileManager.deleteDirectory(file); if (file.exists()) { file.delete(); } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java index d4a85218fa..ae55c2343c 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java @@ -15,14 +15,14 @@ */ public class LocalDBDownloadService implements DBSource { - private final String prefixDir; + private final String destinationDirectory; /** * LocalDBDownloadService constructor for initialisation of attributes - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ - public LocalDBDownloadService(final String prefixDir) { - this.prefixDir = prefixDir; + public LocalDBDownloadService(final String destinationDirectory) { + this.destinationDirectory = destinationDirectory; } /** @@ -31,10 +31,8 @@ public LocalDBDownloadService(final String prefixDir) { */ @Override public void initiateDownload(List config) throws Exception { - String destPath = tempFolderPath + File.separator + prefixDir; - DBSource.createFolderIfNotExist(destPath); File srcDatabaseConfigPath = new File(config.get(0)); - File destDatabaseConfigPath = new File(destPath); + File destDatabaseConfigPath = new File(destinationDirectory); FileUtils.copyDirectory(srcDatabaseConfigPath, destDatabaseConfigPath); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java new file mode 100644 index 0000000000..f29bd3c2fb --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class Manifest { + @JsonProperty("url") + private String url; + @JsonProperty("db_name") + private String dbName; + @JsonProperty("sha256_hash") + private String sha256Hash; + @JsonProperty("valid_for_in_days") + private int validForInDays; + @JsonProperty("updated_at_in_epoch_milli") + private long updatedAt; + @JsonProperty("provider") + private String provider; + + public String getUrl() { + return url; + } + + public String getDbName() { + return dbName; + } + + public String getSha256Hash() { + return sha256Hash; + } + + public int getValidForInDays() { + return validForInDays; + } + + public Long getUpdatedAt() { + return updatedAt; + } + + public String getProvider() { + return provider; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java new file mode 100644 index 0000000000..6cc12acee6 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java @@ -0,0 +1,162 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; + +import java.io.BufferedInputStream; +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URL; +import java.util.List; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; + +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.MAXMIND_DATABASE_EXTENSION; + + +public class ManifestDownloadService implements DBSource { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final int DEFAULT_BYTE_SIZE = 1024; + private static final String CITY_DATABASE_NAME = "geolite2-city.mmdb"; + private static final String COUNTRY_DATABASE_NAME = "geolite2-country.mmdb"; + private static final String ASN_DATABASE_NAME = "geolite2-asn.mmdb"; + private static final String ZIP_FILE_EXTENSION = ".zip"; + private final String directoryName; + + private Manifest cityManifest; + private Manifest countryManifest; + private Manifest asnManifest; + + public ManifestDownloadService(final String directoryName) { + this.directoryName = directoryName; + } + + @Override + public void initiateDownload(final List databasePaths) { + for(final String url : databasePaths) { + final Manifest manifest = deserializeManifestFile(url); + populateManifestObjects(manifest); + + final String manifestFilePath = manifest.getDbName(); + final String zipFileName = manifestFilePath.substring(0, manifestFilePath.lastIndexOf(".")).concat(ZIP_FILE_EXTENSION); + final String zipFilePath = directoryName + File.separator + zipFileName; + + downloadZipFile(manifest.getUrl(), zipFilePath); + unzipDownloadedFile(zipFilePath, directoryName); + } + } + + private Manifest deserializeManifestFile(final String CDNEndpoint) { + HttpURLConnection httpURLConnection = null; + try { + final URL url = new URL(CDNEndpoint); + httpURLConnection = (HttpURLConnection) url.openConnection(); + httpURLConnection.addRequestProperty("User-Agent", "Custom-User-Agent"); + + final Manifest manifest = OBJECT_MAPPER.readValue(httpURLConnection.getInputStream(), Manifest.class); + httpURLConnection.disconnect(); + + return manifest; + } catch (final IOException ex) { + if (httpURLConnection != null) { + httpURLConnection.disconnect(); + } + throw new DownloadFailedException("Exception occurred while reading manifest.json file due to: " + ex); + } + } + + private void populateManifestObjects(final Manifest manifest) { + switch (manifest.getDbName()) { + case CITY_DATABASE_NAME: + cityManifest = manifest; + break; + case COUNTRY_DATABASE_NAME: + countryManifest = manifest; + break; + case ASN_DATABASE_NAME: + asnManifest = manifest; + break; + } + } + + private void downloadZipFile(final String databaseUrl, final String destinationPath) { + HttpURLConnection httpURLConnection; + try { + final URL url = new URL(databaseUrl); + httpURLConnection = (HttpURLConnection) url.openConnection(); + // CDN endpoint returns 403 without User Agent. + httpURLConnection.addRequestProperty("User-Agent", "Data Prepper"); + } catch (IOException ex) { + throw new DownloadFailedException("Exception occurred while opening connection due to: " + ex.getMessage()); + } + + try (final BufferedInputStream in = new BufferedInputStream(httpURLConnection.getInputStream()); + final FileOutputStream fileOutputStream = new FileOutputStream(destinationPath)) { + final byte[] dataBuffer = new byte[DEFAULT_BYTE_SIZE]; + int bytesRead; + while ((bytesRead = in.read(dataBuffer, 0, DEFAULT_BYTE_SIZE)) != -1) { + fileOutputStream.write(dataBuffer, 0, bytesRead); + } + httpURLConnection.disconnect(); + } catch (final IOException ex) { + httpURLConnection.disconnect(); + throw new DownloadFailedException("Exception occurred while downloading MaxMind database due to: " + ex.getMessage()); + } + } + + private void unzipDownloadedFile(final String zipFilePath, final String outputFilePath) { + final File inputFile = new File(zipFilePath); + final File outputDir = new File(outputFilePath); + if (!outputDir.exists()) { + outputDir.mkdirs(); + } + + final byte[] buffer = new byte[DEFAULT_BYTE_SIZE]; + + try (final FileInputStream fileInputStream = new FileInputStream(inputFile); + final ZipInputStream zipInputStream = new ZipInputStream(fileInputStream)) { + + ZipEntry zipEntry = zipInputStream.getNextEntry(); + while (zipEntry != null && zipEntry.getName().endsWith(MAXMIND_DATABASE_EXTENSION)) { + final String fileName = zipEntry.getName(); + final File newFile = new File(outputDir + File.separator + fileName); + + final FileOutputStream fileOutputStream = new FileOutputStream(newFile); + int len; + while ((len = zipInputStream.read(buffer)) > 0) { + fileOutputStream.write(buffer, 0, len); + } + fileOutputStream.close(); + zipInputStream.closeEntry(); + zipEntry = zipInputStream.getNextEntry(); + } + + // deleting zip file after unzipping + inputFile.delete(); + } catch (final IOException e) { + inputFile.delete(); + throw new DownloadFailedException("Exception occurred while unzipping the database file due to: " + e.getMessage()); + } + } + + + public Manifest getCityManifest() { + return cityManifest; + } + + public Manifest getCountryManifest() { + return countryManifest; + } + + public Manifest getAsnManifest() { + return asnManifest; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java index 11793a8189..144facd034 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java @@ -5,7 +5,7 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; import org.opensearch.dataprepper.plugins.processor.extension.AwsAuthenticationOptionsConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,7 +14,6 @@ import software.amazon.awssdk.transfer.s3.model.DirectoryDownload; import software.amazon.awssdk.transfer.s3.model.DownloadDirectoryRequest; -import java.io.File; import java.net.URI; import java.net.URISyntaxException; import java.nio.file.Paths; @@ -29,30 +28,29 @@ public class S3DBService implements DBSource { private String bucketName; private String bucketPath; private final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig; - private final String prefixDir; + private final String destinationDirectory; /** * S3DBService constructor for initialisation of attributes * - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ public S3DBService(final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig, - final String prefixDir) { + final String destinationDirectory) { this.awsAuthenticationOptionsConfig = awsAuthenticationOptionsConfig; - this.prefixDir = prefixDir; + this.destinationDirectory = destinationDirectory; } /** * Initialisation of Download through Url * @param s3URLs s3URLs */ - public void initiateDownload(List s3URLs) { + public void initiateDownload(final List s3URLs) { for (String s3Url : s3URLs) { try { URI uri = new URI(s3Url); bucketName = uri.getHost(); bucketPath = removeTrailingSlash(removeLeadingSlash(uri.getPath())); - DBSource.createFolderIfNotExist(tempFolderPath + File.separator + prefixDir); buildRequestAndDownloadFile(bucketName, bucketPath); } catch (URISyntaxException ex) { LOG.info("Initiate Download Exception", ex); @@ -96,7 +94,7 @@ public void buildRequestAndDownloadFile(String... path) { DirectoryDownload directoryDownload = transferManager.downloadDirectory( DownloadDirectoryRequest.builder() - .destination(Paths.get(tempFolderPath + File.separator + prefixDir)) + .destination(Paths.get(destinationDirectory)) .bucket(path[0]) .listObjectsV2RequestTransformer(l -> l.prefix(path[1])) .build()); diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java similarity index 52% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java index e7eaffe75d..1afa8fda88 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; +import java.io.File; import java.net.MalformedURLException; import java.net.URI; import java.net.URISyntaxException; @@ -17,20 +18,21 @@ /** * Implementation of class for checking whether URL type is S3 or file path */ -public class DbSourceIdentification { +public class DatabaseSourceIdentification { - private DbSourceIdentification() { + private DatabaseSourceIdentification() { } - private static String s3DomainPattern = "[a-zA-Z0-9-]+\\.s3\\.amazonaws\\.com"; + private static final String S3_DOMAIN_PATTERN = "[a-zA-Z0-9-]+\\.s3\\.amazonaws\\.com"; + private static final String MANIFEST_ENDPOINT_PATH = "manifest.json"; /** * Check for database path is valid S3 URI or not * @param uriString uriString * @return boolean */ - public static boolean isS3Uri(String uriString) { + public static boolean isS3Uri(final String uriString) { try { URI uri = new URI(uriString); if (uri.getScheme() != null && uri.getScheme().equalsIgnoreCase("s3")) { @@ -43,44 +45,48 @@ public static boolean isS3Uri(String uriString) { } /** - * Check for database path is valid S3 URL or not - * @param urlString urlString + * Check for database path is valid URL or not + * @param input input * @return boolean */ - public static boolean isS3Url(String urlString) { + public static boolean isURL(final String input) { try { - URL url = new URL(urlString); - if (Pattern.matches(s3DomainPattern, url.getHost())) { - return true; - } - } catch (Exception e) { + final URI uri = new URI(input); + final URL url = new URL(input); + return !input.endsWith(MANIFEST_ENDPOINT_PATH) && + uri.getScheme() != null && + !Pattern.matches(S3_DOMAIN_PATTERN, url.getHost()) && + (uri.getScheme().equals("http") || uri.getScheme().equals("https")); + } catch (URISyntaxException | MalformedURLException e) { return false; } - return false; } /** - * Check for database path is valid URL or not + * Check for database path is local file path or not * @param input input * @return boolean */ - public static boolean isURL(String input) { - try { - URI uri = new URI(input); - URL url = new URL(input); - return uri.getScheme() != null && !Pattern.matches(s3DomainPattern, url.getHost()) &&(uri.getScheme().equals("http") || uri.getScheme().equals("https")); - } catch (URISyntaxException | MalformedURLException e) { - return false; - } + public static boolean isFilePath(final String input) { + final File file = new File(input); + return file.exists() && file.isDirectory(); } /** - * Check for database path is local file path or not + * Check for database path is CDN endpoint * @param input input * @return boolean */ - public static boolean isFilePath(String input) { - return input.startsWith("/") || input.startsWith("./") || input.startsWith("\\") || (input.length() > 1 && input.charAt(1) == ':'); + public static boolean isCDNEndpoint(final String input) { + if (input.endsWith(MANIFEST_ENDPOINT_PATH)) { + try { + final URI uri = new URI(input); + return uri.getScheme().equals("http") || uri.getScheme().equals("https"); + } catch (final URISyntaxException e) { + return false; + } + } + return false; } /** @@ -88,18 +94,21 @@ public static boolean isFilePath(String input) { * @param databasePaths - List of database paths to get databases data from * @return DBSourceOptions */ - public static DBSourceOptions getDatabasePathType(List databasePaths) { + public static DBSourceOptions getDatabasePathType(final List databasePaths) { DBSourceOptions downloadSourceOptions = null; for(final String databasePath : databasePaths) { - if(DbSourceIdentification.isFilePath(databasePath)) { + if(DatabaseSourceIdentification.isFilePath(databasePath)) { return DBSourceOptions.PATH; } - else if(DbSourceIdentification.isURL(databasePath)) + else if (DatabaseSourceIdentification.isCDNEndpoint(databasePath)) { + downloadSourceOptions = DBSourceOptions.HTTP_MANIFEST; + } + else if(DatabaseSourceIdentification.isURL(databasePath)) { downloadSourceOptions = DBSourceOptions.URL; } - else if(DbSourceIdentification.isS3Uri(databasePath) || (DbSourceIdentification.isS3Url(databasePath))) + else if(DatabaseSourceIdentification.isS3Uri(databasePath)) { downloadSourceOptions = DBSourceOptions.S3; } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java index d717ebc711..dceb1bc591 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.utils; +import org.opensearch.dataprepper.plugins.processor.exception.InvalidIPAddressException; + import java.net.Inet4Address; import java.net.Inet6Address; import java.net.InetAddress; @@ -20,10 +22,15 @@ public class IPValidationCheck { * Check for IP is valid or not * @param ipAddress ipAddress * @return boolean - * @throws UnknownHostException UnknownHostException + * @throws InvalidIPAddressException InvalidIPAddressException */ - public static boolean isPublicIpAddress(final String ipAddress) throws UnknownHostException { - InetAddress address = InetAddress.getByName(ipAddress); + public static boolean isPublicIpAddress(final String ipAddress) { + InetAddress address; + try { + address = InetAddress.getByName(ipAddress); + } catch (final UnknownHostException e) { + throw new InvalidIPAddressException(e.getMessage()); + } if (address instanceof Inet6Address || address instanceof Inet4Address) { return !address.isSiteLocalAddress() && !address.isLoopbackAddress(); } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java index 1e75d4c4ca..907225a4d8 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java @@ -10,23 +10,19 @@ import java.io.File; import java.nio.file.Path; import java.nio.file.Paths; -import java.util.Arrays; -import java.util.List; -/**cls +/** * * Implementation of class logic to check maxmind database id free or enterprise version */ public class LicenseTypeCheck { - - protected static final String[] DEFAULT_DATABASE_FILENAMES = new String[] { "GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb" }; - protected static final List geoLite2Database = Arrays.asList(DEFAULT_DATABASE_FILENAMES); - private static final String geoIP2EnterpriseDB = "GeoIP2-Enterprise.mmdb"; + private static final String GEOIP2_DATABASE = "geoip2"; + private static final String GEOLITE2_DATABASE = "geolite2"; + private static final String MMDB = "mmdb"; public LicenseTypeCheck() { - } /** @@ -36,31 +32,30 @@ public LicenseTypeCheck() { * @return license type options */ public LicenseTypeOptions isGeoLite2OrEnterpriseLicense(final String databasePath) { - LicenseTypeOptions licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; - File directory = new File(databasePath); - // list all files present in the directory - File[] files = directory.listFiles(); - - for(File file : files) { - // convert the file name into string - String fileName = file.toString(); - - int index = fileName.lastIndexOf('.'); - if(index > 0) { - String extension = fileName.substring(index + 1); - Path onlyFileName = Paths.get(fileName).getFileName(); - - if((extension.equals("mmdb")) && (geoIP2EnterpriseDB.equals(onlyFileName.toString()))) { - licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; - break; - } - else if((extension.equals("mmdb")) && (geoLite2Database.contains(onlyFileName.toString()))) - { - licenseTypeOptions = LicenseTypeOptions.FREE; + LicenseTypeOptions licenseTypeOptions = null; + final File directory = new File(databasePath); + if (directory.isDirectory()) { + // list all files present in the directory + final File[] files = directory.listFiles(); + + for (final File file : files) { + // convert the file name into string + final String fileName = file.toString(); + + int index = fileName.lastIndexOf('.'); + if (index > 0) { + String extension = fileName.substring(index + 1); + Path onlyFileName = Paths.get(fileName).getFileName(); + + if ((extension.equals(MMDB)) && (onlyFileName.toString().toLowerCase().contains(GEOIP2_DATABASE))) { + licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; + break; + } else if ((extension.equals(MMDB)) && (onlyFileName.toString().toLowerCase().contains(GEOLITE2_DATABASE))) { + licenseTypeOptions = LicenseTypeOptions.FREE; + } } } } return licenseTypeOptions; } } - diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java new file mode 100644 index 0000000000..4a4896304c --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java @@ -0,0 +1,26 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +import org.junit.jupiter.api.Test; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; + +class GeoIPFieldTest { + + @Test + void test_findByName_should_return_geoip_field_if_valid() { + final GeoIPField geoIPField = GeoIPField.findByName("city_name"); + assertThat(geoIPField, equalTo(GeoIPField.CITY_NAME)); + } + + @Test + void test_findByName_should_return_null_if_invalid() { + final GeoIPField geoIPField = GeoIPField.findByName("coordinates"); + assertThat(geoIPField, equalTo(null)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java index 55d28b566b..1ff06cfc99 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java @@ -10,7 +10,10 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; @@ -19,30 +22,59 @@ import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; import org.opensearch.dataprepper.plugins.processor.extension.GeoIpConfigSupplier; -import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import org.opensearch.dataprepper.plugins.processor.utils.IPValidationCheck; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.stream.Collectors; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LATITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LONGITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_FAILED_LOOKUP; import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_PROCESSED; @@ -66,10 +98,15 @@ class GeoIPProcessorTest { private Counter geoIpEventsProcessed; @Mock private Counter geoIpEventsFailedLookup; + @Mock + private GeoIPDatabaseReader geoIPDatabaseReader; + @Captor + private ArgumentCaptor> geoIPFieldCaptor; @BeforeEach void setUp() { - when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(geoIPProcessorService); + when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(Optional.of(geoIPProcessorService)); + lenient().when(geoIPProcessorService.getGeoIPDatabaseReader()).thenReturn(geoIPDatabaseReader); lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_PROCESSED)).thenReturn(geoIpEventsProcessed); lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_FAILED_LOOKUP)).thenReturn(geoIpEventsFailedLookup); } @@ -85,20 +122,18 @@ private GeoIPProcessor createObjectUnderTest() { } @Test - void doExecuteTest_with_when_condition_should_only_enrich_events_that_match_when_condition() throws NoSuchFieldException, IllegalAccessException { + void doExecuteTest_with_when_condition_should_only_enrich_events_that_match_when_condition() { final String whenCondition = "/peer/status == success"; when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(geoIPProcessorConfig.getWhenCondition()).thenReturn(whenCondition); when(entry.getSource()).thenReturn("/peer/ip"); when(entry.getTarget()).thenReturn(TARGET); - when(entry.getFields()).thenReturn(setFields()); + when(entry.getIncludeFields()).thenReturn(setFields()); final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - when(geoIPProcessorService.getGeoData(any(), any())).thenReturn(prepareGeoData()); - - ReflectivelySetField.setField(GeoIPProcessor.class, geoIPProcessor, "geoIPProcessorService", geoIPProcessorService); + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); final Record record1 = createCustomRecord("success"); final Record record2 = createCustomRecord("failed"); @@ -126,38 +161,154 @@ void doExecuteTest_with_when_condition_should_only_enrich_events_that_match_when } @Test - void doExecuteTest_should_add_geo_data_to_event_if_source_is_non_null() throws NoSuchFieldException, IllegalAccessException { + void doExecuteTest_should_add_geo_data_to_event_if_source_is_non_null() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getTarget()).thenReturn(TARGET); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); + assertThat(event.containsKey(TARGET), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + } + } + + @Test + void doExecuteTest_should_add_geo_data_with_expected_fields_to_event_when_include_fields_is_configured() { when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(entry.getSource()).thenReturn(SOURCE); when(entry.getTarget()).thenReturn(TARGET); - when(entry.getFields()).thenReturn(setFields()); + + final List includeFields = List.of("city_name", "asn"); + final List includeFieldsResult = List.of(GeoIPField.CITY_NAME, GeoIPField.ASN); + when(entry.getIncludeFields()).thenReturn(includeFields); final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - when(geoIPProcessorService.getGeoData(any(), any())).thenReturn(prepareGeoData()); - ReflectivelySetField.setField(GeoIPProcessor.class, geoIPProcessor, - "geoIPProcessorService", geoIPProcessorService); + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); Collection> records = geoIPProcessor.doExecute(setEventQueue()); + verify(geoIPDatabaseReader).getGeoData(any(), geoIPFieldCaptor.capture(), any()); + for (final Record record : records) { final Event event = record.getData(); assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); - assertThat(event.containsKey("geolocation"), equalTo(true)); + assertThat(event.containsKey(TARGET), equalTo(true)); verify(geoIpEventsProcessed).increment(); } + + final List value = geoIPFieldCaptor.getValue(); + assertThat(value, containsInAnyOrder(includeFieldsResult.toArray())); } @Test - void doExecuteTest_should_not_add_geo_data_to_event_if_source_is_null() throws NoSuchFieldException, IllegalAccessException { + void doExecuteTest_should_add_geo_data_with_expected_fields_to_event_when_exclude_fields_is_configured() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getTarget()).thenReturn(TARGET); + + final List excludeFields = List.of("country_name", "country_iso_code", "city_name", "asn", "asn_organization", "network", "ip"); + final List excludeFieldsResult = List.of(CONTINENT_NAME, CONTINENT_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE, LOCATION, LOCATION_ACCURACY_RADIUS, LATITUDE, LONGITUDE, METRO_CODE, TIME_ZONE, POSTAL_CODE, + MOST_SPECIFIED_SUBDIVISION_NAME, MOST_SPECIFIED_SUBDIVISION_ISO_CODE, LEAST_SPECIFIED_SUBDIVISION_NAME, + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE, COUNTRY_CONFIDENCE, CITY_CONFIDENCE, MOST_SPECIFIED_SUBDIVISION_CONFIDENCE, + LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE, POSTAL_CODE_CONFIDENCE); + when(entry.getExcludeFields()).thenReturn(excludeFields); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + verify(geoIPDatabaseReader).getGeoData(any(), geoIPFieldCaptor.capture(), any()); + + for (final Record record : records) { + final Event event = record.getData(); + assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); + assertThat(event.containsKey(TARGET), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + } + + final List value = geoIPFieldCaptor.getValue(); + assertThat(value, containsInAnyOrder(excludeFieldsResult.toArray())); + } + + @Test + void doExecuteTest_should_not_add_geo_data_to_event_if_source_is_null() { when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(entry.getSource()).thenReturn("ip"); - when(entry.getFields()).thenReturn(setFields()); + when(entry.getIncludeFields()).thenReturn(setFields()); final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - ReflectivelySetField.setField(GeoIPProcessor.class, geoIPProcessor, - "geoIPProcessorService", geoIPProcessorService); Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + } + + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailedLookup).increment(); + } + + @Test + void doExecuteTest_should_not_add_geo_data_to_event_if_returned_data_is_empty() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(Collections.EMPTY_MAP); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + } + + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailedLookup).increment(); + } + + @Test + void doExecuteTest_should_not_add_geodata_if_database_is_expired() { + when(geoIPDatabaseReader.isExpired()).thenReturn(true); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + final Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + } + + verify(geoIpEventsProcessed).increment(); + } + + @Test + void doExecuteTest_should_not_add_geodata_if_ip_address_is_not_public() { + try (final MockedStatic ipValidationCheckMockedStatic = mockStatic(IPValidationCheck.class)) { + ipValidationCheckMockedStatic.when(() -> IPValidationCheck.isPublicIpAddress(any())).thenReturn(false); + } + + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + final Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + } + verify(geoIpEventsProcessed).increment(); verify(geoIpEventsFailedLookup).increment(); } @@ -165,17 +316,15 @@ void doExecuteTest_should_not_add_geo_data_to_event_if_source_is_null() throws N @Test void test_tags_when_enrich_fails() { when(entry.getSource()).thenReturn(SOURCE); - when(entry.getFields()).thenReturn(setFields()); + when(entry.getIncludeFields()).thenReturn(setFields()); List testTags = List.of("tag1", "tag2"); when(geoIPProcessorConfig.getTagsOnFailure()).thenReturn(testTags); when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); - when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(geoIPProcessorService); - GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - doThrow(EnrichFailedException.class).when(geoIPProcessorService).getGeoData(any(), any()); + doThrow(EnrichFailedException.class).when(geoIPDatabaseReader).getGeoData(any(), any(), any()); Collection> records = geoIPProcessor.doExecute(setEventQueue()); diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java index 129dcb7370..8b095ffb63 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java @@ -5,7 +5,6 @@ package org.opensearch.dataprepper.plugins.processor.configuration; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; @@ -19,32 +18,77 @@ import static org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig.DEFAULT_TARGET; class EntryConfigTest { - private EntryConfig entryConfig; - @BeforeEach - void setUp() { - entryConfig = new EntryConfig(); + private EntryConfig createObjectUnderTest() { + return new EntryConfig(); } @Test void testDefaultConfig() { + final EntryConfig entryConfig = createObjectUnderTest(); + assertThat(entryConfig.getSource(), is(nullValue())); assertThat(entryConfig.getTarget(), equalTo(DEFAULT_TARGET)); - assertThat(entryConfig.getFields(), is(nullValue())); + assertThat(entryConfig.getIncludeFields(), equalTo(null)); + assertThat(entryConfig.getExcludeFields(), equalTo(null)); } @Test void testCustomConfig() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final String sourceValue = "source"; final String targetValue = "target"; - final List fieldsValue = List.of("city", "country"); + final List includeFieldsValue = List.of("city_name"); + final List excludeFieldsValue = List.of("asn"); ReflectivelySetField.setField(EntryConfig.class, entryConfig, "source", sourceValue); ReflectivelySetField.setField(EntryConfig.class, entryConfig, "target", targetValue); - ReflectivelySetField.setField(EntryConfig.class, entryConfig, "fields", fieldsValue); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFieldsValue); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFieldsValue); assertThat(entryConfig.getSource(), equalTo(sourceValue)); assertThat(entryConfig.getTarget(), equalTo(targetValue)); - assertThat(entryConfig.getFields(), equalTo(fieldsValue)); + assertThat(entryConfig.getIncludeFields(), equalTo(includeFieldsValue)); + assertThat(entryConfig.getExcludeFields(), equalTo(excludeFieldsValue)); + } + + @Test + void test_areFieldsValid_should_return_true_if_only_include_fields_is_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List includeFields = List.of("city_name", "continent_code"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(true)); + } + + @Test + void test_areFieldsValid_should_return_true_if_only_exclude_fields_is_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List excludeFields = List.of("city_name", "continent_code"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(true)); + } + + @Test + void test_areFieldsValid_should_return_false_if_both_include_and_exclude_fields_are_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List includeFields = List.of("city_name", "continent_code"); + final List excludeFields = List.of("asn"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFields); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(false)); + } + + @Test + void test_areFieldsValid_should_return_false_if_both_include_and_exclude_fields_are_not_configured() { + final EntryConfig entryConfig = createObjectUnderTest(); + + assertThat(entryConfig.areFieldsValid(), equalTo(false)); } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java new file mode 100644 index 0000000000..c438f3c0b2 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java @@ -0,0 +1,353 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.db.Metadata; +import com.maxmind.db.Network; +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.EnterpriseResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.util.Collections; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IP; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.NETWORK; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN_ORGANIZATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LAT; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LON; + +@ExtendWith(MockitoExtension.class) +class GeoIP2DatabaseReaderTest { + public static final String GEOIP2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-ip2"; + public static final long ASN_RESULT = 12345L; + public static final String ASN_ORG_RESULT = "Example Org"; + private static final String NETWORK_RESULT = "1.2.3.0/24"; + private static final String COUNTRY_NAME_RESULT = "United States"; + private static final String COUNTRY_ISO_CODE_RESULT = "US"; + private static final Boolean COUNTRY_IS_IN_EUROPEAN_UNION_RESULT = false; + private static final String CITY_NAME_RESULT = "Los Angeles"; + private static final Double LATITUDE_RESULT = 12.34; + private static final Double LONGITUDE_RESULT = 56.78; + private static final String TIME_ZONE_RESULT = "America/Los_Angeles"; + private static final Integer METRO_CODE_RESULT = 807; + private static final String POSTAL_CODE_RESULT = "90001"; + private static final String LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String MOST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String REGISTERED_COUNTRY_NAME_RESULT = "Argentina"; + private static final String REGISTERED_COUNTRY_ISO_CODE_RESULT = "AR"; + private static final String REPRESENTED_COUNTRY_NAME_RESULT = "Belgium"; + private static final String REPRESENTED_COUNTRY_ISO_CODE_RESULT = "BE"; + private static final String REPRESENTED_COUNTRY_TYPE_RESULT = "military"; + private static final String CONTINENT_NAME_RESULT = "North America"; + private static final String CONTINENT_CODE_RESULT = "123456"; + private static final Map LOCATION_RESULT = Map.of(LAT, LATITUDE_RESULT, LON, LONGITUDE_RESULT); + private static final Integer COUNTRY_CONFIDENCE_RESULT = 100; + private static final Integer CITY_CONFIDENCE_RESULT = 90; + private static final Integer LOCATION_ACCURACY_RADIUS_RESULT = 10; + private static final Integer POSTAL_CODE_CONFIDENCE_RESULT = 85; + private static final Integer MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT = 75; + private static final Integer LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT = 60; + private static final String IP_RESULT = "1.2.3.4"; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + @Mock + private DatabaseReader enterpriseDatabaseReader; + @Mock + private Metadata metadata; + @Mock + private EnterpriseResponse enterpriseResponse; + @Mock + private AsnResponse asnResponse; + @Mock + private Continent continent; + @Mock + private Country country; + @Mock + private RepresentedCountry representedCountry; + @Mock + private Country registeredCountry; + @Mock + private City city; + @Mock + private Location location; + @Mock + private Network network; + @Mock + private Postal postal; + @Mock + private Subdivision leastSpecificSubdivision; + @Mock + private Subdivision mostSpecificSubdivision; + @Mock + private InetAddress inetAddress; + @Mock + private GeoIPFileManager geoIPFileManager; + + @BeforeEach + void setUp() throws IOException { + when(databaseReaderBuilder.buildReader(Path.of(GEOIP2_TEST_MMDB_FILES + File.separator + "GeoIP2-Enterprise-Test.mmdb"), 0)) + .thenReturn(enterpriseDatabaseReader); + + when(enterpriseDatabaseReader.getMetadata()).thenReturn(metadata); + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + + lenient().when(enterpriseResponse.getContinent()).thenReturn(continent); + lenient().when(continent.getName()).thenReturn(CONTINENT_NAME_RESULT); + lenient().when(continent.getCode()).thenReturn(CONTINENT_CODE_RESULT); + + lenient().when(enterpriseResponse.getCountry()).thenReturn(country); + lenient().when(country.getName()).thenReturn(COUNTRY_NAME_RESULT); + lenient().when(country.getIsoCode()).thenReturn(COUNTRY_ISO_CODE_RESULT); + lenient().when(country.isInEuropeanUnion()).thenReturn(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT); + lenient().when(country.getConfidence()).thenReturn(COUNTRY_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(registeredCountry.getName()).thenReturn(REGISTERED_COUNTRY_NAME_RESULT); + lenient().when(registeredCountry.getIsoCode()).thenReturn(REGISTERED_COUNTRY_ISO_CODE_RESULT); + + lenient().when(enterpriseResponse.getRepresentedCountry()).thenReturn(representedCountry); + lenient().when(representedCountry.getName()).thenReturn(REPRESENTED_COUNTRY_NAME_RESULT); + lenient().when(representedCountry.getIsoCode()).thenReturn(REPRESENTED_COUNTRY_ISO_CODE_RESULT); + lenient().when(representedCountry.getType()).thenReturn(REPRESENTED_COUNTRY_TYPE_RESULT); + + lenient().when(enterpriseResponse.getCity()).thenReturn(city); + lenient().when(city.getName()).thenReturn(CITY_NAME_RESULT); + lenient().when(city.getConfidence()).thenReturn(CITY_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getLocation()).thenReturn(location); + lenient().when(location.getLatitude()).thenReturn(LATITUDE_RESULT); + lenient().when(location.getLongitude()).thenReturn(LONGITUDE_RESULT); + lenient().when(location.getMetroCode()).thenReturn(METRO_CODE_RESULT); + lenient().when(location.getTimeZone()).thenReturn(TIME_ZONE_RESULT); + lenient().when(location.getAccuracyRadius()).thenReturn(LOCATION_ACCURACY_RADIUS_RESULT); + + lenient().when(enterpriseResponse.getPostal()).thenReturn(postal); + lenient().when(postal.getCode()).thenReturn(POSTAL_CODE_RESULT); + lenient().when(postal.getConfidence()).thenReturn(POSTAL_CODE_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getLeastSpecificSubdivision()).thenReturn(leastSpecificSubdivision); + lenient().when(leastSpecificSubdivision.getName()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(leastSpecificSubdivision.getIsoCode()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + lenient().when(leastSpecificSubdivision.getConfidence()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getMostSpecificSubdivision()).thenReturn(mostSpecificSubdivision); + lenient().when(mostSpecificSubdivision.getName()).thenReturn(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(mostSpecificSubdivision.getIsoCode()).thenReturn(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + lenient().when(mostSpecificSubdivision.getConfidence()).thenReturn(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT); + + lenient().when(asnResponse.getAutonomousSystemNumber()).thenReturn(ASN_RESULT); + lenient().when(asnResponse.getAutonomousSystemOrganization()).thenReturn(ASN_ORG_RESULT); + lenient().when(asnResponse.getNetwork()).thenReturn(network); + lenient().when(asnResponse.getIpAddress()).thenReturn(IP_RESULT); + lenient().when(network.toString()).thenReturn(NETWORK_RESULT); + } + + @AfterEach + void tearDown() { + verifyNoMoreInteractions(enterpriseDatabaseReader); + } + + GeoIP2DatabaseReader createObjectUnderTest() { + return new GeoIP2DatabaseReader(databaseReaderBuilder, geoIPFileManager, GEOIP2_TEST_MMDB_FILES, 0); + } + + @Test + void test_getGeoData_for_all_fields_in_enterprise_database_add_only_fields_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + COUNTRY_CONFIDENCE, REPRESENTED_COUNTRY_ISO_CODE, REGISTERED_COUNTRY_ISO_CODE, CITY_NAME, CITY_CONFIDENCE, LOCATION, + LOCATION_ACCURACY_RADIUS, METRO_CODE, TIME_ZONE, POSTAL_CODE, POSTAL_CODE_CONFIDENCE); + final Set databases = Set.of(GeoIPDatabase.ENTERPRISE); + + when(enterpriseDatabaseReader.tryEnterprise(inetAddress)).thenReturn(Optional.of(enterpriseResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(COUNTRY_CONFIDENCE.getFieldName()), equalTo(COUNTRY_CONFIDENCE_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(CITY_CONFIDENCE.getFieldName()), equalTo(CITY_CONFIDENCE_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(POSTAL_CODE_CONFIDENCE.getFieldName()), equalTo(POSTAL_CODE_CONFIDENCE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_enterprise_database_when_no_fields_are_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + final List fields = Collections.emptyList(); + final Set databases = Set.of(GeoIPDatabase.ENTERPRISE); + + when(enterpriseDatabaseReader.tryEnterprise(inetAddress)).thenReturn(Optional.of(enterpriseResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_CONFIDENCE.getFieldName()), equalTo(COUNTRY_CONFIDENCE_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(CITY_CONFIDENCE.getFieldName()), equalTo(CITY_CONFIDENCE_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(POSTAL_CODE_CONFIDENCE.getFieldName()), equalTo(POSTAL_CODE_CONFIDENCE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT)); + } + + @Test + void test_getGeoData_for_asn_fields_in_enterprise_database_when_no_fields_are_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + final Set databases = Set.of(GeoIPDatabase.ASN); + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK, IP); + + when(enterpriseDatabaseReader.tryAsn(inetAddress)).thenReturn(Optional.of(asnResponse)); + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(ASN.getFieldName()), equalTo(ASN_RESULT)); + assertThat(geoData.get(ASN_ORGANIZATION.getFieldName()), equalTo(ASN_ORG_RESULT)); + assertThat(geoData.get(NETWORK.getFieldName()), equalTo(NETWORK_RESULT)); + assertThat(geoData.get(IP.getFieldName()), equalTo(IP_RESULT)); + } + + @Test + void test_database_expired_should_return_false_when_expiry_date_is_in_future() { + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_expired_should_return_true_when_expiry_date_is_in_past() throws IOException { + final Date date = new Date(91911199999L); + when(metadata.getBuildDate()).thenReturn(date); + doNothing().when(geoIPFileManager).deleteDirectory(any()); + + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(true)); + verify(enterpriseDatabaseReader).close(); + } + + @Test + void test_database_close_should_not_close_the_reader_unless_close_count_is_zero() { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_close_should_close_the_reader_when_close_count_is_zero() throws IOException { + doNothing().when(geoIPFileManager).deleteDirectory(any()); + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + verify(enterpriseDatabaseReader).close(); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java new file mode 100644 index 0000000000..912a6eed6f --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java @@ -0,0 +1,397 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.db.Metadata; +import com.maxmind.db.Network; +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.CityResponse; +import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.GeoIPDatabase.CITY; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IP; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LATITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LONGITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.NETWORK; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN_ORGANIZATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LAT; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LON; + + +@ExtendWith(MockitoExtension.class) +class GeoLite2DatabaseReaderTest { + public static final String GEOLITE2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-lite2"; + public static final long ASN_RESULT = 12345L; + public static final String ASN_ORG_RESULT = "Example Org"; + private static final String NETWORK_RESULT = "1.2.3.0/24"; + private static final String COUNTRY_NAME_RESULT = "United States"; + private static final String COUNTRY_ISO_CODE_RESULT = "US"; + private static final Boolean COUNTRY_IS_IN_EUROPEAN_UNION_RESULT = false; + private static final String CITY_NAME_RESULT = "Los Angeles"; + private static final Double LATITUDE_RESULT = 12.34; + private static final Double LONGITUDE_RESULT = 56.78; + private static final Integer LOCATION_ACCURACY_RADIUS_RESULT = 10; + private static final String TIME_ZONE_RESULT = "America/Los_Angeles"; + private static final Integer METRO_CODE_RESULT = 807; + private static final String POSTAL_CODE_RESULT = "90001"; + private static final String LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String MOST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String REGISTERED_COUNTRY_NAME_RESULT = "Argentina"; + private static final String REGISTERED_COUNTRY_ISO_CODE_RESULT = "AR"; + private static final String REPRESENTED_COUNTRY_NAME_RESULT = "Belgium"; + private static final String REPRESENTED_COUNTRY_ISO_CODE_RESULT = "BE"; + private static final String REPRESENTED_COUNTRY_TYPE_RESULT = "military"; + private static final String CONTINENT_NAME_RESULT = "North America"; + private static final String CONTINENT_CODE_RESULT = "123456"; + private static final Map LOCATION_RESULT = Map.of(LAT, LATITUDE_RESULT, LON, LONGITUDE_RESULT); + private static final String IP_RESULT = "1.2.3.4"; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + @Mock + private DatabaseReader countryDatabaseReader; + @Mock + private DatabaseReader cityDatabaseReader; + @Mock + private DatabaseReader asnDatabaseReader; + @Mock + private Metadata metadata; + @Mock + private CountryResponse countryResponse; + @Mock + private CityResponse cityResponse; + @Mock + private AsnResponse asnResponse; + @Mock + private Continent continent; + @Mock + private Country country; + @Mock + private RepresentedCountry representedCountry; + @Mock + private Country registeredCountry; + @Mock + private City city; + @Mock + private Location location; + @Mock + private Network network; + @Mock + private Postal postal; + @Mock + private Subdivision leastSpecificSubdivision; + @Mock + private Subdivision mostSpecificSubdivision; + @Mock + private InetAddress inetAddress; + @Mock + private GeoIPFileManager geoIPFileManager; + + @BeforeEach + void setUp() throws IOException { + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-Country-Test.mmdb"), 0)) + .thenReturn(countryDatabaseReader); + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-City-Test.mmdb"), 0)) + .thenReturn(cityDatabaseReader); + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-ASN-Test.mmdb"), 0)) + .thenReturn(asnDatabaseReader); + + when(countryDatabaseReader.getMetadata()).thenReturn(metadata); + when(cityDatabaseReader.getMetadata()).thenReturn(metadata); + when(asnDatabaseReader.getMetadata()).thenReturn(metadata); + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + + lenient().when(countryResponse.getContinent()).thenReturn(continent); + lenient().when(continent.getName()).thenReturn(CONTINENT_NAME_RESULT); + lenient().when(continent.getCode()).thenReturn(CONTINENT_CODE_RESULT); + + lenient().when(countryResponse.getCountry()).thenReturn(country); + lenient().when(country.getName()).thenReturn(COUNTRY_NAME_RESULT); + lenient().when(country.getIsoCode()).thenReturn(COUNTRY_ISO_CODE_RESULT); + lenient().when(country.isInEuropeanUnion()).thenReturn(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT); + + lenient().when(countryResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(registeredCountry.getName()).thenReturn(REGISTERED_COUNTRY_NAME_RESULT); + lenient().when(registeredCountry.getIsoCode()).thenReturn(REGISTERED_COUNTRY_ISO_CODE_RESULT); + + lenient().when(countryResponse.getRepresentedCountry()).thenReturn(representedCountry); + lenient().when(representedCountry.getName()).thenReturn(REPRESENTED_COUNTRY_NAME_RESULT); + lenient().when(representedCountry.getIsoCode()).thenReturn(REPRESENTED_COUNTRY_ISO_CODE_RESULT); + lenient().when(representedCountry.getType()).thenReturn(REPRESENTED_COUNTRY_TYPE_RESULT); + + lenient().when(cityResponse.getCity()).thenReturn(city); + lenient().when(city.getName()).thenReturn(CITY_NAME_RESULT); + + lenient().when(cityResponse.getContinent()).thenReturn(continent); + lenient().when(cityResponse.getCountry()).thenReturn(country); + lenient().when(cityResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(cityResponse.getRepresentedCountry()).thenReturn(representedCountry); + + lenient().when(cityResponse.getLocation()).thenReturn(location); + lenient().when(location.getLatitude()).thenReturn(LATITUDE_RESULT); + lenient().when(location.getLongitude()).thenReturn(LONGITUDE_RESULT); + lenient().when(location.getMetroCode()).thenReturn(METRO_CODE_RESULT); + lenient().when(location.getTimeZone()).thenReturn(TIME_ZONE_RESULT); + lenient().when(location.getAccuracyRadius()).thenReturn(LOCATION_ACCURACY_RADIUS_RESULT); + + lenient().when(cityResponse.getPostal()).thenReturn(postal); + lenient().when(postal.getCode()).thenReturn(POSTAL_CODE_RESULT); + + lenient().when(cityResponse.getLeastSpecificSubdivision()).thenReturn(leastSpecificSubdivision); + lenient().when(leastSpecificSubdivision.getName()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(leastSpecificSubdivision.getIsoCode()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + + lenient().when(cityResponse.getMostSpecificSubdivision()).thenReturn(mostSpecificSubdivision); + lenient().when(mostSpecificSubdivision.getName()).thenReturn(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(mostSpecificSubdivision.getIsoCode()).thenReturn(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + + lenient().when(asnResponse.getAutonomousSystemNumber()).thenReturn(ASN_RESULT); + lenient().when(asnResponse.getAutonomousSystemOrganization()).thenReturn(ASN_ORG_RESULT); + lenient().when(asnResponse.getNetwork()).thenReturn(network); + lenient().when(asnResponse.getIpAddress()).thenReturn(IP_RESULT); + lenient().when(network.toString()).thenReturn(NETWORK_RESULT); + } + + @AfterEach + void tearDown() { + verifyNoMoreInteractions(countryDatabaseReader); + verifyNoMoreInteractions(cityDatabaseReader); + verifyNoMoreInteractions(asnDatabaseReader); + } + + GeoLite2DatabaseReader createObjectUnderTest() { + return new GeoLite2DatabaseReader(databaseReaderBuilder, geoIPFileManager, GEOLITE2_TEST_MMDB_FILES, 0); + } + + @Test + void test_getGeoData_for_all_fields_in_country_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE); + final Set databases = Set.of(GeoIPDatabase.COUNTRY); + + when(countryDatabaseReader.tryCountry(inetAddress)).thenReturn(Optional.of(countryResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_city_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE, CITY_NAME, LOCATION, LOCATION_ACCURACY_RADIUS, LATITUDE, LONGITUDE, METRO_CODE, TIME_ZONE, POSTAL_CODE, + MOST_SPECIFIED_SUBDIVISION_NAME, MOST_SPECIFIED_SUBDIVISION_ISO_CODE, LEAST_SPECIFIED_SUBDIVISION_NAME, + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE); + final Set databases = Set.of(CITY); + + when(cityDatabaseReader.tryCity(inetAddress)).thenReturn(Optional.of(cityResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LATITUDE.getFieldName()), equalTo(LATITUDE_RESULT)); + assertThat(geoData.get(LONGITUDE.getFieldName()), equalTo(LONGITUDE_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_asn_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK, IP); + final Set databases = Set.of(GeoIPDatabase.ASN); + + when(asnDatabaseReader.tryAsn(inetAddress)).thenReturn(Optional.of(asnResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(ASN.getFieldName()), equalTo(ASN_RESULT)); + assertThat(geoData.get(ASN_ORGANIZATION.getFieldName()), equalTo(ASN_ORG_RESULT)); + assertThat(geoData.get(NETWORK.getFieldName()), equalTo(NETWORK_RESULT)); + assertThat(geoData.get(IP.getFieldName()), equalTo(IP_RESULT)); + } + + @Test + void test_getGeoData_for_country_database_should_not_add_any_fields_if_country_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_getGeoData_for_city_database_should_not_add_any_fields_if_city_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_getGeoData_for_asn_database_should_not_add_any_fields_if_asn_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_database_expired_should_return_false_when_expiry_date_is_in_future() { + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_expired_should_return_true_when_expiry_date_is_in_past() throws IOException { + final Date date = new Date(91911199999L); + when(metadata.getBuildDate()).thenReturn(date); + doNothing().when(geoIPFileManager).deleteFile(any()); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(true)); + verify(countryDatabaseReader).close(); + verify(cityDatabaseReader).close(); + verify(asnDatabaseReader).close(); + } + + @Test + void test_database_close_should_not_close_the_reader_unless_close_count_is_zero() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_close_should_close_the_reader_when_close_count_is_zero() throws IOException { + doNothing().when(geoIPFileManager).deleteDirectory(any()); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + verify(countryDatabaseReader).close(); + verify(cityDatabaseReader).close(); + verify(asnDatabaseReader).close(); + + } + +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java deleted file mode 100644 index 5493082d08..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIpServiceConfig; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; - -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GetGeoIP2DataTest { - - private static final String PATH = "./src/test/resources/mmdb-file/geo-enterprise"; - private String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; - private static final String PREFIX_DIR = "first_database"; - public static final int REFRESH_SCHEDULE = 10; - public static final String IP = "2001:4860:4860::8888"; - @Mock - private GeoIPProcessorConfig geoIPProcessorConfig; - @Mock - private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private MaxMindConfig maxMindConfig; - @Mock - private DBSource downloadSource; - private GetGeoIP2Data getGeoIP2Data; - private final int cacheSize = 4068; - - @BeforeEach - void setUp() { - when(geoIpServiceConfig.getMaxMindConfig()) - .thenReturn(maxMindConfig); - String dbPath = "./src/test/resources/mmdb-file/geo-enterprise"; - getGeoIP2Data = new GetGeoIP2Data(dbPath, cacheSize); - } - - @Disabled("Doesn't have valid GeoIP2-Enterprise.mmdb") - @Test - void getGeoDataTest() throws UnknownHostException { - - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("ip"), equalTo("2001:4860:4860:0:0:0:0:8888")); - } - - @Test - @Disabled - void switchDatabaseReaderTest() { - tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIPMaxmind"; - DBSource.createFolderIfNotExist(tempFolderPath); - getGeoIP2Data = new GetGeoIP2Data(tempFolderPath, cacheSize); - assertDoesNotThrow(() -> { - getGeoIP2Data.switchDatabaseReader(); - }); - assertDoesNotThrow(() -> { - getGeoIP2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_cover_EnrichFailedException() throws UnknownHostException { - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - assertThrows(EnrichFailedException.class, () -> getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } - - @Test - @Disabled - void closeReaderTest() throws UnknownHostException { - getGeoIP2Data.closeReader(); - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - assertThrows(EnrichFailedException.class, () -> getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java deleted file mode 100644 index 7a58317a5d..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java +++ /dev/null @@ -1,129 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIpServiceConfig; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; - -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GetGeoLite2DataTest { - - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; - public static final String IP = "2a02:ec00:0:0:0:0:0:0"; - private static final String PREFIX_DIR = "first_database"; - private String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; - @Mock - private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private MaxMindConfig maxMindConfig; - private GetGeoLite2Data getGeoLite2Data; - private final int cacheSize = 4068; - - @BeforeEach - void setUp() { - when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); - when(maxMindConfig.getCacheSize()).thenReturn(cacheSize); - - final String dbPath = "./src/test/resources/mmdb-file/geo-lite2"; - getGeoLite2Data = new GetGeoLite2Data(dbPath, cacheSize); - } - - @Test - @Disabled - void getGeoDataTest_without_attributes() throws UnknownHostException { - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath); - Assertions.assertNotNull(geoData); - assertThat(geoData.get("country_iso_code"), equalTo("FR")); - assertThat(geoData.get("ip"), equalTo(IP)); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_with_attributes() throws UnknownHostException { - List attributes = List.of("city_name", - "country_name", - "ip", - "country_iso_code", - "continent_name", - "region_iso_code", - "region_name", - "timezone", - "location", - "asn", - "organization_name", "network"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath); - Assertions.assertNotNull(geoData); - assertThat(geoData.get("country_name"), equalTo("France")); - assertThat(geoData.get("ip"), equalTo(IP)); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_cover_EnrichFailedException() throws UnknownHostException { - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - String dbPath = "/src/test/resources/mmdb-file/geo-enterprise"; - getGeoLite2Data = new GetGeoLite2Data(dbPath, cacheSize); - GeoIPProcessorService.downloadReady = false; - assertThrows(EnrichFailedException.class, () -> getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath)); - } - - @Test - @Disabled - void switchDatabaseReaderTest() { - tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIPMaxmind"; - DBSource.createFolderIfNotExist(tempFolderPath); - getGeoLite2Data = new GetGeoLite2Data(tempFolderPath, cacheSize); - assertDoesNotThrow(() -> { - getGeoLite2Data.switchDatabaseReader(); - }); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void closeReaderTest() throws UnknownHostException { - // While closing the readers, all the readers reset to null. - getGeoLite2Data.closeReader(); - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - assertThrows(NullPointerException.class, () -> getGeoLite2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java new file mode 100644 index 0000000000..d5a3edbabc --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class DatabaseReaderInitializationExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private DatabaseReaderInitializationException createObjectUnderTest() { + return new DatabaseReaderInitializationException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java new file mode 100644 index 0000000000..ca12b8a8d0 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class DownloadFailedExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private DownloadFailedException createObjectUnderTest() { + return new DownloadFailedException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java new file mode 100644 index 0000000000..b1a6f4b822 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class EngineFailureExceptionTest { + private String message; + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private EngineFailureException createObjectUnderTest() { + return new EngineFailureException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java new file mode 100644 index 0000000000..af31a9532e --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class EnrichFailedExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private EnrichFailedException createObjectUnderTest() { + return new EnrichFailedException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java new file mode 100644 index 0000000000..d1f9e0c517 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class InvalidIPAddressExceptionTest { + private String message; + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private InvalidIPAddressException createObjectUnderTest() { + return new InvalidIPAddressException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java new file mode 100644 index 0000000000..4463d10141 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class NoValidDatabaseFoundExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private NoValidDatabaseFoundException createObjectUnderTest() { + return new NoValidDatabaseFoundException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java index 9d42f5b6f2..c5f484a769 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java @@ -10,10 +10,15 @@ import org.mockito.Mock; import org.mockito.MockedConstruction; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; + +import java.util.Optional; +import java.util.concurrent.locks.ReentrantReadWriteLock; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mockConstruction; @ExtendWith(MockitoExtension.class) @@ -21,8 +26,14 @@ class DefaultGeoIpConfigSupplierTest { @Mock private GeoIpServiceConfig geoIpServiceConfig; + @Mock + private GeoIPDatabaseManager geoIPDatabaseManager; + + @Mock + private ReentrantReadWriteLock.ReadLock readLock; + private DefaultGeoIpConfigSupplier createObjectUnderTest() { - return new DefaultGeoIpConfigSupplier(geoIpServiceConfig); + return new DefaultGeoIpConfigSupplier(geoIpServiceConfig, geoIPDatabaseManager, readLock); } @Test @@ -30,11 +41,24 @@ void getGeoIpProcessorService_returns_geoIPProcessorService() { try (final MockedConstruction mockedConstruction = mockConstruction(GeoIPProcessorService.class)) { final DefaultGeoIpConfigSupplier objectUnderTest = createObjectUnderTest(); - final GeoIPProcessorService geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); + final Optional geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); assertThat(mockedConstruction.constructed().size(), equalTo(1)); - assertThat(geoIPProcessorService, instanceOf(GeoIPProcessorService.class)); - assertThat(geoIPProcessorService, equalTo(mockedConstruction.constructed().get(0))); + assertTrue(geoIPProcessorService.isPresent()); + assertThat(geoIPProcessorService.get(), instanceOf(GeoIPProcessorService.class)); + assertThat(geoIPProcessorService.get(), equalTo(mockedConstruction.constructed().get(0))); + } + } + + @Test + void getGeoIpProcessorService_returns_empty_if_service_config_is_null() { + try (final MockedConstruction mockedConstruction = + mockConstruction(GeoIPProcessorService.class)) { + final DefaultGeoIpConfigSupplier objectUnderTest = new DefaultGeoIpConfigSupplier(null, geoIPDatabaseManager, readLock); + final Optional geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); + + assertThat(mockedConstruction.constructed().size(), equalTo(0)); + assertTrue(geoIPProcessorService.isEmpty()); } } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java index 232f98edec..db607a6fb0 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java @@ -6,89 +6,101 @@ package org.opensearch.dataprepper.plugins.processor.extension; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.time.Duration; +import java.util.concurrent.locks.ReentrantReadWriteLock; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) -@Disabled class GeoIPProcessorServiceTest { - - private static final String S3_URL = "https://mybucket10012023.s3.amazonaws.com/GeoLite2/"; - private static final String URL = "https://download.maxmind.com/app/geoip_download?edition_" + - "id=GeoLite2-ASN&suffix=tar.gz"; - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; - private static final String S3_REGION = "us-east-1"; - public static final int REFRESH_SCHEDULE = 10; - public static final String IP = "2001:4860:4860::8888"; - String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; @Mock - private GeoIPProcessorConfig geoIPProcessorConfig; + private GeoIPDatabaseReader geoIPDatabaseReaderMock; @Mock - private GetGeoData geoData; + private GeoIPDatabaseReader newGeoIPDatabaseReaderMock; @Mock private MaxMindConfig maxMindConfig; - private GeoIPProcessorService geoIPProcessorService; + @Mock + private GeoIpServiceConfig geoIpServiceConfig; + @Mock + private GeoIPDatabaseManager geoIPDatabaseManager; + @Mock + private ReentrantReadWriteLock.ReadLock readLock; @BeforeEach void setUp() { + when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); + doNothing().when(geoIPDatabaseManager).initiateDatabaseDownload(); + when(geoIPDatabaseManager.getGeoIPDatabaseReader()).thenReturn(geoIPDatabaseReaderMock, newGeoIPDatabaseReaderMock); + } + GeoIPProcessorService createObjectUnderTest() { + return new GeoIPProcessorService(geoIpServiceConfig, geoIPDatabaseManager, readLock); } @Test - void getGeoDataTest_PATH() throws UnknownHostException, NoSuchFieldException, IllegalAccessException { - // TODO: pass in geoIpServiceConfig object - geoIPProcessorService = new GeoIPProcessorService(null); - - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - ReflectivelySetField.setField(GeoIPProcessorService.class, - geoIPProcessorService, "geoData", geoData); - when(geoData.getGeoData(any(), any(), anyString())).thenReturn(prepareGeoData()); - Map geoData = geoIPProcessorService.getGeoData(inetAddress, attributes); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("continent_name"), equalTo("North America")); + void test_getGeoIPDatabaseReader_should_not_trigger_update_when_refresh_interval_is_high() { + when(maxMindConfig.getDatabaseRefreshInterval()).thenReturn(Duration.ofHours(1)); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verifyNoMoreInteractions(geoIPDatabaseManager); } @Test - void getGeoDataTest_URL() throws UnknownHostException, NoSuchFieldException, IllegalAccessException { - // TODO: pass in geoIpServiceConfig object - geoIPProcessorService = new GeoIPProcessorService(null); - - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - ReflectivelySetField.setField(GeoIPProcessorService.class, - geoIPProcessorService, "geoData", geoData); - when(geoData.getGeoData(any(), any(), anyString())).thenReturn(prepareGeoData()); - Map geoData = geoIPProcessorService.getGeoData(inetAddress, attributes); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("continent_name"), equalTo("North America")); + void test_getGeoIPDatabaseReader_should_trigger_update_when_refresh_interval_is_met() throws InterruptedException { + when(maxMindConfig.getDatabaseRefreshInterval()).thenReturn(Duration.ofNanos(1)); + doNothing().when(geoIPDatabaseManager).updateDatabaseReader(); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + + // Wait for update to be called by ExecutorService + Thread.sleep(1000); + verify(geoIPDatabaseManager).updateDatabaseReader(); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verifyNoMoreInteractions(geoIPDatabaseManager); + + final GeoIPDatabaseReader newGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(newGeoIPDatabaseReaderMock, equalTo(newGeoIPDatabaseReader)); } - private Map prepareGeoData() { - Map geoDataMap = new HashMap<>(); - geoDataMap.put("country_iso_code", "US"); - geoDataMap.put("continent_name", "North America"); - geoDataMap.put("timezone", "America/Chicago"); - geoDataMap.put("country_name", "United States"); - return geoDataMap; + @Test + void test_shutdown() throws InterruptedException { + when(maxMindConfig.getDatabaseRefreshInterval()).thenReturn(Duration.ofNanos(1)); + doNothing().when(geoIPDatabaseManager).updateDatabaseReader(); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + + // Wait for update to be called by ExecutorService + Thread.sleep(1000); + verify(geoIPDatabaseManager).updateDatabaseReader(); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verifyNoMoreInteractions(geoIPDatabaseManager); + + final GeoIPDatabaseReader newGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(newGeoIPDatabaseReaderMock, equalTo(newGeoIPDatabaseReader)); + + assertDoesNotThrow(objectUnderTest::shutdown); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java index 91c48322a2..78a3076fb2 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -19,17 +20,22 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) class GeoIpConfigExtensionTest { @Mock private ExtensionPoints extensionPoints; - @Mock private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private ExtensionProvider.Context context; + private MaxMindConfig maxMindConfig; + + + @BeforeEach + void setUp() { + when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); + } private GeoIpConfigExtension createObjectUnderTest() { return new GeoIpConfigExtension(geoIpServiceConfig); @@ -48,17 +54,6 @@ void apply_should_addExtensionProvider() { assertThat(actualExtensionProvider, instanceOf(GeoIpConfigProvider.class)); } - @Test - void extension_should_create_supplier_with_default_config_if_not_configured() { - try (final MockedConstruction mockedConstruction = - mockConstruction(GeoIpServiceConfig.class)) { - final GeoIpConfigExtension geoIpConfigExtension = new GeoIpConfigExtension(null); - - assertThat(mockedConstruction.constructed().size(), equalTo(1)); - assertThat(geoIpConfigExtension, instanceOf(GeoIpConfigExtension.class)); - } - } - @Test void extension_should_create_supplier_with_provided_config() { try (final MockedConstruction mockedConstruction = diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java index 6b22d14bf4..228bfa75ad 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java @@ -7,8 +7,11 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import java.net.URISyntaxException; import java.time.Duration; import java.util.List; @@ -46,4 +49,18 @@ void testCustomConfig() throws NoSuchFieldException, IllegalAccessException { assertThat(maxMindConfig.getAwsAuthenticationOptionsConfig(), equalTo(awsAuthenticationOptionsConfig)); } + @ParameterizedTest + @CsvSource({ + "https://example.com/, false, true", + "http://example.com/, false, false", + "https://example.com/, true, true", + "http://example.com/, true, true"}) + void testSecureEndpoint(final String databasePath, final boolean insecure, final boolean result) + throws NoSuchFieldException, IllegalAccessException, URISyntaxException { + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "databasePaths", List.of(databasePath)); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "insecure", insecure); + + assertThat(maxMindConfig.getDatabasePaths().size(), equalTo(1)); + assertThat(maxMindConfig.isHttpsEndpointOrInsecure(), equalTo(result)); + } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java deleted file mode 100644 index 3b7ed49d87..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; - -import java.io.File; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.mockStatic; - -@ExtendWith(MockitoExtension.class) -class DBSourceTest { - - private final String outputFilePath = System.getProperty("java.io.tmpdir") + "GeoTest"; - - @Test - void createFolderIfNotExistTest() { - try (MockedStatic mockedStatic = mockStatic(DBSource.class)) { - mockedStatic.when(() -> DBSource.createFolderIfNotExist(outputFilePath)).thenReturn(new File(outputFilePath)); - File actualFile = new File(outputFilePath); - assertEquals(actualFile, DBSource.createFolderIfNotExist(outputFilePath)); - } - } - - @Test - void deleteDirectoryTest() { - DBSource.createFolderIfNotExist(outputFilePath); - DBSource.createFolderIfNotExist(outputFilePath + File.separator + "GeoIPz"); - DBSource.createFolderIfNotExist(outputFilePath + File.separator + "GeoIPx"); - assertDoesNotThrow(() -> { - DBSource.deleteDirectory(new File(outputFilePath)); - }); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java new file mode 100644 index 0000000000..1b435fb1f2 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java @@ -0,0 +1,71 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.maxmind.geoip2.DatabaseReader; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@ExtendWith(MockitoExtension.class) +class DatabaseReaderBuilderTest { + public static final String GEOLITE2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-lite2"; + public static final String GEOIP2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-ip2"; + + @ParameterizedTest + @ValueSource(strings = {"geolite2-asn", "geolite2-city", "geolite2-country"}) + void createLoaderTest_for_geolite2_databases(final String databaseName) throws IOException { + String databaseToUse = null; + final File directory = new File(GEOLITE2_TEST_MMDB_FILES); + final String[] list = directory.list(); + for (String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (fileName.endsWith(".mmdb") + && lowerCaseFileName.contains(databaseName)) { + databaseToUse = fileName; + } + } + + final Path path = Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + databaseToUse); + + final DatabaseReaderBuilder databaseReaderBuilder = new DatabaseReaderBuilder(); + + final DatabaseReader databaseReader = databaseReaderBuilder.buildReader(path, 4096); + assertNotNull(databaseReader); + assertTrue(databaseToUse.toLowerCase().contains(databaseReader.getMetadata().getDatabaseType().toLowerCase())); + } + + @ParameterizedTest + @ValueSource(strings = {"geoip2-enterprise"}) + void createLoaderTest_for_geoip2_databases(final String databaseName) throws IOException { + String databaseToUse = null; + final File directory = new File(GEOIP2_TEST_MMDB_FILES); + final String[] list = directory.list(); + for (String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (fileName.endsWith(".mmdb") + && lowerCaseFileName.contains(databaseName)) { + databaseToUse = fileName; + } + } + + final Path path = Path.of(GEOIP2_TEST_MMDB_FILES + File.separator + databaseToUse); + + final DatabaseReaderBuilder databaseReaderBuilder = new DatabaseReaderBuilder(); + + final DatabaseReader databaseReader = databaseReaderBuilder.buildReader(path, 4096); + assertNotNull(databaseReader); + assertTrue(databaseToUse.toLowerCase().contains(databaseReader.getMetadata().getDatabaseType().toLowerCase())); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java deleted file mode 100644 index 621498e8e2..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import com.maxmind.geoip2.DatabaseReader; -import org.apache.commons.io.FileUtils; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -import java.io.BufferedInputStream; -import java.io.File; -import java.io.IOException; -import java.net.URL; -import java.nio.file.Path; - -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class DatabaseReaderCreateTest { - @Mock - private Path path; - - @Test - void createLoaderTest() throws IOException { - final String testFileURL = "https://github.com/maxmind/MaxMind-DB/raw/main/test-data/GeoLite2-City-Test.mmdb"; - final File file = File.createTempFile( "GeoIP2-City-Test", ".mmdb"); - - final BufferedInputStream in = new BufferedInputStream(new URL(testFileURL).openStream()); - FileUtils.copyInputStreamToFile(in, file); - when(path.toFile()).thenReturn(file); - - DatabaseReader databaseReader = DatabaseReaderCreate.createLoader(path, 4096); - Assertions.assertNotNull(databaseReader); - in.close(); - file.deleteOnExit(); - } -} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java deleted file mode 100644 index 83be77f774..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoIP2Data; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoLite2Data; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; - -import static org.junit.jupiter.api.Assertions.assertInstanceOf; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GeoDataFactoryTest { - - @Mock - private MaxMindConfig maxMindConfig; - - @Mock - private LicenseTypeCheck licenseTypeCheck; - - @Test - void testCreateWithFreeLicense() { - when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(anyString())).thenReturn(LicenseTypeOptions.FREE); - final GeoDataFactory geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - final String databasePath = "testPath"; - - final GetGeoData getGeoData = geoDataFactory.create(databasePath); - assertInstanceOf(GetGeoLite2Data.class, getGeoData); - } - - @Test - void testCreateWithEnterpriseLicense() { - when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(anyString())).thenReturn(LicenseTypeOptions.ENTERPRISE); - final GeoDataFactory geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - final String databasePath = "testPath"; - - final GetGeoData getGeoData = geoDataFactory.create(databasePath); - assertInstanceOf(GetGeoIP2Data.class, getGeoData); - } - -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java new file mode 100644 index 0000000000..f19a3ccc82 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java @@ -0,0 +1,231 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedConstruction; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIP2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoLite2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; + +import java.util.List; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class GeoIPDatabaseManagerTest { + private static final String S3_URI = "s3://geoip/data"; + private static final String HTTPS_URL = "https://download.maxmind.com/app/geoip_download?editionid=GeoLite2-ASN&suffix=tar.gz"; + private static final String PATH = "./build/resources/test/mmdb-files/geo-lite2"; + private static final String CDN_ENDPOINT = "https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2/manifest.json"; + + @Mock + private MaxMindConfig maxMindConfig; + + @Mock + private LicenseTypeCheck licenseTypeCheck; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + + @Mock + private ReentrantReadWriteLock.WriteLock writeLock; + + @Mock + private GeoIPFileManager databaseFileManager; + + @BeforeEach + void setUp() { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.FREE); + } + + private GeoIPDatabaseManager createObjectUnderTest() { + return new GeoIPDatabaseManager(maxMindConfig, licenseTypeCheck, databaseReaderBuilder, databaseFileManager, writeLock); + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_file_path_should_use_local_download_service_and_geolite2_reader() throws Exception { + try (final MockedConstruction localDBDownloadServiceMockedConstruction = mockConstruction(LocalDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(PATH))); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class) + ) { + + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(PATH)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(localDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(localDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(PATH)); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_file_path_should_use_local_download_service_and_geoip2_reader() throws Exception { + try (final MockedConstruction localDBDownloadServiceMockedConstruction = mockConstruction(LocalDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(PATH))); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(PATH)); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(localDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(localDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(PATH)); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoIP2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_s3_uri_should_use_s3_download_service_and_geolite2_reader() { + try (final MockedConstruction s3DBServiceMockedConstruction = mockConstruction(S3DBService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(S3_URI))); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class) + ) { + + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(S3_URI)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(s3DBServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(s3DBServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(S3_URI)); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_s3_uri_should_use_s3_download_service_and_geoip2_reader() { + try (final MockedConstruction s3DBServiceMockedConstruction = mockConstruction(S3DBService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(S3_URI))); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(S3_URI)); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(s3DBServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(s3DBServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(S3_URI)); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoIP2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_url_should_use_http_download_service_and_geolite2_reader() { + try (final MockedConstruction httpDBDownloadServiceMockedConstruction = mockConstruction(HttpDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(HTTPS_URL))); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class) + ) { + + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(HTTPS_URL)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(httpDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(httpDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(HTTPS_URL)); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_url_should_use_http_download_service_and_geoip2_reader() { + try (final MockedConstruction httpDBDownloadServiceMockedConstruction = mockConstruction(HttpDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(HTTPS_URL))); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(HTTPS_URL)); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(httpDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(httpDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(HTTPS_URL)); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoIP2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_cdn_should_use_cdn_download_service_and_geolite2_reader() { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(CDN_ENDPOINT))); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class) + ) { + + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(CDN_ENDPOINT)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(cdnDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(List.of(CDN_ENDPOINT)); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_updateDatabaseReader_with_geolite2_cdn_should_use_cdn_download_service_and_geolite2_reader_and_get_new_reader() { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload(List.of(CDN_ENDPOINT))); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class) + ) { + when(maxMindConfig.getDatabasePaths()).thenReturn(List.of(CDN_ENDPOINT)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + + objectUnderTest.initiateDatabaseDownload(); + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(0))); + + objectUnderTest.updateDatabaseReader(); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(2)); + for (ManifestDownloadService manifestDownloadService : cdnDownloadServiceMockedConstruction.constructed()) { + verify(manifestDownloadService).initiateDownload(List.of(CDN_ENDPOINT)); + } + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(2)); + // verify if first instance is closed + verify(geoLite2DatabaseReaderMockedConstruction.constructed().get(0)).close(); + final GeoIPDatabaseReader updatedGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(updatedGeoIPDatabaseReader, equalTo(geoLite2DatabaseReaderMockedConstruction.constructed().get(1))); + } + } +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java new file mode 100644 index 0000000000..c7c74acf8d --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java @@ -0,0 +1,36 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.Test; + +import java.io.File; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +class GeoIPFileManagerTest { + private final String outputFilePath = "./src/test/resources/geoip/test"; + + @Test + void createFolderIfNotExistTest() { + final GeoIPFileManager geoIPFileManager = new GeoIPFileManager(); + geoIPFileManager.createDirectoryIfNotExist(outputFilePath); + + final File file = new File(outputFilePath); + assertTrue(file.exists()); + } + + @Test + void deleteDirectoryTest() { + final GeoIPFileManager geoIPFileManager = new GeoIPFileManager(); + geoIPFileManager.createDirectoryIfNotExist(outputFilePath); + + final File file = new File(outputFilePath); + assertTrue(file.isDirectory()); + geoIPFileManager.deleteDirectory(file); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java index 8b2bbf47cd..87152004e8 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java @@ -7,6 +7,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import java.util.List; @@ -18,6 +19,8 @@ class HttpDBDownloadServiceTest { private static final String PREFIX_DIR = "first_database"; private HttpDBDownloadService downloadThroughUrl; + @Mock + private GeoIPFileManager geoIPFileManager; @Test void initiateDownloadTest() { @@ -30,6 +33,6 @@ void initiateDownloadTest() { } private HttpDBDownloadService createObjectUnderTest() { - return new HttpDBDownloadService(PREFIX_DIR); + return new HttpDBDownloadService(PREFIX_DIR, geoIPFileManager); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java new file mode 100644 index 0000000000..2dfe276104 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; + +import java.io.File; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class ManifestDownloadServiceTest { + private static final String OUTPUT_DIR = "./src/test/resources/geoip"; + + private ManifestDownloadService createObjectUnderTest() { + return new ManifestDownloadService(OUTPUT_DIR); + } + + @Test + void test_with_valid_endpoint_should_download_file() { + final ManifestDownloadService objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDownload(List.of("https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2-city/manifest.json")); + + final File file = new File(OUTPUT_DIR + File.separator + "geolite2-city.mmdb"); + assertTrue(file.exists()); + + file.deleteOnExit(); + final File directory = new File(OUTPUT_DIR); + directory.deleteOnExit(); + } + + @Test + void test_with_invalid_endpoint_should_throw_exception() { + final ManifestDownloadService objectUnderTest = createObjectUnderTest(); + assertThrows(DownloadFailedException.class, () -> objectUnderTest + .initiateDownload(List.of("https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2-enterprise/manifest.json"))); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java index 815d2eb624..42e35164e7 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java @@ -11,7 +11,7 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; import java.util.List; diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java index 7be0b00db9..64b3a8f24c 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java @@ -20,65 +20,81 @@ class DbSourceIdentificationTest { private static final String S3_URI = "s3://dataprepper/logdata/22833bd46b8e0.mmdb"; - private static final String S3_URL = "https://dataprepper.s3.amazonaws.com/logdata/22833bd46b8e0.json"; private static final String URL = "https://www.dataprepper.com"; - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; + private static final String DIRECTORY_PATH = "./build/resources/test/mmdb-files/geo-lite2"; + private static final String FILE_PATH = "./build/resources/test/mmdb-files/geo-lite2/GeoLite2-ASN-Test.mmdb"; + private static final String CDN_ENDPOINT_HOST = "https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2/manifest.json"; @Test void test_positive_case() { - assertTrue(DbSourceIdentification.isS3Uri(S3_URI)); - assertTrue(DbSourceIdentification.isS3Url(S3_URL)); - assertTrue(DbSourceIdentification.isURL(URL)); - assertTrue(DbSourceIdentification.isFilePath(PATH)); + assertTrue(DatabaseSourceIdentification.isS3Uri(S3_URI)); + assertTrue(DatabaseSourceIdentification.isURL(URL)); + assertTrue(DatabaseSourceIdentification.isFilePath(DIRECTORY_PATH)); + assertTrue(DatabaseSourceIdentification.isCDNEndpoint(CDN_ENDPOINT_HOST)); } @Test void test_negative_case() { - assertFalse(DbSourceIdentification.isS3Uri(S3_URL)); - assertFalse(DbSourceIdentification.isS3Uri(URL)); - assertFalse(DbSourceIdentification.isS3Uri(PATH)); + assertFalse(DatabaseSourceIdentification.isS3Uri(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isS3Uri(URL)); + assertFalse(DatabaseSourceIdentification.isS3Uri(DIRECTORY_PATH)); - assertFalse(DbSourceIdentification.isS3Url(S3_URI)); - assertFalse(DbSourceIdentification.isS3Url(URL)); - assertFalse(DbSourceIdentification.isS3Url(PATH)); + assertFalse(DatabaseSourceIdentification.isURL(S3_URI)); + assertFalse(DatabaseSourceIdentification.isURL(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isURL(DIRECTORY_PATH)); - assertFalse(DbSourceIdentification.isURL(S3_URI)); - assertFalse(DbSourceIdentification.isURL(S3_URL)); - assertFalse(DbSourceIdentification.isURL(PATH)); + assertFalse(DatabaseSourceIdentification.isFilePath(S3_URI)); + assertFalse(DatabaseSourceIdentification.isFilePath(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isFilePath(URL)); - assertFalse(DbSourceIdentification.isFilePath(S3_URI)); - assertFalse(DbSourceIdentification.isFilePath(S3_URL)); - assertFalse(DbSourceIdentification.isFilePath(URL)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(S3_URI)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(DIRECTORY_PATH)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(URL)); } @Test - void getDatabasePathTypeTest_PATH() throws NoSuchFieldException, IllegalAccessException { - List databasePath = List.of("./src/test/resources/mmdb-file/geo-lite2"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + void getDatabasePathTypeTest_PATH() { + List databasePath = List.of(DIRECTORY_PATH); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.PATH)); } @Test - void getDatabasePathTypeTest_URL() throws NoSuchFieldException, IllegalAccessException { + void getDatabasePathTypeTest_should_return_null_if_not_directory() { + List databasePath = List.of(FILE_PATH); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); + Assertions.assertNull(dbSourceOptions); + } + + @Test + void getDatabasePathTypeTest_URL() { List databasePath = List.of("https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-ASN&suffix=tar.gz"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.URL)); } @Test - void getDatabasePathTypeTest_S3() throws NoSuchFieldException, IllegalAccessException { + void getDatabasePathTypeTest_S3() { List databasePath = List.of("s3://mybucket10012023/GeoLite2/"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.S3)); } + @Test + void getDatabasePathTypeTest_CDN() { + List databasePath = List.of(CDN_ENDPOINT_HOST); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); + Assertions.assertNotNull(dbSourceOptions); + assertThat(dbSourceOptions, equalTo(DBSourceOptions.HTTP_MANIFEST)); + } + @Test void isS3Uri_NullPointerException_test() { assertDoesNotThrow(() -> { - DbSourceIdentification.isS3Uri(null); + DatabaseSourceIdentification.isS3Uri(null); }); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java index 487f131f81..03dc47aac0 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java @@ -9,28 +9,29 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.exception.InvalidIPAddressException; -import java.net.UnknownHostException; +import static org.junit.jupiter.api.Assertions.assertThrows; @ExtendWith(MockitoExtension.class) class IPValidationCheckTest { private static final String PRIVATE_IP_ADDRESS = "192.168.29.233"; private static final String PUBLIC_IP_ADDRESS = "2001:4860:4860::8888"; - private static final String INVALID_IP_ADDRESS = "255.255.255.0"; + private static final String INVALID_IP_ADDRESS = "255.255.255.999"; @Test - void ipValidationcheckTest_positive() throws UnknownHostException { + void ipValidationcheckTest_public() { Assertions.assertTrue(IPValidationCheck.isPublicIpAddress(PUBLIC_IP_ADDRESS)); } @Test - void ipValidationcheckTest_negative() throws UnknownHostException { + void ipValidationcheckTest_negative() { Assertions.assertFalse(IPValidationCheck.isPublicIpAddress(PRIVATE_IP_ADDRESS)); } @Test - void ipValidationcheckTest_invalid() throws UnknownHostException { - Assertions.assertTrue(IPValidationCheck.isPublicIpAddress(INVALID_IP_ADDRESS)); + void ipValidationcheckTest_invalid() { + assertThrows(InvalidIPAddressException.class, () -> IPValidationCheck.isPublicIpAddress(INVALID_IP_ADDRESS)); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java index 02d40cdf1e..5530367d7c 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java @@ -17,30 +17,24 @@ @ExtendWith(MockitoExtension.class) class LicenseTypeCheckTest { - private static final String FOLDER_PATH_GEO_LITE2 = "./src/test/resources/mmdb-file/geo-lite2"; - private static final String FOLDER_PATH_GEO_ENTERPRISE = "./src/test/resources/mmdb-file/geo-enterprise"; + private static final String FOLDER_PATH_GEO_LITE2 = "./build/resources/test/mmdb-files/geo-lite2"; + private static final String FOLDER_PATH_GEO_ENTERPRISE = "./build/resources/test/mmdb-files/geo-enterprise"; private LicenseTypeCheck createObjectUnderTest() { return new LicenseTypeCheck(); } @Test - void isGeoLite2OrEnterpriseLicenseTest_positive() { + void test_isGeoLite2OrEnterpriseLicenseTest_should_return_free_when_geolite2_databases_are_used() { final LicenseTypeCheck objectUnderTest = createObjectUnderTest(); LicenseTypeOptions licenseTypeOptionsFree = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_LITE2); assertThat(licenseTypeOptionsFree, equalTo(LicenseTypeOptions.FREE)); - - LicenseTypeOptions licenseTypeOptionsEnterprise = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_ENTERPRISE); - assertThat(licenseTypeOptionsEnterprise, equalTo(LicenseTypeOptions.ENTERPRISE)); } @Test - void isGeoLite2OrEnterpriseLicenseTest_negative() { + void test_isGeoLite2OrEnterpriseLicenseTest_should_return_enterprise_when_geoip2_databases_are_used() { final LicenseTypeCheck objectUnderTest = createObjectUnderTest(); LicenseTypeOptions licenseTypeOptionsFree = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_ENTERPRISE); - assertThat(licenseTypeOptionsFree, not(equalTo(LicenseTypeOptions.FREE))); - - LicenseTypeOptions licenseTypeOptionsEnterprise = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_LITE2); - assertThat(licenseTypeOptionsEnterprise, not(equalTo(LicenseTypeOptions.ENTERPRISE))); + assertThat(licenseTypeOptionsFree, not(equalTo(LicenseTypeOptions.ENTERPRISE))); } } \ No newline at end of file