#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>

#include "esphome/components/ble_device_base/ble_aes_ccm.h"

namespace esphome::ble_device_base::testing {

// Reference vector generated with Python `cryptography` AESCCM(tag_length=4),
// using the same AES-128-CCM parameters BTHome advertisements use: a 16-byte
// key, a 13-byte nonce, a 4-byte authentication tag and no associated data.
namespace {
const uint8_t KEY[16] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                         0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
const uint8_t NONCE[13] = {0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c};
const uint8_t CIPHERTEXT[7] = {0x68, 0xb4, 0xf6, 0xc5, 0x2b, 0xf8, 0xaf};
const uint8_t TAG[4] = {0x48, 0x4d, 0xaa, 0x56};
const uint8_t PLAINTEXT[7] = {0x02, 0x01, 0x64, 0x03, 0x10, 0x8a, 0x01};
}  // namespace

TEST(BleAesCcm, DecryptsAndAuthenticatesKnownVector) {
  uint8_t out[sizeof(PLAINTEXT)] = {};
  EXPECT_TRUE(aes_ccm_auth_decrypt(KEY, NONCE, sizeof(NONCE), nullptr, 0, CIPHERTEXT, sizeof(CIPHERTEXT), out, TAG,
                                   sizeof(TAG)));
  EXPECT_EQ(0, memcmp(out, PLAINTEXT, sizeof(PLAINTEXT)));
}

TEST(BleAesCcm, RejectsTamperedTag) {
  uint8_t bad_tag[sizeof(TAG)];
  memcpy(bad_tag, TAG, sizeof(TAG));
  bad_tag[0] ^= 0x01;
  uint8_t out[sizeof(PLAINTEXT)] = {};
  EXPECT_FALSE(aes_ccm_auth_decrypt(KEY, NONCE, sizeof(NONCE), nullptr, 0, CIPHERTEXT, sizeof(CIPHERTEXT), out, bad_tag,
                                    sizeof(bad_tag)));
}

TEST(BleAesCcm, RejectsTamperedCiphertext) {
  uint8_t bad_ct[sizeof(CIPHERTEXT)];
  memcpy(bad_ct, CIPHERTEXT, sizeof(CIPHERTEXT));
  bad_ct[0] ^= 0x01;
  uint8_t out[sizeof(PLAINTEXT)] = {};
  EXPECT_FALSE(
      aes_ccm_auth_decrypt(KEY, NONCE, sizeof(NONCE), nullptr, 0, bad_ct, sizeof(bad_ct), out, TAG, sizeof(TAG)));
}

TEST(BleAesCcm, RejectsWrongKey) {
  uint8_t bad_key[sizeof(KEY)];
  memcpy(bad_key, KEY, sizeof(KEY));
  bad_key[0] ^= 0xFF;
  uint8_t out[sizeof(PLAINTEXT)] = {};
  EXPECT_FALSE(aes_ccm_auth_decrypt(bad_key, NONCE, sizeof(NONCE), nullptr, 0, CIPHERTEXT, sizeof(CIPHERTEXT), out, TAG,
                                    sizeof(TAG)));
}

}  // namespace esphome::ble_device_base::testing