/*
 * Copyright 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include <openssl/ec.h>
#include <openssl/ecdh.h>
#include <openssl/obj_mac.h>

#define ECDH_KEY_LEN 33
#define ECDH_SHARED_SECRET_LEN 32

/* Computes P256 ECDH shared secret using |private_key| as generated by
 * generate_p256_key() as the private key, and |other_public_key| as the public
 * key. Writes ECDH_SHARED_SECRET_LEN bytes to |shared_secret| and returns 0 on
 * success, returns -1 on failure.
 */
int shared_secret_compute(const uint8_t* private_key,
                          uint32_t private_key_len,
                          const uint8_t other_public_key[ECDH_KEY_LEN],
                          uint8_t shared_secret[ECDH_SHARED_SECRET_LEN]) {
  int ret = -1;
  EC_GROUP* group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
  EC_GROUP_set_point_conversion_form(group, POINT_CONVERSION_COMPRESSED);
  EC_POINT* other_point = EC_POINT_new(group);
  EC_KEY* pkey = EC_KEY_new();

  if (!EC_POINT_oct2point(group, other_point, other_public_key, ECDH_KEY_LEN,
                          NULL)) {
    fprintf(stderr, "Deserializing other_public_key failed\n");
    goto end;
  }
  if (!d2i_ECPrivateKey(&pkey, &private_key, private_key_len)) {
    fprintf(stderr, "Deserializing private_key failed\n");
    goto end;
  }
  EC_KEY_set_group(pkey, group);
  ret = ECDH_compute_key(shared_secret, ECDH_SHARED_SECRET_LEN, other_point,
                         pkey, NULL);
  if (ret != ECDH_SHARED_SECRET_LEN) {
    fprintf(stderr, "Failed to compute shared secret: %d\n", ret);
    ret = -1;
    goto end;
  }
  ret = 0;

end:
  EC_POINT_free(other_point);
  EC_GROUP_free(group);
  EC_KEY_free(pkey);
  return ret;
}

/* Generates a new EC Key to be used for computing a shared secret with the
 * device. On success, returns 0 and writes key material and number of bytes
 * allocated for |*private_key| to |*private_key_len|, and writes ECDH_KEY_LEN
 * bytes to public_key. Returns -1 on failure.
 */
int generate_p256_key(uint8_t** private_key, uint32_t* private_key_len,
                      uint8_t public_key[ECDH_KEY_LEN]) {
  int pkey_len = 0, ret = -1;
  EC_GROUP* group = NULL;
  const EC_POINT* point = NULL;
  EC_KEY* pkey = NULL;
  uint8_t* tmp;

  if (!private_key || !private_key_len || !public_key) {
    fprintf(stderr, "Invalid input parameters\n");
    return -1;
  }
  pkey = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
  if (!EC_KEY_generate_key(pkey)) {
    fprintf(stderr, "Failed to generate key\n");
    return -1;
  }
  pkey_len = i2d_ECPrivateKey(pkey, NULL);
  if (pkey_len == -1) {
    fprintf(stderr, "Failed to get private key length\n");
    goto end;
  }
  *private_key_len = pkey_len;
  *private_key = (uint8_t*)malloc(pkey_len);
  tmp = *private_key;
  if (!i2d_ECPrivateKey(pkey, &tmp)) {
    fprintf(stderr, "Failed to serialize private key\n");
    goto end;
  }
  group = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
  point = EC_KEY_get0_public_key(pkey);
  if (!EC_POINT_point2oct(group, point, POINT_CONVERSION_COMPRESSED,
      public_key, ECDH_KEY_LEN, NULL)) {
    fprintf(stderr, "Failed to serialize public key\n");
    goto end;
  }
  ret = 0;

end:
  EC_GROUP_free(group);
  EC_KEY_free(pkey);
  if (ret == -1 && *private_key) {
    free(*private_key);
  }
  return ret;
}