#include "ayla/crypto.hpp"

#include <cstdio>
#include <cstring>

#include "mbedtls/platform_util.h"

namespace fgl::ayla {

namespace {

void hmac_sha256(const uint8_t* key, size_t key_len, const uint8_t* data,
                 size_t data_len, uint8_t out[kSha256Len]) {
  const mbedtls_md_info_t* md = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
  mbedtls_md_hmac(md, key, key_len, data, data_len, out);
}

}  // namespace

bool DirectionCrypto::derive(const uint8_t* key, size_t key_len,
                             const uint8_t* msg, size_t msg_len,
                             uint8_t out[kSha256Len]) {
  // key = HMAC_k( HMAC_k(msg) || msg )
  if (msg_len > sizeof(uint8_t) * (kSha256Len + 192) - kSha256Len) {
    return false;  // невозможно по протоколу (<=73 байт)
  }
  uint8_t inner[kSha256Len];
  hmac_sha256(key, key_len, msg, msg_len, inner);
  uint8_t buf[kSha256Len + 192];
  memcpy(buf, inner, kSha256Len);
  memcpy(buf + kSha256Len, msg, msg_len);
  hmac_sha256(key, key_len, buf, kSha256Len + msg_len, out);
  return true;
}

bool DirectionCrypto::init(const char* lanip_key, const char* rnd_a,
                           const char* rnd_b, int64_t t_a, int64_t t_b) {
  if (lanip_key == nullptr || rnd_a == nullptr || rnd_b == nullptr) return false;
  initialized_ = false;
  mode_ = 0;
  // затираем предыдущий материал
  mbedtls_platform_zeroize(sign_key_, sizeof(sign_key_));
  mbedtls_platform_zeroize(aes_key_, sizeof(aes_key_));
  mbedtls_platform_zeroize(chain_iv_, sizeof(chain_iv_));

  uint8_t msg[128];
  size_t off = 0;
  auto append = [&](const char* s) {
    size_t n = strlen(s);
    if (off + n > sizeof(msg)) return false;
    memcpy(msg + off, s, n);
    off += n;
    return true;
  };
  char t_a_buf[24], t_b_buf[24];
  snprintf(t_a_buf, sizeof(t_a_buf), "%lld", static_cast<long long>(t_a));
  snprintf(t_b_buf, sizeof(t_b_buf), "%lld", static_cast<long long>(t_b));
  if (!append(rnd_a) || !append(rnd_b) || !append(t_a_buf) || !append(t_b_buf) ||
      off + 1 > sizeof(msg)) {
    return false;
  }

  const uint8_t* k = reinterpret_cast<const uint8_t*>(lanip_key);
  size_t klen = strlen(lanip_key);

  uint8_t block[kSha256Len];
  msg[off] = 0x30;  // sign key
  if (!derive(k, klen, msg, off + 1, block)) return false;
  memcpy(sign_key_, block, kHmacKeyLen);

  msg[off] = 0x31;  // aes key
  if (!derive(k, klen, msg, off + 1, block)) return false;
  memcpy(aes_key_, block, kAesKeyLen);

  msg[off] = 0x32;  // iv seed (первые 16 байт)
  if (!derive(k, klen, msg, off + 1, block)) return false;
  memcpy(chain_iv_, block, kAesBlockLen);

  mbedtls_aes_init(&aes_enc_);
  mbedtls_aes_init(&aes_dec_);
  if (mbedtls_aes_setkey_enc(&aes_enc_, aes_key_, 256) != 0) return false;
  if (mbedtls_aes_setkey_dec(&aes_dec_, aes_key_, 256) != 0) return false;

  initialized_ = true;
  return true;
}

long DirectionCrypto::encrypt(uint8_t* out, size_t out_cap, const uint8_t* in,
                              size_t in_len) {
  if (!initialized_) return -1;
  if (mode_ == 0) mode_ = 1;
  if (mode_ != 1) return -1;  // направление уже в режиме расшифровки
  // Java-паддинг: >=1 NUL, кратно 16.
  size_t padded = ((in_len + 1 + kAesBlockLen - 1) / kAesBlockLen) * kAesBlockLen;
  if (padded > out_cap) return -1;
  uint8_t buf[kAesBlockLen * 128];  // до 2048 байт (kHttpdMaxBody)
  if (padded > sizeof(buf)) return -1;
  memcpy(buf, in, in_len);
  memset(buf + in_len, 0, padded - in_len);
  // mbedtls_aes_crypt_cbc использует и обновляет chain_iv_ на месте —
  // это и есть состояние цепочки.
  if (mbedtls_aes_crypt_cbc(&aes_enc_, MBEDTLS_AES_ENCRYPT, padded, chain_iv_,
                            buf, out) != 0) {
    return -1;
  }
  return static_cast<long>(padded);
}

long DirectionCrypto::decrypt(uint8_t* out, size_t out_cap, const uint8_t* in,
                              size_t in_len) {
  if (!initialized_) return -1;
  if (mode_ == 0) mode_ = 2;
  if (mode_ != 2) return -1;  // направление уже в режиме шифрования
  if (in_len == 0 || in_len % kAesBlockLen != 0 || in_len > out_cap) return -1;
  if (mbedtls_aes_crypt_cbc(&aes_dec_, MBEDTLS_AES_DECRYPT, in_len, chain_iv_,
                            in, out) != 0) {
    return -1;
  }
  // Снять завершающие NUL (по крайней мере один в Java-паддинге, но
  // принимаем и строго-кратные длины без NUL).
  size_t len = in_len;
  while (len > 0 && out[len - 1] == 0) len--;
  return static_cast<long>(len);
}

void DirectionCrypto::sign(uint8_t out[kSha256Len], const uint8_t* data,
                           size_t len) const {
  hmac_sha256(sign_key_, kHmacKeyLen, data, len, out);
}

bool SessionCrypto::init(const char* lanip_key, const char* random_1,
                         const char* random_2, int64_t time_1, int64_t time_2) {
  return app.init(lanip_key, random_1, random_2, time_1, time_2) &&
         dev.init(lanip_key, random_2, random_1, time_2, time_1);
}

}  // namespace fgl::ayla