metron-platform/metron-data-management/src/test/java/org/apache/metron/dataloads/nonbulk/geo/MaxmindDbEnrichmentLoaderTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.metron.dataloads.nonbulk.geo;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.PosixParser;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.util.GenericOptionsParser;
import org.apache.metron.common.utils.CompressionStrategies;
import org.apache.metron.integration.utils.TestUtils;
import org.junit.Rule;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.migrationsupport.rules.EnableRuleMigrationSupport;
import org.junit.rules.TemporaryFolder;

import java.io.File;

import static org.junit.jupiter.api.Assertions.*;

@EnableRuleMigrationSupport
public class MaxmindDbEnrichmentLoaderTest {
  private class MockMaxmindDbEnrichmentLoader extends MaxmindDbEnrichmentLoader {
    @Override
    protected void pushConfig(Path srcPath, Path dstPath, String configName, String zookeeper) {
    }
  }

  @Rule
  public TemporaryFolder testFolder = new TemporaryFolder();

  private File remoteDir;
  private File tmpDir;

  @BeforeEach
  public void setup() throws Exception {
      testFolder.create();
      remoteDir = testFolder.newFolder("remoteDir");
      tmpDir = testFolder.newFolder("tmpDir");
  }

  @Test
  public void testCommandLineShortOpts() throws Exception {
    String[] argv = {
        "-g testGeoUrl",
        "-a testAsnUrl",
        "-r /test/remoteDirGeo",
        "-ra", "/test/remoteDirAsn",
        "-t /test/tmpDir",
        "-z test:2181"
    };
    String[] otherArgs = new GenericOptionsParser(argv).getRemainingArgs();

    CommandLine cli = MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.parse(new PosixParser(), otherArgs);
    assertEquals("testGeoUrl", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.GEO_URL.get(cli).trim());
    assertEquals("testAsnUrl", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.ASN_URL.get(cli).trim());
    assertEquals("/test/remoteDirGeo", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.REMOTE_GEO_DIR.get(cli).trim());
    assertEquals("/test/remoteDirAsn", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.REMOTE_ASN_DIR.get(cli).trim());
    assertEquals("/test/tmpDir", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.TMP_DIR.get(cli).trim());
    assertEquals("test:2181", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.ZK_QUORUM.get(cli).trim());
  }

  @Test
  public void testCommandLineLongOpts() throws Exception {
    String[] argv = {
        "--geo_url", "testGeoUrl",
        "--remote_dir", "/test/remoteDir",
        "-ra", "/test/remoteDir",
        "--tmp_dir", "/test/tmpDir",
        "--zk_quorum", "test:2181"
    };
    String[] otherArgs = new GenericOptionsParser(argv).getRemainingArgs();

    CommandLine cli = MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.parse(new PosixParser(), otherArgs);
    assertEquals("testGeoUrl", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.GEO_URL.get(cli).trim());
    assertEquals("/test/remoteDir", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.REMOTE_GEO_DIR.get(cli).trim());
    assertEquals("/test/tmpDir", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.TMP_DIR.get(cli).trim());
    assertEquals("test:2181", MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.ZK_QUORUM.get(cli).trim());
  }

  @Test
  public void testLoadGeoIpDatabase() throws Exception {
    File dbPlainTextFile = new File(remoteDir.getAbsolutePath() + "/MaxmindDbEnrichmentLoaderTest.mmdb");
    TestUtils.write(dbPlainTextFile, "hello world");
    File dbFile = new File(remoteDir.getAbsolutePath() + "/MaxmindDbEnrichmentLoaderTest.mmdb.gz");
    CompressionStrategies.GZIP.compress(dbPlainTextFile, dbFile);
    String[] argv = {
        "--geo_url", "file://" + dbFile.getAbsolutePath(),
        "--remote_dir", remoteDir.getAbsolutePath(),
        "--remote_asn_dir", remoteDir.getAbsolutePath(),
        "--tmp_dir", tmpDir.getAbsolutePath(),
        "--zk_quorum", "test:2181"
    };
    String[] otherArgs = new GenericOptionsParser(argv).getRemainingArgs();
    CommandLine cli = MaxmindDbEnrichmentLoader.GeoEnrichmentOptions.parse(new PosixParser(), otherArgs);

    MaxmindDbEnrichmentLoader loader = new MockMaxmindDbEnrichmentLoader();
    loader.loadGeoLiteDatabase(cli);
    Configuration config = new Configuration();
    FileSystem fs = FileSystem.get(config);
    assertTrue(fs.exists(new Path(remoteDir + "/" + dbFile.getName())));
  }

  @Test
  public void loader_throws_exception_on_bad_gzip_file() throws Exception {
    File dbFile = new File(remoteDir.getAbsolutePath() + "/MaxmindDbEnrichmentLoaderTest.mmdb");
    dbFile.createNewFile();

    String geoUrl = "file://" + dbFile.getAbsolutePath();
    int numRetries = 2;
    MaxmindDbEnrichmentLoader loader = new MockMaxmindDbEnrichmentLoader();
    IllegalStateException e = assertThrows(IllegalStateException.class, () -> loader.downloadGeoFile(geoUrl, tmpDir.getAbsolutePath(), numRetries));
    assertEquals("Unable to download geo enrichment database.", e.getMessage());
  }

}