package com.dano.test1

import java.io.File
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.Mac
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec
import kotlin.math.min

object AES256Helper {

    private fun hkdfFromToken(tokenHex: String, info: String = "qdb-aes", len: Int = 32): ByteArray {
        val ikm = hexToBytes(tokenHex)
        val mac = Mac.getInstance("HmacSHA256")
        val zeroSalt = ByteArray(32) { 0 }
        mac.init(SecretKeySpec(zeroSalt, "HmacSHA256"))
        val prk = mac.doFinal(ikm)

        var previous = ByteArray(0)
        val okm = ByteArray(len)
        var generated = 0
        var counter = 1
        while (generated < len) {
            mac.init(SecretKeySpec(prk, "HmacSHA256"))
            mac.update(previous)
            mac.update(info.toByteArray(Charsets.UTF_8))
            mac.update(counter.toByte())
            val t = mac.doFinal()
            val toCopy = min(len - generated, t.size)
            System.arraycopy(t, 0, okm, generated, toCopy)
            previous = t
            generated += toCopy
            counter++
        }
        return okm
    }

    private fun hexToBytes(hex: String): ByteArray {
        val clean = hex.trim()
        val len = clean.length
        val out = ByteArray(len / 2)
        var i = 0
        while (i < len) {
            out[i / 2] = ((Character.digit(clean[i], 16) shl 4) + Character.digit(clean[i + 1], 16)).toByte()
            i += 2
        }
        return out
    }

    fun encryptFileWithToken(inFile: File, outFile: File, token: String) {
        val key = hkdfFromToken(token)
        val iv = ByteArray(16).also { SecureRandom().nextBytes(it) }
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
        cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
        val plain = inFile.readBytes()
        val enc = cipher.doFinal(plain)
        outFile.writeBytes(iv + enc)
    }

    fun decryptFileWithToken(inFile: File, token: String): ByteArray {
        val key = hkdfFromToken(token)
        val data = inFile.readBytes()
        require(data.size >= 16) { "cipher too short" }
        val iv = data.copyOfRange(0, 16)
        val ct = data.copyOfRange(16, data.size)
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
        return cipher.doFinal(ct)
    }
}
