#include "com_android_server_ActivationTeeService.h"

#include <stdio.h>
#include <activation.h>

#define PLATFORM_LOG_TAG "Activation JNI"
#include <tees_log.h>

static void throwJavaException(JNIEnv *env, ActivationResult res);

#define kMaxMessageSize 1024

/*
 * Class:     com_android_server_ActivationTeeService
 * Method:    generateSessionCertificate
 * Signature: ()Lcom/samsung/android/service/activationteeservice/ActivationTeeCertificates;
 */
JNIEXPORT jobject JNICALL Java_com_android_server_ActivationTeeService_generateSessionCertificate(
    JNIEnv *env, jobject this) {
  uint8_t rsa_cert[kActivationMaxBufferSize];
  size_t rsa_cert_size = sizeof(rsa_cert);
  uint8_t drk_cert[kActivationMaxBufferSize];
  size_t drk_cert_size = sizeof(drk_cert);

  ssize_t res = ActivationGenerateSessionCertificate(drk_cert, &drk_cert_size,
                                                     rsa_cert, &rsa_cert_size);

  if (res != kActivationSuccess) {
    throwJavaException(env, res);
    return NULL;
  }

  // Session certificate was generated successfully
  jclass java_class = (*env)->FindClass(env,
      "com/samsung/android/service/activationteeservice/ActivationTeeCertificates");
  if (java_class == NULL) {
    return NULL;
  }

  jmethodID constructor = (*env)->GetMethodID(env, java_class, "<init>",
      "([B[B)V");
  if (constructor == NULL) {
    return NULL;
  }

  jbyteArray j_drk_cert = (*env)->NewByteArray(env, drk_cert_size);
  jbyteArray j_rsa_cert = (*env)->NewByteArray(env, rsa_cert_size);

  (*env)->SetByteArrayRegion(env, j_drk_cert, 0, drk_cert_size, (jbyte *)drk_cert);
  (*env)->SetByteArrayRegion(env, j_rsa_cert, 0, rsa_cert_size, (jbyte *)rsa_cert);

  // Finaly, creating new ActivationGenerateCredentials java object
  return (*env)->NewObject(env, java_class, constructor, j_drk_cert, j_rsa_cert);
}

/*
 * Class:     com_android_server_ActivationTeeService
 * Method:    storeServerKey
 * Signature: ([B)V
 */
JNIEXPORT void JNICALL Java_com_android_server_ActivationTeeService_storeServerKey(
    JNIEnv *env, jobject this, jbyteArray encryptedKey) {
  jboolean isCopyData;

  if (!encryptedKey) {
    throwJavaException(env, kActivationErrorBadParameters);
    return;
  }

  jbyte *data = (*env)->GetByteArrayElements(env, encryptedKey, &isCopyData);
  size_t data_len = (*env)->GetArrayLength(env, encryptedKey);

  ssize_t res = ActivationStoreServerKey(data, data_len);

  if (isCopyData) {
    (*env)->ReleaseByteArrayElements(env, encryptedKey, data, JNI_ABORT);
  }

  if (res != kActivationSuccess) {
    throwJavaException(env, res);
  }
}

/*
 * Class:     com_android_server_ActivationTeeService
 * Method:    generateCredentials
 * Signature: ([B)Lcom/samsung/android/service/activationteeservice/ActivationTeeCredentials;
 */
JNIEXPORT jobject JNICALL Java_com_android_server_ActivationTeeService_generateCredentials(
    JNIEnv *env, jobject this, jbyteArray jinput) {
  jboolean isCopyInput;

  if (!jinput) {
    throwJavaException(env, kActivationErrorBadParameters);
    return NULL;
  }

  jbyte *input = (*env)->GetByteArrayElements(env, jinput, &isCopyInput);
  jsize input_size = (*env)->GetArrayLength(env, jinput);

  uint8_t encrypted[kActivationMaxBufferSize];
  uint8_t iv[kActivationIvSize];
  uint8_t tag[kActivationTagSize];
  size_t encrypted_size = sizeof(encrypted);
  size_t iv_size = sizeof(iv);
  size_t tag_size = sizeof(tag);

  // Calling activation-lib function
  ssize_t res = ActivationGenerateCredentials(
      (const uint8_t *)input, (size_t)input_size,
      encrypted, &encrypted_size,
      iv, &iv_size,
      tag, &tag_size);

  // Now we can release input
  if (isCopyInput) {
    (*env)->ReleaseByteArrayElements(env, jinput, input, JNI_ABORT);
  }

  // Checking activation-lib return code
  if (res != 0) {
    throwJavaException(env, res);
    return NULL;
  }

  // ActivationGenerateCredentials was executed successfully.
  // So we are creating ActivationGenerateCredentials java object with all data.
  jclass java_class = (*env)->FindClass(env,
      "com/samsung/android/service/activationteeservice/ActivationTeeCredentials");
  if (java_class == NULL) {
    return NULL;
  }

  jmethodID constructor = (*env)->GetMethodID(env, java_class, "<init>",
      "([B[B[B)V");
  if (constructor == NULL) {
    return NULL;
  }

  // Allocating java byte arrays and copy data to it
  jbyteArray jencrypted = (*env)->NewByteArray(env, encrypted_size);
  (*env)->SetByteArrayRegion(env, jencrypted, 0, encrypted_size,
      (jbyte *)encrypted);
  jbyteArray jiv = (*env)->NewByteArray(env, iv_size);
  (*env)->SetByteArrayRegion(env, jiv, 0, iv_size, (jbyte *)iv);
  jbyteArray jtag = (*env)->NewByteArray(env, tag_size);
  (*env)->SetByteArrayRegion(env, jtag, 0, tag_size, (jbyte *)tag);

  // Finaly, creating new ActivationGenerateCredentials java object
  return (*env)->NewObject(env, java_class, constructor, jencrypted, jiv, jtag);
}

static void throwJavaException(JNIEnv *env, ActivationResult res) {
  jclass exception = (*env)->FindClass(
      env, "com/samsung/android/service/activationteeservice/ActivationTeeException");
  if (exception == NULL) {
    MB_LOGE("Exception class is not found\n");
    return;
  }

  char message[kMaxMessageSize] = {0};
  snprintf(message, sizeof(message), "Native error: %d\n", (int)res);
  (*env)->ThrowNew(env, exception, message);
}