From 51f037410904481ae55b216f73379ef37abbc7c3 Mon Sep 17 00:00:00 2001
From: nielsandriesse <andriesseniels@gmail.com>
Date: Tue, 2 Jun 2020 15:00:37 +1000
Subject: [PATCH] Implement IP2Country

---
 build.gradle                                  |   2 +
 .../securesms/loki/utilities/IP2Country.kt    | 116 ++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 src/org/thoughtcrime/securesms/loki/utilities/IP2Country.kt

diff --git a/build.gradle b/build.gradle
index ad5c0ac04d..accac8f19d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -8,6 +8,7 @@ buildscript {
     ext.rss_parser_version = "2.0.4"
     ext.google_services_version = "4.3.3"
     ext.firebase_messaging_version = "18.0.0"
+    ext.opencsv_version = "4.6"
 
     repositories {
         mavenLocal()
@@ -188,6 +189,7 @@ dependencies {
     implementation "com.jakewharton.rxbinding3:rxbinding:3.1.0"
     implementation "com.github.tbruyelle:rxpermissions:0.10.2"
     implementation "com.github.ybq:Android-SpinKit:1.4.0"
+    implementation "com.opencsv:opencsv:$opencsv_version"
 }
 
 def canonicalVersionCode = 52
diff --git a/src/org/thoughtcrime/securesms/loki/utilities/IP2Country.kt b/src/org/thoughtcrime/securesms/loki/utilities/IP2Country.kt
new file mode 100644
index 0000000000..2181e70b45
--- /dev/null
+++ b/src/org/thoughtcrime/securesms/loki/utilities/IP2Country.kt
@@ -0,0 +1,116 @@
+package org.thoughtcrime.securesms.loki.utilities
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.support.v4.content.LocalBroadcastManager
+import android.util.Log
+import com.opencsv.CSVReader
+import org.whispersystems.signalservice.loki.api.onionrequests.OnionRequestAPI
+import java.io.File
+import java.io.FileOutputStream
+import java.io.FileReader
+
+class IP2Country private constructor(private val context: Context) {
+    private val pathsBuiltEventReceiver: BroadcastReceiver
+    private val countryNamesCache = mutableMapOf<String, String>()
+
+    private val ipv4Table by lazy {
+        loadFile("geolite2_country_blocks_ipv4.csv")
+    }
+
+    private val countryNamesTable by lazy {
+        loadFile("geolite2_country_locations_english.csv")
+    }
+
+    // region Initialization
+    companion object {
+
+        public lateinit var shared: IP2Country
+
+        public fun configureIfNeeded(context: Context) {
+            if (::shared.isInitialized) { return; }
+            shared = IP2Country(context)
+        }
+    }
+
+    init {
+        preloadCountriesIfNeeded()
+        pathsBuiltEventReceiver = object : BroadcastReceiver() {
+
+            override fun onReceive(context: Context, intent: Intent) {
+                preloadCountriesIfNeeded()
+            }
+        }
+        LocalBroadcastManager.getInstance(context).registerReceiver(pathsBuiltEventReceiver, IntentFilter("pathsBuilt"))
+    }
+
+    // TODO: Deinit?
+    // endregion
+
+    // region Implementation
+    private fun loadFile(fileName: String): File {
+        val directory = File(context.applicationInfo.dataDir)
+        val file = File(directory, fileName)
+        if (directory.list().contains(fileName)) { return file }
+        val inputStream = context.assets.open("csv/$fileName")
+        val outputStream = FileOutputStream(file)
+        val buffer = ByteArray(1024)
+        while (true) {
+            val count = inputStream.read(buffer)
+            if (count < 0) { break }
+            outputStream.write(buffer, 0, count)
+        }
+        inputStream.close()
+        outputStream.close()
+        return file
+    }
+
+    private fun getCountry(ip: String): String {
+        var truncatedIP = ip
+        fun getCountryInternal(): String {
+            val country = countryNamesCache[ip]
+            if (country != null) { return country }
+            val ipv4TableReader = CSVReader(FileReader(ipv4Table.absoluteFile))
+            val countryNamesTableReader = CSVReader(FileReader(ipv4Table.absoluteFile))
+            var ipv4TableLine = ipv4TableReader.readNext()
+            while (ipv4TableLine != null) {
+                if (!ipv4TableLine[0].startsWith(truncatedIP)) {
+                    ipv4TableLine = ipv4TableReader.readNext()
+                    continue
+                }
+                val countryID = ipv4TableLine[1]
+                var countryNamesTableLine = countryNamesTableReader.readNext()
+                while (countryNamesTableLine != null) {
+                    if (countryNamesTableLine[0] != countryID) {
+                        countryNamesTableLine = countryNamesTableReader.readNext()
+                        continue
+                    }
+                    @Suppress("NAME_SHADOWING") val country = countryNamesTableLine[5]
+                    countryNamesCache[ip] = country
+                    return country
+                }
+            }
+            if (truncatedIP.contains(".") && !truncatedIP.endsWith(".")) { // The fuzziest we want to go is xxx.x
+                truncatedIP.dropLast(1)
+                if (truncatedIP.endsWith(".")) { truncatedIP.dropLast(1) }
+                return getCountryInternal()
+            } else {
+                return "Unknown Country"
+            }
+        }
+        return getCountryInternal()
+    }
+
+    private fun preloadCountriesIfNeeded() {
+        Thread {
+            val path = OnionRequestAPI.paths.firstOrNull() ?: return@Thread
+            path.forEach { snode ->
+                getCountry(snode.ip) // Preload if needed
+            }
+            Log.d("Loki", "Finished preloading onion request path countries.")
+        }.start()
+    }
+    // endregion
+}
\ No newline at end of file