maxmind integration (#1924)

Author: ag060

.github/workflows/prod.yml

@@ -66,6 +66,14 @@ jobs:
           wget -O filetypes.json https://raw.githubusercontent.com/akto-api-security/akto/master/pii-types/filetypes.json
           wget -O automated_api_groups.csv https://raw.githubusercontent.com/akto-api-security/akto/master/automated-api-groups/automated-api-groups.csv
 
+
+      - name: Create maxmind directory
+        run: mkdir -p ./apps/threat-detection-backend/src/main/resources/maxmind
+      - name: Download Maxmind Country database
+        working-directory: ./apps/threat-detection-backend/src/main/resources/maxmind
+        run: |
+          wget -O Geo-Country.mmdb https://raw.githubusercontent.com/akto-api-security/tests-library/refs/heads/master/resources/Geo-Country.mmdb
+
       - name: Prepare Dashboard polaris UI
         working-directory: ./apps/dashboard/web/polaris_web
         run: npm install && export RELEASE_VERSION=${{github.event.inputs.release_version}} && npm run build

.github/workflows/staging.yml

@@ -41,6 +41,12 @@ jobs:
           wget -O general.json https://raw.githubusercontent.com/akto-api-security/pii-types/master/general.json
           wget -O fintech.json https://raw.githubusercontent.com/akto-api-security/akto/master/pii-types/fintech.json
           wget -O filetypes.json https://raw.githubusercontent.com/akto-api-security/akto/master/pii-types/filetypes.json
+      - name: Create maxmind directory
+        run: mkdir -p ./apps/threat-detection-backend/src/main/resources/maxmind
+      - name: Download Maxmind Country database
+        working-directory: ./apps/threat-detection-backend/src/main/resources/maxmind
+        run: |
+          wget -O Geo-Country.mmdb https://raw.githubusercontent.com/akto-api-security/tests-library/refs/heads/master/resources/Geo-Country.mmdb
 
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v1

apps/threat-detection-backend/.gitignore

@@ -0,0 +1 @@
+*.mmdb

apps/threat-detection-backend/pom.xml

@@ -79,6 +79,13 @@
             <version>${vertex.version}</version>
         </dependency>
 
+        <!-- Maxmind GeoIP2 -->
+        <dependency>
+            <groupId>com.maxmind.geoip2</groupId>
+            <artifactId>geoip2</artifactId>
+            <version>2.15.0</version>
+        </dependency>
+
     </dependencies>
     <build>
         <plugins>
@@ -121,10 +128,6 @@
         <resources>
             <resource>
                 <directory>src/main/resources</directory>
-                <filtering>true</filtering>
-                <includes>
-                    <include>**/version.txt</include>
-                </includes>
             </resource>
         </resources>
     </build>
@@ -204,4 +207,4 @@
         </profile>
     </profiles>
 
-</project>
+</project>

apps/threat-detection-backend/src/main/java/com/akto/threat/backend/Main.java

@@ -7,6 +7,7 @@
 import com.akto.kafka.KafkaConfig;
 import com.akto.kafka.KafkaConsumerConfig;
 import com.akto.kafka.KafkaProducerConfig;
+import com.akto.threat.backend.client.IPLookupClient;
 import com.akto.threat.backend.service.MaliciousEventService;
 import com.akto.threat.backend.service.ThreatActorService;
 import com.akto.threat.backend.service.ThreatApiService;
@@ -17,6 +18,10 @@
 import com.mongodb.WriteConcern;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import org.apache.commons.io.IOUtils;
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.codecs.pojo.PojoCodecProvider;
 
@@ -53,14 +58,28 @@ public static void main(String[] args) throws Exception {
                 KafkaProducerConfig.newBuilder().setBatchSize(100).setLingerMs(1000).build())
             .build();
 
+    IPLookupClient ipLookupClient = new IPLookupClient(getMaxmindFile());
+
     new FlushMessagesToDB(internalKafkaConfig, threatProtectionMongo).run();
 
     MaliciousEventService maliciousEventService =
-        new MaliciousEventService(internalKafkaConfig, threatProtectionMongo);
+        new MaliciousEventService(internalKafkaConfig, threatProtectionMongo, ipLookupClient);
 
     ThreatActorService threatActorService = new ThreatActorService(threatProtectionMongo);
     ThreatApiService threatApiService = new ThreatApiService(threatProtectionMongo);
 
     new BackendVerticle(maliciousEventService, threatActorService, threatApiService).start();
   }
+
+  private static File getMaxmindFile() throws IOException {
+    File maxmindTmpFile = File.createTempFile("tmp-geo-country", ".mmdb");
+    maxmindTmpFile.deleteOnExit();
+
+    try (FileOutputStream fos = new FileOutputStream(maxmindTmpFile)) {
+      IOUtils.copy(
+          Main.class.getClassLoader().getResourceAsStream("maxmind/Geo-Country.mmdb"), fos);
+    }
+
+    return maxmindTmpFile;
+  }
 }

apps/threat-detection-backend/src/main/java/com/akto/threat/backend/client/IPLookupClient.java

@@ -0,0 +1,26 @@
+package com.akto.threat.backend.client;
+
+import com.maxmind.geoip2.DatabaseReader;
+import com.maxmind.geoip2.model.CountryResponse;
+import java.io.File;
+import java.io.IOException;
+import java.net.InetAddress;
+import java.util.Optional;
+
+public class IPLookupClient {
+  private final DatabaseReader db;
+
+  public IPLookupClient(File dbFile) throws IOException {
+    this.db = new DatabaseReader.Builder(dbFile).build();
+  }
+
+  public Optional<String> getCountryISOCodeGivenIp(String ip) {
+    try {
+      InetAddress ipAddr = InetAddress.getByName(ip);
+      CountryResponse resp = db.country(ipAddr);
+      return Optional.of(resp.getCountry().getIsoCode());
+    } catch (Exception e) {
+      return Optional.empty();
+    }
+  }
+}

apps/threat-detection-backend/src/main/java/com/akto/threat/backend/service/MaliciousEventService.java

@@ -11,6 +11,7 @@
 import com.akto.proto.generated.threat_detection.service.dashboard_service.v1.ListMaliciousRequestsRequest;
 import com.akto.proto.generated.threat_detection.service.dashboard_service.v1.ListMaliciousRequestsResponse;
 import com.akto.proto.generated.threat_detection.service.malicious_alert_service.v1.RecordMaliciousEventRequest;
+import com.akto.threat.backend.client.IPLookupClient;
 import com.akto.threat.backend.constants.KafkaTopic;
 import com.akto.threat.backend.constants.MongoDBCollection;
 import com.akto.threat.backend.db.AggregateSampleMaliciousEventModel;
@@ -33,10 +34,13 @@ public class MaliciousEventService {
 
   private final Kafka kafka;
   private MongoClient mongoClient;
+  private IPLookupClient ipLookupClient;
 
-  public MaliciousEventService(KafkaConfig kafkaConfig, MongoClient mongoClient) {
+  public MaliciousEventService(
+      KafkaConfig kafkaConfig, MongoClient mongoClient, IPLookupClient ipLookupClient) {
     this.kafka = new Kafka(kafkaConfig);
     this.mongoClient = mongoClient;
+    this.ipLookupClient = ipLookupClient;
   }
 
   public void recordMaliciousEvent(String accountId, RecordMaliciousEventRequest request) {
@@ -64,7 +68,8 @@ public void recordMaliciousEvent(String accountId, RecordMaliciousEventRequest r
             .setLatestApiCollectionId(evt.getLatestApiCollectionId())
             .setEventType(maliciousEventType)
             .setLatestApiIp(evt.getLatestApiIp())
-            .setCountry("US")
+            .setCountry(
+                this.ipLookupClient.getCountryISOCodeGivenIp(evt.getLatestApiIp()).orElse(""))
             .setCategory(evt.getCategory())
             .setSubCategory(evt.getSubCategory())
             .build();

apps/threat-detection-backend/src/main/java/com/akto/threat/backend/tasks/FlushMessagesToDB.java

@@ -71,8 +71,8 @@ public void run() {
   private void processRecords(ConsumerRecords<String, String> records) {
     records.forEach(
         r -> {
-          String message = r.value();
           try {
+            String message = r.value();
             writeMessage(message);
           } catch (JsonProcessingException e) {
             System.out.println("Error while parsing message" + e);