#include <cstdio>
#include <cstring>
#include "Libraries/Crypto.hpp"
using namespace Kt::Crypto;

static int fails = 0;
static void hexdump(const char* label, const uint8_t* p, size_t n) {
    printf("%s: ", label);
    for (size_t i = 0; i < n; i++) printf("%02x", p[i]);
    printf("\n");
}
static void check(const char* name, const uint8_t* got, const char* wantHex) {
    size_t n = strlen(wantHex) / 2;
    uint8_t want[128];
    for (size_t i = 0; i < n; i++) { unsigned v; sscanf(wantHex + 2*i, "%2x", &v); want[i] = (uint8_t)v; }
    if (memcmp(got, want, n) == 0) { printf("PASS %s\n", name); }
    else { printf("FAIL %s\n", name); hexdump("  got ", got, n); printf("  want: %s\n", wantHex); fails++; }
}

int main() {
    uint8_t out[64];

    Sha1("abc", 3, out);
    check("sha1(abc)", out, "a9993e364706816aba3e25717850c26c9cd0d89d");
    Sha1("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq", 56, out);
    check("sha1(448bit)", out, "84983e441c3bd26ebaae4aa1f95129e5e54670f1");

    Sha256("abc", 3, out);
    check("sha256(abc)", out, "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
    Sha256("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq", 56, out);
    check("sha256(448bit)", out, "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1");

    { uint8_t k[20]; memset(k, 0x0b, 20);
      HmacSha1(k, 20, "Hi There", 8, out);
      check("hmac-sha1 rfc2202#1", out, "b617318655057264e28bc0b6fb378c8ef146be00"); }
    { HmacSha1((const uint8_t*)"Jefe", 4, "what do ya want for nothing?", 28, out);
      check("hmac-sha1 rfc2202#2", out, "effcdf6ae5eb2fa2d27416d5f184df9c259a7c79"); }
    { uint8_t k[20]; memset(k, 0x0b, 20);
      HmacSha256(k, 20, "Hi There", 8, out);
      check("hmac-sha256 rfc4231#1", out, "b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7"); }

    // FIPS-197 AES-128 / AES-256
    { AesCtx c; uint8_t key[16], pt[16], ct[16], back[16];
      for (int i = 0; i < 16; i++) key[i] = (uint8_t)i;
      for (int i = 0; i < 16; i++) pt[i] = (uint8_t)(i * 0x11);
      AesInit(c, key, 16); AesEncryptBlock(c, pt, ct);
      check("aes128 fips197", ct, "69c4e0d86a7b0430d8cdb78070b4c55a");
      AesDecryptBlock(c, ct, back);
      check("aes128 decrypt", back, "00112233445566778899aabbccddeeff"); }
    { AesCtx c; uint8_t key[32], pt[16], ct[16], back[16];
      for (int i = 0; i < 32; i++) key[i] = (uint8_t)i;
      for (int i = 0; i < 16; i++) pt[i] = (uint8_t)(i * 0x11);
      AesInit(c, key, 32); AesEncryptBlock(c, pt, ct);
      check("aes256 fips197", ct, "8ea2b7ca516745bfeafc49904b496089");
      AesDecryptBlock(c, ct, back);
      check("aes256 decrypt", back, "00112233445566778899aabbccddeeff"); }

    // RFC 3394 section 4.1 (128-bit KEK, 128-bit key) and 4.6 (256/256)
    { uint8_t kek[16], kd[16], wrapped[24], unwrapped[16];
      for (int i = 0; i < 16; i++) kek[i] = (uint8_t)i;
      for (int i = 0; i < 16; i++) kd[i] = (uint8_t)(i * 0x11);
      AesKeyWrap(kek, 16, kd, 16, wrapped);
      check("keywrap rfc3394 4.1", wrapped, "1fa68b0a8112b447aef34bd8fb5a7b829d3e862371d2cfe5");
      bool ok = AesKeyUnwrap(kek, 16, wrapped, 24, unwrapped);
      printf("%s keyunwrap integrity\n", ok ? "PASS" : "FAIL"); if (!ok) fails++;
      check("keyunwrap rfc3394 4.1", unwrapped, "00112233445566778899aabbccddeeff"); }
    { // 256-bit KEK, 256-bit key data (RFC3394 4.6)
      uint8_t kek[32], kd[32], wrapped[40], unwrapped[32];
      for (int i = 0; i < 32; i++) kek[i] = (uint8_t)i;
      const char* kdhex = "00112233445566778899AABBCCDDEEFF000102030405060708090A0B0C0D0E0F";
      for (int i = 0; i < 32; i++) { unsigned v; sscanf(kdhex + 2*i, "%2x", &v); kd[i] = (uint8_t)v; }
      AesKeyWrap(kek, 32, kd, 32, wrapped);
      check("keywrap rfc3394 4.6", wrapped, "28c9f404c4b810f4cbccb35cfb87f8263f5786e2d80ed326cbc7f0e71a99f43bfb988b9b7a02dd21");
      bool ok = AesKeyUnwrap(kek, 32, wrapped, 40, unwrapped);
      printf("%s keyunwrap256 integrity\n", ok ? "PASS" : "FAIL"); if (!ok) fails++; }

    // RFC 4493 AES-CMAC
    { uint8_t key[16] = {0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c};
      const uint8_t* parts[1]; size_t lens[1];
      uint8_t empty[1] = {0};
      parts[0] = empty; lens[0] = 0;
      AesCmac(key, 16, parts, lens, 1, out);
      check("cmac rfc4493 len0", out, "bb1d6929e95937287fa37d129b756746");
      uint8_t msg[64] = {0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a,
                         0xae,0x2d,0x8a,0x57,0x1e,0x03,0xac,0x9c,0x9e,0xb7,0x6f,0xac,0x45,0xaf,0x8e,0x51,
                         0x30,0xc8,0x1c,0x46,0xa3,0x5c,0xe4,0x11,0xe5,0xfb,0xc1,0x19,0x1a,0x0a,0x52,0xef,
                         0xf6,0x9f,0x24,0x45,0xdf,0x4f,0x9b,0x17,0xad,0x2b,0x41,0x7b,0xe6,0x6c,0x37,0x10};
      parts[0] = msg; lens[0] = 16;
      AesCmac(key, 16, parts, lens, 1, out);
      check("cmac rfc4493 len16", out, "070a16b46b4d4144f79bdd9dd04a287c");
      parts[0] = msg; lens[0] = 40;
      AesCmac(key, 16, parts, lens, 1, out);
      check("cmac rfc4493 len40", out, "dfa66747de9ae63030ca32611497c827");
      parts[0] = msg; lens[0] = 64;
      AesCmac(key, 16, parts, lens, 1, out);
      check("cmac rfc4493 len64", out, "51f0bebf7e3b9d92fc49741779363cfe");
      // split across parts to exercise the streaming path
      const uint8_t* sp[3] = {msg, msg+10, msg+33}; size_t sl[3] = {10, 23, 31};
      AesCmac(key, 16, sp, sl, 3, out);
      check("cmac split len64", out, "51f0bebf7e3b9d92fc49741779363cfe"); }

    // WPA PSK vectors (IEEE 802.11i Annex H.4)
    { uint8_t psk[32];
      Pbkdf2Sha1("password", 8, (const uint8_t*)"IEEE", 4, 4096, psk, 32);
      check("pbkdf2 wpa 'password'/IEEE", psk, "f42c6fc52df0ebef9ebb4b90b38a5f902e83fe1b135a70e23aed762e9710a12e");
      Pbkdf2Sha1("ThisIsAPassword", 15, (const uint8_t*)"ThisIsASSID", 11, 4096, psk, 32);
      check("pbkdf2 wpa 'ThisIsAPassword'", psk, "0dc0d6eb90555ed6419756b9a15ec3e3209b63df707dd508d14581f8982721af"); }

    printf(fails ? "\n%d FAILURES\n" : "\nall vectors pass\n", fails);
    return fails != 0;
}