Kecalek/app/src/main/java/com/kecalek/chat/crypto/RSACrypto.kt

package com.kecalek.chat.crypto

import java.security.KeyFactory
import java.security.KeyPairGenerator
import java.security.Signature
import java.security.interfaces.RSAPrivateKey
import java.security.interfaces.RSAPublicKey
import java.security.spec.MGF1ParameterSpec
import java.security.spec.PKCS8EncodedKeySpec
import java.security.spec.PSSParameterSpec
import java.security.spec.X509EncodedKeySpec

/**
 * RSA-4096 for login challenge-response only.
 * Uses RSA-PSS with SHA-256, MGF1-SHA256.
 *
 * Private key storage: DER PKCS8 raw bytes encrypted via ECP1.
 * Public key: DER SubjectPublicKeyInfo (X.509).
 *
 * Compatible with Python generate_rsa_keypair, rsa_sign, rsa_verify.
 * Sign uses PSS with salt_length=MAX. Verify accepts MAX or hash-length salt.
 */
object RSACrypto {

    private const val KEY_SIZE = 4096

    /**
     * Generate RSA-4096 keypair.
     */
    fun generateKeypair(): Pair<RSAPrivateKey, RSAPublicKey> {
        val kpg = KeyPairGenerator.getInstance("RSA")
        kpg.initialize(KEY_SIZE)
        val kp = kpg.generateKeyPair()
        return Pair(kp.private as RSAPrivateKey, kp.public as RSAPublicKey)
    }

    /**
     * Serialize private key to DER PKCS8 format.
     * Optionally encrypt with password using ECP1.
     */
    fun serializePrivate(key: RSAPrivateKey, password: String? = null): ByteArray {
        val der = key.encoded // PKCS8 DER
        return if (password != null) {
            KeyEncryption.encrypt(der, password)
        } else {
            der
        }
    }

    /**
     * Serialize public key to DER X.509 format.
     */
    fun serializePublic(key: RSAPublicKey): ByteArray {
        return key.encoded // X.509 DER
    }

    /**
     * Load private key from DER bytes (optionally ECP1-encrypted).
     */
    fun loadPrivate(data: ByteArray, password: String? = null): RSAPrivateKey {
        val der = if (KeyEncryption.isEcp1Format(data) && password != null) {
            KeyEncryption.decrypt(data, password)
        } else {
            data
        }
        val keyFactory = KeyFactory.getInstance("RSA")
        return keyFactory.generatePrivate(PKCS8EncodedKeySpec(der)) as RSAPrivateKey
    }

    /**
     * Load public key from DER X.509 bytes.
     */
    fun loadPublic(data: ByteArray): RSAPublicKey {
        val keyFactory = KeyFactory.getInstance("RSA")
        return keyFactory.generatePublic(X509EncodedKeySpec(data)) as RSAPublicKey
    }

    /**
     * Sign data with RSA-PSS (SHA-256, MGF1-SHA256, max salt length).
     * Compatible with Python rsa_sign.
     */
    fun sign(privateKey: RSAPrivateKey, data: ByteArray): ByteArray {
        // Max salt length = key size in bytes - hash size - 2
        val maxSaltLen = privateKey.modulus.bitLength() / 8 - 32 - 2
        val pssSpec = PSSParameterSpec(
            "SHA-256",
            "MGF1",
            MGF1ParameterSpec.SHA256,
            maxSaltLen,
            1, // trailer field
        )
        val sig = Signature.getInstance("RSASSA-PSS")
        sig.setParameter(pssSpec)
        sig.initSign(privateKey)
        sig.update(data)
        return sig.sign()
    }

    /**
     * Verify RSA-PSS signature.
     * Uses salt_length = max for verification (Java PSS handles this internally).
     * For cross-platform compat, we try max salt first, then hash-length salt.
     */
    fun verify(publicKey: RSAPublicKey, signature: ByteArray, data: ByteArray): Boolean {
        // Try with max salt length first (Python's default for signing)
        val maxSaltLen = publicKey.modulus.bitLength() / 8 - 32 - 2
        if (verifyWithSaltLen(publicKey, signature, data, maxSaltLen)) return true
        // Try with hash-length salt (iOS compatibility)
        if (verifyWithSaltLen(publicKey, signature, data, 32)) return true
        return false
    }

    private fun verifyWithSaltLen(
        publicKey: RSAPublicKey,
        signature: ByteArray,
        data: ByteArray,
        saltLen: Int,
    ): Boolean {
        return try {
            val pssSpec = PSSParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec.SHA256,
                saltLen,
                1,
            )
            val sig = Signature.getInstance("RSASSA-PSS")
            sig.setParameter(pssSpec)
            sig.initVerify(publicKey)
            sig.update(data)
            sig.verify(signature)
        } catch (_: Exception) {
            false
        }
    }
}