session-android/src/org/thoughtcrime/securesms/loki/database/LokiPreKeyBundleDatabase.kt

package org.thoughtcrime.securesms.loki.database

import android.content.ContentValues
import android.content.Context
import net.sqlcipher.Cursor
import org.thoughtcrime.securesms.crypto.IdentityKeyUtil
import org.thoughtcrime.securesms.crypto.PreKeyUtil
import org.thoughtcrime.securesms.database.Database
import org.thoughtcrime.securesms.database.DatabaseFactory
import org.thoughtcrime.securesms.database.helpers.SQLCipherOpenHelper
import org.thoughtcrime.securesms.logging.Log
import org.thoughtcrime.securesms.loki.utilities.get
import org.thoughtcrime.securesms.loki.utilities.getBase64EncodedData
import org.thoughtcrime.securesms.loki.utilities.getInt
import org.thoughtcrime.securesms.loki.utilities.insertOrUpdate
import org.thoughtcrime.securesms.util.Base64
import org.thoughtcrime.securesms.util.TextSecurePreferences
import org.whispersystems.libsignal.IdentityKey
import org.whispersystems.libsignal.InvalidKeyException
import org.whispersystems.libsignal.ecc.Curve
import org.whispersystems.libsignal.state.PreKeyBundle
import org.whispersystems.libsignal.util.KeyHelper
import org.whispersystems.signalservice.api.push.SignalServiceAddress
import org.whispersystems.signalservice.loki.database.LokiPreKeyBundleDatabaseProtocol

class LokiPreKeyBundleDatabase(context: Context, helper: SQLCipherOpenHelper) : Database(context, helper), LokiPreKeyBundleDatabaseProtocol {

    companion object {
        private val table = "loki_pre_key_bundle_database"
        private val publicKey = "public_key"
        private val preKeyID = "pre_key_id"
        private val preKeyPublic = "pre_key_public"
        private val signedPreKeyID = "signed_pre_key_id"
        private val signedPreKeyPublic = "signed_pre_key_public"
        private val signedPreKeySignature = "signed_pre_key_signature"
        private val identityKey = "identity_key"
        private val deviceID = "device_id"
        private val registrationID = "registration_id"
        @JvmStatic val createTableCommand = "CREATE TABLE $table (" + "$publicKey TEXT PRIMARY KEY," + "$preKeyID INTEGER," +
            "$preKeyPublic TEXT NOT NULL," + "$signedPreKeyID INTEGER," + "$signedPreKeyPublic TEXT NOT NULL," +
            "$signedPreKeySignature TEXT," + "$identityKey TEXT NOT NULL," + "$deviceID INTEGER," + "$registrationID INTEGER" + ");"
    }

    fun generatePreKeyBundle(publicKey: String): PreKeyBundle? {
        var failureCount = 0
        while (failureCount < 3) {
            try {
                val preKey = generatePreKeyBundle(publicKey, failureCount > 0) ?: return null
                // Verify the bundle is correct
                if (!Curve.verifySignature(preKey.identityKey.publicKey, preKey.signedPreKey.serialize(), preKey.signedPreKeySignature)) {
                    throw InvalidKeyException()
                }
                return preKey;
            } catch (e: InvalidKeyException) {
                failureCount += 1
            }
        }
        Log.w("Loki", "Failed to generate a valid pre key bundle for: $publicKey.")
        return null
    }

    private fun generatePreKeyBundle(publicKey: String, forceClean: Boolean): PreKeyBundle? {
        if (publicKey.isEmpty()) return null
        var registrationID = TextSecurePreferences.getLocalRegistrationId(context)
        if (registrationID == 0) {
            registrationID = KeyHelper.generateRegistrationId(false)
            TextSecurePreferences.setLocalRegistrationId(context, registrationID)
        }
        val deviceID = SignalServiceAddress.DEFAULT_DEVICE_ID
        val preKeyRecord = DatabaseFactory.getLokiPreKeyRecordDatabase(context).getOrCreatePreKeyRecord(publicKey)
        val identityKeyPair = IdentityKeyUtil.getIdentityKeyPair(context)
        if (!forceClean && TextSecurePreferences.isSignedPreKeyRegistered(context)) {
            Log.d("Loki", "A signed pre key has already been registered.")
        } else {
            Log.d("Loki", "Registering a new signed pre key.")
            PreKeyUtil.generateSignedPreKey(context, identityKeyPair, true)
            TextSecurePreferences.setSignedPreKeyRegistered(context, true)
        }
        val activeSignedPreKey = PreKeyUtil.getActiveSignedPreKey(context) ?: return null
        return PreKeyBundle(registrationID, deviceID, preKeyRecord.id, preKeyRecord.keyPair.publicKey, activeSignedPreKey.id, activeSignedPreKey.keyPair.publicKey, activeSignedPreKey.signature, identityKeyPair.publicKey)
    }

    override fun getPreKeyBundle(publicKey: String): PreKeyBundle? {
        val database = databaseHelper.readableDatabase
        return database.get(table, "${Companion.publicKey} = ?", arrayOf( publicKey )) { cursor ->
            val registrationID = cursor.getInt(registrationID)
            val deviceID = cursor.getInt(deviceID)
            val preKeyID = cursor.getInt(preKeyID)
            val preKey = Curve.decodePoint(cursor.getBase64EncodedData(preKeyPublic), 0)
            val signedPreKeyID = cursor.getInt(signedPreKeyID)
            val signedPreKey = Curve.decodePoint(cursor.getBase64EncodedData(signedPreKeyPublic), 0)
            val signedPreKeySignature = cursor.getBase64EncodedData(signedPreKeySignature)
            val identityKey = IdentityKey(cursor.getBase64EncodedData(identityKey), 0)
            PreKeyBundle(registrationID, deviceID, preKeyID, preKey, signedPreKeyID, signedPreKey, signedPreKeySignature, identityKey)
        }
    }

    fun setPreKeyBundle(publicKey: String, preKeyBundle: PreKeyBundle) {
        val database = databaseHelper.writableDatabase
        val values = ContentValues(9)
        values.put(registrationID, preKeyBundle.registrationId)
        values.put(deviceID, preKeyBundle.deviceId)
        values.put(preKeyID, preKeyBundle.preKeyId)
        values.put(preKeyPublic, Base64.encodeBytes(preKeyBundle.preKey.serialize()))
        values.put(signedPreKeyID, preKeyBundle.signedPreKeyId)
        values.put(signedPreKeyPublic, Base64.encodeBytes(preKeyBundle.signedPreKey.serialize()))
        values.put(signedPreKeySignature, Base64.encodeBytes(preKeyBundle.signedPreKeySignature))
        values.put(identityKey, Base64.encodeBytes(preKeyBundle.identityKey.serialize()))
        values.put(Companion.publicKey, publicKey)
        database.insertOrUpdate(table, values, "${Companion.publicKey} = ?", arrayOf( publicKey ))
    }

    override fun removePreKeyBundle(publicKey: String) {
        val database = databaseHelper.writableDatabase
        database.delete(table, "${Companion.publicKey} = ?", arrayOf( publicKey ))
    }

    fun hasPreKeyBundle(publicKey: String): Boolean {
        val database = databaseHelper.readableDatabase
        var cursor: Cursor? = null
        return try {
            cursor = database.query(table, null, "${Companion.publicKey} = ?", arrayOf( publicKey ), null, null, null)
            cursor != null && cursor.count > 0
        } catch (e: Exception) {
            false
        } finally {
          cursor?.close()
        }
    }
}