#!/usr/bin/env python3
# Known Answer Test generator

import json
import hashlib
import secrets

def blake2s(w, key):
  ctx = hashlib.blake2s(key=key)
  ctx.update(w)
  return ctx.digest().hex()

def out(s):
  o = [s[i:i+2] for i in range(0, len(s), 2)]     # split into pairs
  o = [f'0x{i}' for i in o]                       # prepend 0x and join
  return ', '.join(o)

if __name__ == '__main__':
  w = b''
  for i in range(0, 256):
    w += i.to_bytes(1, 'little')
  k = secrets.token_bytes(32)

  print('#pragma once\n')
  print(f'static const unsigned KATs_len = 256;')
  print(f'static const uint8_t KAT_secret[32] = {{ {", ".join([hex(i) for i in k])} }};')

  print(f'static const uint8_t KATs[256][32] = {{')
  for i in range(0, 256):
    o = blake2s(w[0:i], b'')
    print(f'  // {i}')
    print(f'  {{ {out(o)} }},')
  print(f'}};')

  print(f'static const uint8_t secret_KATs[256][32] = {{')
  for i in range(0, 256):
    o = blake2s(w[0:i], k)
    print(f'  // {i}')
    print(f'  {{ {out(o)} }},')
  print(f'}};')