Geoip processor update

graylog2-server/src/main/java/org/graylog/plugins/map/MapWidgetModule.java

@@ -16,6 +16,17 @@
  */
 package org.graylog.plugins.map;
 
+import com.google.inject.TypeLiteral;
+import com.google.inject.assistedinject.FactoryModuleBuilder;
+import com.google.inject.name.Names;
+import org.graylog.plugins.map.geoip.GeoAsnInformation;
+import org.graylog.plugins.map.geoip.GeoIpResolver;
+import org.graylog.plugins.map.geoip.GeoIpResolverFactory;
+import org.graylog.plugins.map.geoip.GeoLocationInformation;
+import org.graylog.plugins.map.geoip.IpInfoIpAsnResolver;
+import org.graylog.plugins.map.geoip.IpInfoLocationResolver;
+import org.graylog.plugins.map.geoip.MaxMindIpAsnResolver;
+import org.graylog.plugins.map.geoip.MaxMindIpLocationResolver;
 import org.graylog.plugins.map.geoip.MaxmindDataAdapter;
 import org.graylog.plugins.map.geoip.processor.GeoIpProcessor;
 import org.graylog2.plugin.PluginModule;
@@ -28,5 +39,18 @@ protected void configure() {
                 MaxmindDataAdapter.class,
                 MaxmindDataAdapter.Factory.class,
                 MaxmindDataAdapter.Config.class);
+
+        //Create TypeLiterals to specify method type parameters
+        TypeLiteral<GeoIpResolver<GeoLocationInformation>> mmCityTl = new TypeLiteral<GeoIpResolver<GeoLocationInformation>>() {};
+        TypeLiteral<GeoIpResolver<GeoAsnInformation>> mmAsnTl = new TypeLiteral<GeoIpResolver<GeoAsnInformation>>() {};
+        TypeLiteral<GeoIpResolver<GeoLocationInformation>> ipinfoCityTl = new TypeLiteral<GeoIpResolver<GeoLocationInformation>>() {};
+        TypeLiteral<GeoIpResolver<GeoAsnInformation>> ipInfoAsnTl = new TypeLiteral<GeoIpResolver<GeoAsnInformation>>() {};
+
+        install(new FactoryModuleBuilder()
+                .implement(mmCityTl, Names.named("MAXMIND_CITY"), MaxMindIpLocationResolver.class)
+                .implement(mmAsnTl, Names.named("MAXMIND_ASN"), MaxMindIpAsnResolver.class)
+                .implement(ipinfoCityTl, Names.named("IPINFO_CITY"), IpInfoLocationResolver.class)
+                .implement(ipInfoAsnTl, Names.named("IPINFO_ASN"), IpInfoIpAsnResolver.class)
+                .build(GeoIpResolverFactory.class));
     }
 }

graylog2-server/src/main/java/org/graylog/plugins/map/config/DatabaseVendorType.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+
+package org.graylog.plugins.map.config;
+
+public enum DatabaseVendorType {
+    MAXMIND(DatabaseType.MAXMIND_CITY, DatabaseType.MAXMIND_ASN),
+    IPINFO(DatabaseType.IPINFO_STANDARD_LOCATION, DatabaseType.IPINFO_ASN);
+
+    private final DatabaseType cityDbType;
+    private final DatabaseType asnDbType;
+
+    DatabaseVendorType(DatabaseType cityDbType, DatabaseType asnDbType) {
+        this.cityDbType = cityDbType;
+        this.asnDbType = asnDbType;
+    }
+
+    public DatabaseType getCityDbType() {
+        return cityDbType;
+    }
+
+    public DatabaseType getAsnDbType() {
+        return asnDbType;
+    }
+}

graylog2-server/src/main/java/org/graylog/plugins/map/config/GeoIpResolverConfig.java

@@ -16,12 +16,11 @@
  */
 package org.graylog.plugins.map.config;
 
-import com.google.auto.value.AutoValue;
-
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.auto.value.AutoValue;
 
 @JsonAutoDetect
 @JsonIgnoreProperties(ignoreUnknown = true)
@@ -31,28 +30,40 @@ public abstract class GeoIpResolverConfig {
     @JsonProperty("enabled")
     public abstract boolean enabled();
 
-    @JsonProperty("db_type")
-    public abstract DatabaseType dbType();
+    @JsonProperty("enforce_graylog_schema")
+    public abstract boolean enforceGraylogSchema();
+
+    @JsonProperty("db_vendor_type")
+    public abstract DatabaseVendorType databaseVendorType();
+
+    @JsonProperty("city_db_path")
+    public abstract String cityDbPath();
 
-    @JsonProperty("db_path")
-    public abstract String dbPath();
+    @JsonProperty("asn_db_path")
+    public abstract String asnDbPath();
 
     @JsonCreator
-    public static GeoIpResolverConfig create(@JsonProperty("enabled") boolean enabled,
-                                             @JsonProperty("db_type") DatabaseType dbType,
-                                             @JsonProperty("db_path") String dbPath) {
+    public static GeoIpResolverConfig create(@JsonProperty("enabled") boolean cityEnabled,
+                                             @JsonProperty("enforce_graylog_schema") boolean enforceGraylogSchema,
+                                             @JsonProperty("db_vendor_type") DatabaseVendorType databaseVendorType,
+                                             @JsonProperty("city_db_path") String cityDbPath,
+                                             @JsonProperty("asn_db_path") String asnDbPath) {
         return builder()
-                .enabled(enabled)
-                .dbType(dbType)
-                .dbPath(dbPath)
+                .enabled(cityEnabled)
+                .enforceGraylogSchema(enforceGraylogSchema)
+                .databaseVendorType(databaseVendorType == null ? DatabaseVendorType.MAXMIND : databaseVendorType)
+                .cityDbPath(cityDbPath)
+                .asnDbPath(asnDbPath)
                 .build();
     }
 
     public static GeoIpResolverConfig defaultConfig() {
        return builder()
                .enabled(false)
-               .dbType(DatabaseType.MAXMIND_CITY)
-               .dbPath("/etc/graylog/server/GeoLite2-City.mmdb")
+               .databaseVendorType(DatabaseVendorType.MAXMIND)
+               .enforceGraylogSchema(false)
+               .cityDbPath("/etc/graylog/server/GeoLite2-City.mmdb")
+               .asnDbPath("/etc/graylog/server/GeoLite2-ASN.mmdb")
                .build();
     }
 
@@ -63,11 +74,17 @@ public static Builder builder() {
     public abstract Builder toBuilder();
 
     @AutoValue.Builder
-    public static abstract class Builder {
+    public abstract static class Builder {
         public abstract Builder enabled(boolean enabled);
-        public abstract Builder dbType(DatabaseType dbType);
-        public abstract Builder dbPath(String dbPath);
+
+        public abstract Builder enforceGraylogSchema(boolean enforce);
+
+        public abstract Builder databaseVendorType(DatabaseVendorType type);
+
+        public abstract Builder cityDbPath(String dbPath);
+
+        public abstract Builder asnDbPath(String asnDBPath);
 
         public abstract GeoIpResolverConfig build();
     }
-}
+}

graylog2-server/src/main/java/org/graylog/plugins/map/geoip/GeoAsnInformation.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+
+package org.graylog.plugins.map.geoip;
+
+import com.google.auto.value.AutoValue;
+
+@AutoValue
+public abstract class GeoAsnInformation {
+
+    public abstract String organization();
+
+    public abstract String type();
+
+    public abstract String asn();
+
+    public static GeoAsnInformation create(String organization, String type, String asn) {
+        return new AutoValue_GeoAsnInformation(organization, type, asn);
+    }
+}

graylog2-server/src/main/java/org/graylog/plugins/map/geoip/GeoIpResolver.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+
+package org.graylog.plugins.map.geoip;
+
+import com.codahale.metrics.Timer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.net.InetAddress;
+import java.nio.file.Files;
+import java.util.Optional;
+
+public abstract class GeoIpResolver<V> {
+
+    private static final Logger LOG = LoggerFactory.getLogger(GeoIpResolver.class);
+
+    protected String lastError = null;
+    protected final Timer resolveTime;
+    private final boolean enabled;
+
+    GeoIpResolver(Timer resolveTime, String configPath, boolean enabled) {
+
+        this.resolveTime = resolveTime;
+        if (enabled) {
+            final File configFile = new File(configPath);
+            if (Files.exists(configFile.toPath())) {
+                this.enabled = createDataProvider(configFile);
+            } else {
+                LOG.warn("'{}' database file does not exist: {}", getClass().getName(), configPath);
+                this.enabled = false;
+            }
+        } else {
+            this.enabled = false;
+        }
+    }
+
+    public boolean isEnabled() {
+        return enabled;
+    }
+
+    abstract boolean createDataProvider(File configFile);
+
+    public Optional<V> getGeoIpData(InetAddress address) {
+        lastError = null;
+        if (!enabled || address == null) {
+            return Optional.empty();
+        }
+        return doGetGeoIpData(address);
+    }
+
+    /**
+     * Get the last error, if any, produced after having called {@link #getGeoIpData(InetAddress)}.
+     *
+     * @return optional error message
+     */
+    public Optional<String> getLastError() {
+        return Optional.ofNullable(lastError);
+    }
+
+    protected abstract Optional<V> doGetGeoIpData(InetAddress address);
+}