Chatto-Desktop-ScalaFX/src/main/scala/wow/doge/chatto/service/EncryptionServiceImpl.scala

package wow.doge.chatto.service

import wow.doge.chatto.model.MessageCipher;
import java.util.Base64;

import javax.crypto.Cipher;;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom
import scala.util.Try

class EncryptionServiceImpl extends EncryptionService {

  override def encrypt(
      password: String,
      plainText: String
  ): String Either MessageCipher = {
    val factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
    val secureRandom = new SecureRandom()
    val saltLength = 12
    val keyLength = 128
    val iterationCount = 10000
    val tagSize = 128

    val encode = (bytes: Array[Byte]) =>
      Base64.getEncoder().encodeToString(bytes)

    val salt = Array(saltLength.toByte)
    secureRandom.nextBytes(salt)
    val nonce = Array(12.toByte)
    secureRandom.nextBytes(nonce)

    // val spec =
    //   new PBEKeySpec(password.toCharArray(), salt, iterationCount, keyLength)

    // val tmp = factory.generateSecret(spec)
    // val secretKey = new SecretKeySpec(tmp.getEncoded(), "AES")
    // val cipher = Cipher.getInstance("AES/GCM/NoPadding")
    // cipher.init(
    //   Cipher.ENCRYPT_MODE,
    //   secretKey,
    //   new GCMParameterSpec(128, nonce)
    // )

    // val cipherTextByte = cipher.doFinal(plainText.getBytes)

    val messageCipher = for {
      factory <- {
        Try(SecretKeyFactory.getInstance("PBKDFWithHmacSHA56")).toOption
          .toRight("Failed to get skf instance")
      }

      spec <- {
        Try(
          new PBEKeySpec(
            password.toCharArray(),
            salt,
            iterationCount,
            keyLength
          )
        ).toOption.toRight("Failed to get pbekeyspec")
      }

      secret <- Try(factory.generateSecret(spec)).toOption
        .toRight("Failed to get secret")

      secretKey <- Try(new SecretKeySpec(secret.getEncoded(), "AES")).toOption
        .toRight("Failed to get secret key")

      cipher <- Try(Cipher.getInstance("AES/GCM/NoPadding")).toOption
        .toRight("Failed to get cipher instance")
      _ <- Right(
        cipher.init(
          Cipher.ENCRYPT_MODE,
          secretKey,
          new GCMParameterSpec(128, nonce)
        )
      )

      cipherTextByte <- Try(cipher.doFinal(plainText.getBytes())).toOption
        .toRight("Failed to generate cipher")

      messageCipher = MessageCipher(
        v = 1,
        salt = encode(salt),
        mode = "gcm",
        iterations = iterationCount,
        cipher = "aes",
        adata = "",
        cipherText = encode(cipherTextByte),
        iv = encode(nonce),
        keySize = keyLength,
        tagSize = tagSize
      )
    } yield (messageCipher)

    messageCipher
  }

  override def decrypt(
      password: String,
      messageCipher: MessageCipher
  ): Try[String] = {
    val decode = (text: String) => {
      Base64.getDecoder().decode(text)
    }

    Try {
      val factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
      val spec = new PBEKeySpec(
        password.toCharArray(),
        decode(messageCipher.salt),
        messageCipher.iterations,
        messageCipher.keySize
      );
      val tmp = factory.generateSecret(spec);
      val secretKey = new SecretKeySpec(tmp.getEncoded(), "AES");

      val cipher = Cipher.getInstance("AES/GCM/NoPadding")
      cipher.init(
        Cipher.DECRYPT_MODE,
        secretKey,
        new GCMParameterSpec(128, decode(messageCipher.iv))
      );
      new String(cipher.doFinal(decode(messageCipher.cipherText)));
    }
  }

}

object EncryptionServiceImpl {
  def apply() = new EncryptionServiceImpl()
}