// Copyright 2024 The Pigweed Authors
//
// 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
//
//     https://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 "pw_crypto/aes_mbedtls.h"

#include <mbedtls/aes.h>
#include <mbedtls/cmac.h>

#include "pw_assert/check.h"
#include "pw_crypto/aes.h"
#include "pw_status/status.h"

namespace {
/// Number of bits in a byte.
constexpr size_t kBits = 8;
}  // namespace

namespace pw::crypto::aes::backend {
namespace {
using CmacCtx = NativeCmacContext;
}  // namespace

Status DoInit(CmacCtx& ctx, ConstByteSpan key) {
  const auto key_data = reinterpret_cast<const unsigned char*>(key.data());

  const mbedtls_cipher_info_t* info;
  switch (key.size()) {
    case kKey128SizeBytes:
      info = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_ECB);
      break;
    case kKey192SizeBytes:
      info = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_192_ECB);
      break;
    case kKey256SizeBytes:
      info = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_256_ECB);
      break;
    default:
      PW_CRASH("Unsupported key size for Cmac (%zu bit)", key.size() * kBits);
  }

  if (mbedtls_cipher_setup(&ctx.cipher, info)) {
    return Status::Internal();
  }
  if (mbedtls_cipher_cmac_starts(&ctx.cipher, key_data, key.size() * kBits)) {
    return Status::Internal();
  }

  return OkStatus();
}

Status DoUpdate(CmacCtx& ctx, ConstByteSpan data) {
  const auto data_in = reinterpret_cast<const unsigned char*>(data.data());
  if (mbedtls_cipher_cmac_update(&ctx.cipher, data_in, data.size())) {
    return Status::Internal();
  }

  return OkStatus();
}

Status DoFinal(CmacCtx& ctx, BlockSpan out_mac) {
  const auto data_out = reinterpret_cast<unsigned char*>(out_mac.data());
  if (mbedtls_cipher_cmac_finish(&ctx.cipher, data_out)) {
    return Status::Internal();
  }

  return OkStatus();
}

Status DoEncryptBlock(ConstByteSpan key,
                      ConstBlockSpan plaintext,
                      BlockSpan out_ciphertext) {
  const auto key_data = reinterpret_cast<const unsigned char*>(key.data());
  const auto in = reinterpret_cast<const unsigned char*>(plaintext.data());
  const auto out = reinterpret_cast<unsigned char*>(out_ciphertext.data());

  mbedtls_aes_context aes;
  mbedtls_aes_init(&aes);

  if (mbedtls_aes_setkey_enc(&aes, key_data, key.size() * kBits)) {
    mbedtls_aes_free(&aes);
    return Status::Internal();
  }
  if (mbedtls_aes_crypt_ecb(&aes, MBEDTLS_AES_ENCRYPT, in, out)) {
    mbedtls_aes_free(&aes);
    return Status::Internal();
  }

  mbedtls_aes_free(&aes);
  return OkStatus();
}

}  // namespace pw::crypto::aes::backend