/******************************************************************************
 *
 *  Copyright (C) 2016 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 <stdarg.h>

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "bt_trace.h"
#include "hcidefs.h"
#include "stack/include/smp_api.h"
#include "stack/smp/smp_int.h"

/*
 * This test verifies various key distribution methods in SMP works using the
 * following parameter set:
 *
 * When testing target as Master (Initiator is local, Responder is remote)
 *
 * Initiator's Pairing Request: 0x070710000001(01)
 * Responder's Pairing Response: 0x050008000003(02)
 * Initiator's Bluetooth Address: 0xA1A2A3A4A5A6
 * Initiator's Bluetooth Address Type: 0x01
 * Responder's Bluetooth Address: 0xB1B2B3B4B5B6
 * Responder's Bluetooth Address Type: 0x00
 * Initiator's Random Number: 0x5783D52156AD6F0E6388274EC6702EE0
 * TK Encryption Key: 0x0
 *
 * Correct values:
 *
 * p1: 0x05000800000302070710000001010001
 * p1 XOR r: 0x5283dd2156ae6d096498274ec7712ee1
 * p1 prime: 0x02c7aa2a9857ac866ff91232df0e3c95
 * p2: 0x00000000a1a2a3a4a5a6b1b2b3b4b5b6
 * MConfirm (c1): 0x1e1e3fef878988ead2a74dc5bef13b86
 *
 * NOTE: All these values are presented in mathematical reasonable canonical
 * form that has MSB on the left and LSB on the right. In Bluetooth packets,
 * they are mostly reversed to be Little Endian which have LSB on the left and
 * MSB on the right.
 */

// Set remote bda to 0xB1B2B3B4B5B6
bool BTM_ReadRemoteConnectionAddr(BD_ADDR pseudo_addr, BD_ADDR conn_addr,
                                  tBLE_ADDR_TYPE* p_addr_type) {
  const uint8_t local_bda[] = {0xB1, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6};
  memcpy(conn_addr, local_bda, sizeof(local_bda));
  *p_addr_type = 0x00;
  return true;
}

// Set local_bda to 0xA1A2A3A4A5A6
void BTM_ReadConnectionAddr(BD_ADDR remote_bda, BD_ADDR local_conn_addr,
                            tBLE_ADDR_TYPE* p_addr_type) {
  const uint8_t local_bda[] = {0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6};
  memcpy(local_conn_addr, local_bda, sizeof(local_bda));
  *p_addr_type = 0x01;
}

// Require bte_logmsg.cc to run, here is just to fake it as we don't care about
// trace in unit test
void LogMsg(uint32_t trace_set_mask, const char* fmt_str, ...) {
  va_list args;
  va_start(args, fmt_str);
  vprintf(fmt_str, args);
  va_end(args);
}

extern void smp_gen_p1_4_confirm(tSMP_CB* p_cb,
                                 tBLE_ADDR_TYPE remote_bd_addr_type,
                                 BT_OCTET16 p1);

extern void smp_gen_p2_4_confirm(tSMP_CB* p_cb, BD_ADDR remote_bda,
                                 BT_OCTET16 p2);

extern tSMP_STATUS smp_calculate_comfirm(tSMP_CB* p_cb, BT_OCTET16 rand,
                                         tSMP_ENC* output);

namespace testing {

void dump_uint128(BT_OCTET16 a, char* buffer) {
  for (unsigned int i = 0; i < sizeof(BT_OCTET16); ++i) {
    snprintf(buffer, 3, "%02x", a[i]);
    buffer += 2;
  }
  *buffer = '\0';
}

void dump_uint128_reverse(BT_OCTET16 a, char* buffer) {
  for (int i = (int)(sizeof(BT_OCTET16) - 1); i >= 0; --i) {
    snprintf(buffer, 3, "%02x", a[i]);
    buffer += 2;
  }
  *buffer = '\0';
}

void print_uint128(BT_OCTET16 a) {
  for (unsigned int i = 0; i < sizeof(BT_OCTET16); ++i) {
    printf("%02x", a[i]);
  }
  printf("\n");
}

void parse_uint128(const char* input, BT_OCTET16 output) {
  memset(output, 0, sizeof(BT_OCTET16));
  for (unsigned int count = 0; count < sizeof(BT_OCTET16); count++) {
    sscanf(input, "%2hhx", &output[count]);
    input += 2;
  }
}

void reverse_array_inplace(BT_OCTET16 a) {
  uint8_t tmp;
  uint8_t* a_end = a + sizeof(BT_OCTET16) - 1;
  while (a_end > a) {
    tmp = *a_end;
    *a_end = *a;
    *a = tmp;
    ++a;
    --a_end;
  }
}

class SmpCalculateConfirmTest : public Test {
 protected:
  tSMP_CB p_cb_;
  // Set random to 0x5783D52156AD6F0E6388274EC6702EE0
  BT_OCTET16 rand_ = {0x57, 0x83, 0xD5, 0x21, 0x56, 0xAD, 0x6F, 0x0E,
                      0x63, 0x88, 0x27, 0x4E, 0xC6, 0x70, 0x2E, 0xE0};

  void SetUp() {
    memset(p_cb_.tk, 0, sizeof(p_cb_.tk));
    // Set pairing request packet to 0x070710000001(01)
    p_cb_.local_io_capability = 0x01;
    p_cb_.loc_oob_flag = 0x00;
    p_cb_.loc_auth_req = 0x00;
    p_cb_.loc_enc_size = 0x10;
    p_cb_.local_i_key = 0x07;
    p_cb_.local_r_key = 0x07;
    // Set pairing response packet to 0x050008000003(02)
    p_cb_.peer_io_caps = 0x03;
    p_cb_.peer_oob_flag = 0x00;
    p_cb_.peer_auth_req = 0x00;
    p_cb_.peer_enc_size = 0x08;
    p_cb_.peer_i_key = 0x00;
    p_cb_.peer_r_key = 0x05;
    // Set role to master
    p_cb_.role = HCI_ROLE_MASTER;
    reverse_array_inplace(rand_);
  }
  void TearDown() {}

 public:
};

// Test smp_gen_p2_4_confirm function implementation
TEST_F(SmpCalculateConfirmTest, test_smp_gen_p2_4_confirm_as_master) {
  BT_OCTET16 p2;
  BD_ADDR remote_bda;
  tBLE_ADDR_TYPE remote_bd_addr_type = 0;
  BTM_ReadRemoteConnectionAddr(p_cb_.pairing_bda, remote_bda,
                               &remote_bd_addr_type);
  BTM_ReadConnectionAddr(p_cb_.pairing_bda, p_cb_.local_bda, &p_cb_.addr_type);
  smp_gen_p2_4_confirm(&p_cb_, remote_bda, p2);
  // Correct p2 is 0x00000000a1a2a3a4a5a6b1b2b3b4b5b6
  const char expected_p2_str[] = "00000000a1a2a3a4a5a6b1b2b3b4b5b6";
  char p2_str[2 * sizeof(BT_OCTET16) + 1];
  dump_uint128_reverse(p2, p2_str);
  ASSERT_THAT(p2_str, StrEq(expected_p2_str));
}

// Test smp_gen_p1_4_confirm and SMP_Encrypt function implementation
TEST_F(SmpCalculateConfirmTest, test_SMP_Encrypt_as_master) {
  BT_OCTET16 p1;
  BD_ADDR remote_bda;
  tBLE_ADDR_TYPE remote_bd_addr_type = 0;
  BTM_ReadRemoteConnectionAddr(p_cb_.pairing_bda, remote_bda,
                               &remote_bd_addr_type);
  BTM_ReadConnectionAddr(p_cb_.pairing_bda, p_cb_.local_bda, &p_cb_.addr_type);
  smp_gen_p1_4_confirm(&p_cb_, remote_bd_addr_type, p1);
  // Correct p1 is 0x05000800000302070710000001010001
  const char expected_p1_str[] = "05000800000302070710000001010001";
  char p1_str[2 * sizeof(BT_OCTET16) + 1];
  dump_uint128_reverse(p1, p1_str);
  ASSERT_THAT(p1_str, StrEq(expected_p1_str));
  smp_xor_128(p1, rand_);
  // Correct p1 xor r is 0x5283dd2156ae6d096498274ec7712ee1
  const char expected_p1_xor_r_str[] = "5283dd2156ae6d096498274ec7712ee1";
  char p1_xor_r_str[2 * sizeof(BT_OCTET16) + 1];
  dump_uint128_reverse(p1, p1_xor_r_str);
  ASSERT_THAT(p1_xor_r_str, StrEq(expected_p1_xor_r_str));
  tSMP_ENC output;
  memset(&output, 0, sizeof(tSMP_ENC));
  ASSERT_TRUE(
      SMP_Encrypt(p_cb_.tk, BT_OCTET16_LEN, p1, BT_OCTET16_LEN, &output));
  const char expected_p1_prime_str[] = "02c7aa2a9857ac866ff91232df0e3c95";
  char p1_prime_str[2 * sizeof(BT_OCTET16) + 1];
  dump_uint128_reverse(output.param_buf, p1_prime_str);
  ASSERT_THAT(p1_prime_str, StrEq(expected_p1_prime_str));
}

// Test smp_calculate_comfirm function implementation
TEST_F(SmpCalculateConfirmTest, test_smp_calculate_comfirm_as_master) {
  tSMP_ENC output;
  tSMP_STATUS status = smp_calculate_comfirm(&p_cb_, rand_, &output);
  EXPECT_EQ(status, SMP_SUCCESS);
  // Correct MConfirm is 0x1e1e3fef878988ead2a74dc5bef13b86
  const char expected_confirm_str[] = "1e1e3fef878988ead2a74dc5bef13b86";
  char confirm_str[2 * sizeof(BT_OCTET16) + 1];
  dump_uint128_reverse(output.param_buf, confirm_str);
  ASSERT_THAT(confirm_str, StrEq(expected_confirm_str));
}
}
