package com.jihe.punchnet.sdlan.utils

import com.jihe.punchnet.sdlan.config.RSAConfig
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.openssl.PEMParser
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter
import org.bouncycastle.openssl.jcajce.JcaPEMWriter
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.FileReader
import java.io.FileWriter
import java.security.KeyPair
import java.security.KeyPairGenerator
import java.security.PrivateKey
import java.security.PublicKey
import java.security.Security
import javax.crypto.Cipher

// 添加BouncyCastle提供者
fun setupBouncyCastle() {
    Security.addProvider(BouncyCastleProvider())
}

fun savePublicKeyToPem(publicKey: PublicKey, fileName: String) {
    FileWriter(fileName).use { fileWriter ->
        JcaPEMWriter(fileWriter).use { pemWriter ->
            pemWriter.writeObject(publicKey)
        }
    }
}

fun savePrivateKeyToPem(privateKey: PrivateKey, fileName: String) {
    FileWriter(fileName).use { fileWriter ->
        JcaPEMWriter(fileWriter).use { pemWriter ->
            pemWriter.writeObject(privateKey)
        }
    }
}

fun loadPublicKeyFromPem(fileName: String): PublicKey {
    FileReader(fileName).use { fileReader ->
        PEMParser(fileReader).use { pemParser ->
            val converter = JcaPEMKeyConverter()
            val keypair = pemParser.readObject()
            return converter.getPublicKey(keypair as SubjectPublicKeyInfo)
        }
    }
}

fun loadPrivateKeyFromPem(fileName: String): PrivateKey {
    FileReader(fileName).use { fileReader ->
        PEMParser(fileReader).use { pemParser ->
            val converter = JcaPEMKeyConverter()
            val keypair = pemParser.readObject() as org.bouncycastle.openssl.PEMKeyPair
            return converter.getPrivateKey(keypair.privateKeyInfo)
        }
    }

}

// 生成RSA密钥对
fun generateRSAKeyPair(): KeyPair {
    val keyPairGenerator = KeyPairGenerator.getInstance("RSA", "BC")
    keyPairGenerator.initialize(2048)
    return keyPairGenerator.generateKeyPair()
}

class RSA private constructor () {
    private var pubkey: PublicKey? = null
    private var prikey: PrivateKey? = null
    private var pubkeyStr: String? = null

    companion object {
        private var instance: RSA? = null
            get() {
                if (field == null) {
                    field = RSA()
                }
                return field
            }

        @Synchronized
        fun getRSA(): RSA {
            return instance!!
        }
    }

    fun generateKeyPair(dirname: String) {
        setupBouncyCastle()
        val pubFile = File(dirname, RSAConfig.PUBLIC_FILE_NAME)
        val priFile = File(dirname, RSAConfig.PRIVATE_FILE_NAME)

        val dirpath = File(dirname)
        dirpath.mkdirs()

        if (pubFile.exists() && priFile.exists()) {
            pubkey = loadPublicKeyFromPem(pubFile.path)
            prikey = loadPrivateKeyFromPem(priFile.path)
        } else {
            // generate public and private file
            val keypair = generateRSAKeyPair()

            savePublicKeyToPem(keypair.public, pubFile.path)
            savePrivateKeyToPem(keypair.private, priFile.path)

            pubkey = keypair.public
            prikey = keypair.private
        }

        pubkeyStr = pubFile.readText()
        println("pub key is ${pubkeyStr}")
    }

    fun encrypt(input: ByteArray, use_private_key: Boolean = true): ByteArray {
        val cipher = Cipher.getInstance(RSAConfig.RSAMODE)
        cipher.init(Cipher.ENCRYPT_MODE, if (use_private_key) prikey else pubkey)

        var temp: ByteArray? = null
        var offset = 0
        val outputStream = ByteArrayOutputStream()

        while(input.size - offset > 0) {
            if (input.size - offset >= RSAConfig.ENCRYPT_MAX_SIZE) {
                temp = cipher.doFinal(input, offset, RSAConfig.ENCRYPT_MAX_SIZE)
                offset += RSAConfig.ENCRYPT_MAX_SIZE
            } else {
                temp = cipher.doFinal(input, offset, input.size - offset)
                offset = input.size
            }
            outputStream.write(temp!!)
        }
        outputStream.close()
        return outputStream.toByteArray()
    }

    fun decrypt(ciphered: ByteArray, use_pub_key: Boolean = false): ByteArray? {
        val cipher = Cipher.getInstance(RSAConfig.RSAMODE)
        cipher.init(Cipher.DECRYPT_MODE, if (use_pub_key) pubkey else prikey)

        var temp: ByteArray? = null
        var offset = 0
        val outputStream = ByteArrayOutputStream()

        try {
            while(ciphered.size - offset > 0) {
                if (ciphered.size - offset >= RSAConfig.DECRYPT_MAX_SIZE) {
                    temp = cipher.doFinal(ciphered, offset, RSAConfig.DECRYPT_MAX_SIZE)
                    offset += RSAConfig.DECRYPT_MAX_SIZE
                } else {
                    temp = cipher.doFinal(ciphered, offset, ciphered.size - offset)
                    offset = ciphered.size
                }
                outputStream.write(temp!!)
            }
            outputStream.close()
            return outputStream.toByteArray()
        } catch (e: Exception) {
            return null
        }
    }

    fun getPublicKeyString(): String {
        return pubkeyStr!!
    }

}
