MaxmindDbEnrichmentLoader { protected void loadGeoLiteDatabase(CommandLine cli) throws IOException { System.out.println("Retrieving GeoLite2 archive"); String geo_url = GeoEnrichmentOptions.GEO_URL.get(cli, GEO_CITY_URL_DEFAULT); String asn_url = GeoEnrichmentOptions.ASN_URL.get(cli, ASN_URL_DEFAULT); String tmpDir = GeoEnrichmentOptions.TMP_DIR.get(cli, "/tmp") + "/"; int numRetries = Integer.parseInt(GeoEnrichmentOptions.RETRIES.get(cli, DEFAULT_RETRIES)); File localGeoFile = null; File localAsnFile = null; try { localGeoFile = downloadGeoFile(geo_url, tmpDir, numRetries); localAsnFile = downloadGeoFile(asn_url, tmpDir, numRetries); } catch (IllegalStateException ies) { System.err.println("Failed to download geo db file. Aborting"); System.exit(5); } localGeoFile.deleteOnExit(); localAsnFile.deleteOnExit(); System.out.println("GeoIP files downloaded successfully"); String zookeeper = GeoEnrichmentOptions.ZK_QUORUM.get(cli); long millis = LocalDateTime.now().toInstant(ZoneOffset.UTC).toEpochMilli(); String hdfsGeoLoc = GeoEnrichmentOptions.REMOTE_GEO_DIR.get(cli, "/apps/metron/geo/" + millis); System.out.println("Putting GeoLite City file into HDFS at: " + hdfsGeoLoc); Path srcPath = new Path(localGeoFile.getAbsolutePath()); Path dstPath = new Path(hdfsGeoLoc); putDbFile(srcPath, dstPath); pushConfig(srcPath, dstPath, GeoLiteCityDatabase.GEO_HDFS_FILE, zookeeper); String hdfsAsnLoc = GeoEnrichmentOptions.REMOTE_ASN_DIR.get(cli, "/apps/metron/asn/" + millis); System.out.println("Putting ASN file into HDFS at: " + hdfsAsnLoc); srcPath = new Path(localAsnFile.getAbsolutePath()); dstPath = new Path(hdfsAsnLoc); putDbFile(srcPath, dstPath); pushConfig(srcPath, dstPath, GeoLiteAsnDatabase.ASN_HDFS_FILE, zookeeper); System.out.println("GeoLite2 file placement complete"); System.out.println("Successfully created and updated new GeoLite information"); }  static void main(String... argv);  }
@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: "--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()))); }