#include "Utils.h"
#include <AES.h>
#include <SHA256.h>

#ifdef ARDUINO
  #include <Arduino.h>
#endif

namespace mesh {

extern "C" void MeshCoreWatchdogStage(uint8_t stage) __attribute__((weak));

namespace {
#ifdef AVR_PLATFORM
typedef AESSmall128 MeshAES128;
#else
typedef AES128 MeshAES128;
#endif

void watchdogStage(uint8_t stage) {
  if (MeshCoreWatchdogStage) {
    MeshCoreWatchdogStage(stage);
  }
}
}

uint32_t RNG::nextInt(uint32_t _min, uint32_t _max) {
  uint32_t num;
  random((uint8_t *) &num, sizeof(num));
  return (num % (_max - _min)) + _min;
}

void Utils::sha256(uint8_t *hash, size_t hash_len, const uint8_t* msg, int msg_len) {
  SHA256 sha;
  if (msg && msg_len > 0) {
    sha.update(msg, (size_t)msg_len);
  }
  sha.finalize(hash, hash_len);
}

void Utils::sha256(uint8_t *hash, size_t hash_len, const uint8_t* frag1, int frag1_len, const uint8_t* frag2, int frag2_len) {
  SHA256 sha;
  if (frag1 && frag1_len > 0) {
    sha.update(frag1, (size_t)frag1_len);
  }
  if (frag2 && frag2_len > 0) {
    sha.update(frag2, (size_t)frag2_len);
  }
  sha.finalize(hash, hash_len);
}

int Utils::decrypt(const uint8_t* shared_secret, uint8_t* dest, const uint8_t* src, int src_len) {
  if (src_len <= 0 || (src_len % CIPHER_BLOCK_SIZE) != 0 || src_len > MAX_PACKET_PAYLOAD) {
    return 0;
  }

  MeshAES128 aes;
  uint8_t* dp = dest;
  const uint8_t* sp = src;
  uint8_t block = 0;

  watchdogStage('k');
  aes.setKey(shared_secret, CIPHER_KEY_SIZE);
  while (sp - src < src_len) {
    watchdogStage((uint8_t)('0' + block));
    aes.decryptBlock(dp, sp);
    dp += 16; sp += 16;
    if (block < 9) {
      ++block;
    }
  }
  watchdogStage('d');

  return sp - src;  // will always be multiple of 16
}

int Utils::encrypt(const uint8_t* shared_secret, uint8_t* dest, const uint8_t* src, int src_len) {
  if (src_len < 0 || src_len > MAX_PACKET_PAYLOAD) {
    return 0;
  }
  int enc_len = ((src_len + CIPHER_BLOCK_SIZE - 1) / CIPHER_BLOCK_SIZE) * CIPHER_BLOCK_SIZE;
  if (enc_len > MAX_PACKET_PAYLOAD) {
    return 0;
  }

  MeshAES128 aes;
  uint8_t* dp = dest;

  watchdogStage('K');
  aes.setKey(shared_secret, CIPHER_KEY_SIZE);
  while (src_len >= 16) {
    watchdogStage('E');
    aes.encryptBlock(dp, src);
    dp += 16; src += 16; src_len -= 16;
  }
  if (src_len > 0) {  // remaining partial block
    uint8_t tmp[16];
    memset(tmp, 0, 16);
    memcpy(tmp, src, src_len);
    watchdogStage('e');
    aes.encryptBlock(dp, tmp);
    dp += 16;
  }
  return dp - dest;  // will always be multiple of 16
}

int Utils::encryptThenMAC(const uint8_t* shared_secret, uint8_t* dest, const uint8_t* src, int src_len) {
  int enc_len = encrypt(shared_secret, dest + CIPHER_MAC_SIZE, src, src_len);

  SHA256 sha;
  sha.resetHMAC(shared_secret, PUB_KEY_SIZE);
  sha.update(dest + CIPHER_MAC_SIZE, enc_len);
  sha.finalizeHMAC(shared_secret, PUB_KEY_SIZE, dest, CIPHER_MAC_SIZE);

  return CIPHER_MAC_SIZE + enc_len;
}

int Utils::MACThenDecrypt(const uint8_t* shared_secret, uint8_t* dest, const uint8_t* src, int src_len) {
  if (src_len <= CIPHER_MAC_SIZE) return 0;  // invalid src bytes
  int cipher_len = src_len - CIPHER_MAC_SIZE;
  if ((cipher_len % CIPHER_BLOCK_SIZE) != 0 || cipher_len > MAX_PACKET_PAYLOAD) {
    return 0;
  }

  uint8_t hmac[CIPHER_MAC_SIZE];
  {
    watchdogStage('h');
    SHA256 sha;
    sha.resetHMAC(shared_secret, PUB_KEY_SIZE);
    sha.update(src + CIPHER_MAC_SIZE, cipher_len);
    sha.finalizeHMAC(shared_secret, PUB_KEY_SIZE, hmac, CIPHER_MAC_SIZE);
  }
  if (memcmp(hmac, src, CIPHER_MAC_SIZE) == 0) {
    watchdogStage('a');
    return decrypt(shared_secret, dest, src + CIPHER_MAC_SIZE, cipher_len);
  }
  return 0; // invalid HMAC
}

static const char hex_chars[] = "0123456789ABCDEF";

void Utils::toHex(char* dest, const uint8_t* src, size_t len) {
  while (len > 0) {
    uint8_t b = *src++;
    *dest++ = hex_chars[b >> 4];
    *dest++ = hex_chars[b & 0x0F];
    len--;
  }
  *dest = 0;
}

void Utils::printHex(Stream& s, const uint8_t* src, size_t len) {
  while (len > 0) {
    uint8_t b = *src++;
    s.print(hex_chars[b >> 4]);
    s.print(hex_chars[b & 0x0F]);
    len--;
  }
}

static uint8_t hexVal(char c) {
  if (c >= 'A' && c <= 'F') return c - 'A' + 10;
  if (c >= 'a' && c <= 'f') return c - 'a' + 10;
  if (c >= '0' && c <= '9') return c - '0';
  return 0;
}

bool Utils::isHexChar(char c) {
  return c == '0' || hexVal(c) > 0;
}

bool Utils::fromHex(uint8_t* dest, int dest_size, const char *src_hex) {
  int len = strlen(src_hex);
  if (len != dest_size*2) return false;  // incorrect length

  uint8_t* dp = dest;
  while (dp - dest < dest_size) {
    char ch = *src_hex++;
    char cl = *src_hex++;
    *dp++ = (hexVal(ch) << 4) | hexVal(cl);
  }
  return true;
}

int Utils::parseTextParts(char* text, const char* parts[], int max_num, char separator) {
  int num = 0;
  char* sp = text;
  while (*sp && num < max_num) {
    parts[num++] = sp;
    while (*sp && *sp != separator) sp++;
    if (*sp) {
       *sp++ = 0;  // replace the seperator with a null, and skip past it
    }
  }
  // if we hit the maximum parts, make sure LAST entry does NOT have separator 
  while (*sp && *sp != separator) sp++;
  if (*sp) {
    *sp = 0;  // replace the separator with null
  }
  return num;
}

}